NeurIPS 2026Retrieval-augmented generation

ShiftRAG

Bypassing the Textual Bottleneck
via Decoupled Learning and Continuous Soft Tokens

Bringing the retriever’s semantic signals into the generator’s reasoning.

Youmin Ko¹ Jihong Jeong¹ Hyunjoon Kim^1,2,*

¹ Department of Artificial Intelligence, Hanyang University

² Department of Data Science, Hanyang University

DataMining and Algorithm (DMAL) Lab

* Corresponding author

Hanyang University and Data Mining & Algorithm Laboratory

arXiv Code

Dataset (Emb)

Dataset (Gen) BibTeX

56.1%HotpotQA exact matchEnd-to-end QA · Table 2 +23.0%Relative EM on MuSiQuevs. NV-Embed-v2 · Table 2 3.13×End-to-end speedupvs. NV-Embed-v2 · HotpotQA 0.7 msSST overheadSoft token generation · §4.7

01 / Overview

Retrieval carries more than text.

One backbone. Two task-specific adapters. A continuous interface for evidence.

Most RAG systems retrieve documents and pass their text to a language model. ShiftRAG also transfers the retriever’s internal signals: how the question aligns with each document and how strongly that document is ranked. A Semantic Signal Translator (SST) converts these signals into continuous vectors in the generator’s token-embedding space. The generator reads the question and retrieved passages together with these soft knowledge tokens.

The framework uses one pretrained LLM backbone in two modes. Encoder mode enables bidirectional attention with a retrieval-specific LoRA adapter. Decoder mode uses causal attention with a separate generation adapter. Training the modes in stages lets the retrieval space stabilize before the generator learns to use it.

ShiftRAG inference: encode a question, retrieve documents, translate question–document embeddings and relevance weights into soft knowledge tokens, and generate an answer from text plus soft tokens. — SST preserves retrieval-side information as the backbone switches from evidence retrieval to answer generation.

Retrieve evidence

Encode the question and search the offline document index to select the top-K passages.

Translate the signals

Map question–document embeddings and relevance weights into continuous soft knowledge tokens.

Generate an answer

Condition the decoder on the question, retrieved text, and the translated retrieval signals.

02 / Motivation

Three obstacles between evidence and answers.

Better retrieval needs an interface that preserves what the retriever has learned.

The textual bottleneck

Raw passage text does not directly convey the retriever’s embedding geometry or the strength of question–document alignment. Passing only text discards these signals and leaves the generator to infer document relevance again.

Representation rank collapse

Deep decoder-only representations can concentrate in a narrow subspace. Higher token covariance and lower effective rank indicate reduced diversity, making it harder to separate supporting evidence from distractors.

Competing learning objectives

Retrieval favors contrastive separation; generation favors next-token prediction. Shared adaptation can produce conflicting updates. Even with separate parameters, joint training makes the generator follow a changing retrieval interface.

Layer-wise token covariance and effective rank, with answer-token entropy grouped by correct and incorrect predictions on 2WikiMultihopQA. — Layer-wise covariance and effective rank diagnose representation collapse; answer-token entropy probes how generation responds to evidence.

03 / Method

Separate the learning. Connect the knowledge.

Learn a discriminative retrieval space first, then train the generator against that stable interface.

Separate encoder and decoder training. Encoder-mode LoRA uses LGC and SOA; decoder-mode LoRA and SST use answer supervision and rationale-conditioned self-distillation. — The encoder is trained first and then fixed. Decoder training learns to use its representations through SST, answer supervision, and rationale-conditioned distillation.

Stage 1 · Encoder

Preserve useful distinctions.

Layer-wise Gap-aware Contrastive Learning (LGC) learns from the change in relevance between shallow and final layers. The largest question–document similarity gap becomes the contrastive training score, capturing semantic refinement across depth.

Soft Orthogonal Alignment (SOA) discourages excessive correlation among token representations. It softly regularizes their cosine-similarity Gram matrix toward the identity to preserve diverse directions in the retrieval space.

L_enc = L_LGC + λ_SOA · [L_SOA(q) + L_SOA(d)] / 2

At inference, retrieval uses final-layer cosine similarity, avoiding the overhead of multi-layer scoring.

Stage 2 · Decoder

Learn to use continuous evidence.

The decoder combines next-token supervision on the answer with rationale-conditioned self-distillation. A teacher context contains explicit document-level rationales; the student receives soft knowledge tokens.

Matching the teacher’s answer-token distribution teaches the student to use the translated signals. The answer loss provides direct supervision for generating the ground-truth response.

L_dec = L_NTP + λ_KD · L_KD

Explicit rationales provide training-time supervision. They are not required as inference inputs.

04 / Semantic Signal Translator

A continuous bridge into the generator.

Question–document alignment becomes a small set of generation-friendly vectors.

For each retrieved document, SST concatenates the question and document embeddings, projects them through a learnable encoding layer, and expands the representation into token vectors. A low-rank transformation produces an offset in the generator’s embedding space. The document’s normalized relevance weight controls the strength of that offset.

SST combines question and document embeddings, projects them through a low-rank transformation, weights the offset by retrieval relevance, and adds anchor and positional embeddings to form soft knowledge tokens. — Soft tokens combine learnable anchors, a relevance-weighted offset, and positional embeddings in the decoder’s token-embedding space.

T_k = E + α_k ΔE_k + P

Anchors give the translation a vocabulary-aligned reference. The offset carries question-specific retrieval information. Positions distinguish the resulting tokens.

These vectors supplement the retrieved passages, allowing explicit evidence and embedding-level alignment to guide answer generation together. The reported configuration uses five soft knowledge tokens per question–document pair.

05 / Main Results

Stronger retrieval. Better answers.

Highest recall, Exact Match, and F1 among the methods in the reported three-benchmark comparison.

The main evaluation covers HotpotQA, 2WikiMultihopQA, and MuSiQue, with 1,000 test questions per dataset. Retrieval is measured by Recall@2, @5, and @10; QA uses Exact Match (EM) and F1. Retrieval-only baselines are paired with Gemma2-9B for answer generation.

+4.7 ppHotpotQA EM51.4% → 56.1%

+5.0 pp2WikiMultihopQA EM45.8% → 50.8%

+5.7 ppMuSiQue EM24.8% → 30.5%

Absolute percentage-point gains over NV-Embed-v2. The MuSiQue gain is 23.0% relative.

Table 2: ShiftRAG achieves EM/F1 of 0.5610/0.6790 on HotpotQA, 0.5080/0.5528 on 2WikiMultihopQA, and 0.3050/0.4129 on MuSiQue, leading every reported column. — ShiftRAG improves both exact answers and token-level F1 across all three multi-hop QA benchmarks.

65.98%HotpotQA Recall@2NV-Embed-v2: 63.18%

60.80%2WikiMultihopQA Recall@2NV-Embed-v2: 59.56%

41.86%MuSiQue Recall@2NV-Embed-v2: 37.81%

LLM2Vec, OneGen, and ShiftRAG use Gemma2-9B in this reported comparison.

Table 1: ShiftRAG leads Recall@2, Recall@5, and Recall@10 on HotpotQA, 2WikiMultihopQA, and MuSiQue. Recall@10 is 0.8806, 0.7530, and 0.6516 respectively. — Gap-aware, collapse-resilient encoder training improves evidence retrieval across all reported recall cutoffs.

06 / Component Analysis

What makes the difference?

Controlled ablations examine the contribution of each training objective and decoder component.

Complementary objectives for retrieval.

Removing either LGC or SOA lowers retrieval performance in the reported comparison. Replacing LGC with standard InfoNCE causes the larger drop; SOA provides a consistent complementary benefit. On MuSiQue, Recall@10 falls from 65.16% to 61.40% without LGC and to 60.51% without both objectives.

Table 3 compares full ShiftRAG with removal of SOA, LGC, or both. Both objectives contribute to retrieval recall; removing both produces the largest degradation. — LGC supplies the stronger relevance signal, while SOA helps preserve the representational diversity needed to use it.

Translate the evidence, then learn to use it.

Removing SST reduces MuSiQue EM from 30.5% to 24.5%. Removing decoder-side LoRA lowers it to 26.5%; removing both lowers it to 19.7%. Translation and decoder adaptation contribute together to the strongest result.

Table 4: MuSiQue EM is 0.305 with full ShiftRAG, 0.245 without SST, 0.265 without decoder LoRA, and 0.197 without both. — SST and decoder-side LoRA each improve answer quality; their combination performs best in this controlled comparison.

07 / Latency & Efficiency

A lightweight bridge. A practical trade-off.

More accurate answers with a small translation overhead.

2.79 sEnd-to-end / example

24.14 GiBPeak GPU memory

0.7 msSoft token generation

On HotpotQA, ShiftRAG takes 2.79 seconds per example, compared with 8.72 seconds for the NV-Embed-v2 pipeline: a 3.13× speedup, or 68.0% lower latency. Peak VRAM drops from 38.82 to 24.14 GiB while EM rises from 51.4% to 56.1%.

GritLM remains faster at 2.42 seconds and uses less memory at 20.25 GiB. ShiftRAG adds 0.37 seconds while raising EM from 44.3% to 56.1%, offering a different accuracy–efficiency trade-off.

NV-Embed-v28.72 s

GritLM2.42 s

ShiftRAG2.79 s

Most time is spent generating the answer. The reported breakdown is 0.13 s for question encoding, 0.02 s for Faiss retrieval, 0.0007 s for SST, and 2.64 s for answer generation. Values are rounded measurements from the paper’s inference setup.

08 / Resources

Explore and build with ShiftRAG.

Code and datasets for retrieval and generation.

Source codeTraining, retrieval, and evaluationmeaningful96 / ShiftRAG

Embedding datasetEncoder-mode retrieval trainingShiftRAG_Emb · Hugging Face

Generation datasetDecoder-mode QA trainingShiftRAG_Gen · Hugging Face

Citation

Cite this work.

BibTeX