Skip to content

Back to the catalog

ShiftRAG: Bypassing the Textual Bottleneck via Decoupled Learning and Continuous Soft Tokens
arxiv.orgpaper

ShiftRAG: Bypassing the Textual Bottleneck via Decoupled Learning and Continuous Soft Tokens

ShiftRAG is a decoupled retrieval-augmented generation framework that uses one LLM backbone but separates discriminative retrieval learning from autoregressive generation. It projects retrieval alignments into continuous soft tokens so the generator receives retriever-side signals, and adds a soft orthogonality objective against rank collapse in decoder-only models. The abstract reports gains in retrieval and question answering over state-of-the-art baselines, without numbers.

TL;DR Proposes a decoupled RAG framework that passes retrieval signals to the generator as continuous soft tokens instead of text alone.
Document

What it claims

ShiftRAG (Hanyang University, NeurIPS 2026) runs retrieval and generation on one decoder-only LLM using two separate LoRA adapters. The two modes are trained independently and connected only at inference, through "soft knowledge tokens" built from retriever embeddings and relevance scores. The authors argue that unified retrieval-generation models share adaptation parameters, which causes objective conflict and worsens rank collapse. They support this with a diagnostic figure that uses Gemma 2 as the base model. The paper reports state-of-the-art retrieval and end-to-end QA on HotpotQA, 2WikiMultihopQA, MuSiQue and NaturalQuestions, but the fetched text gives no numbers. Code is at https://github.com/meaningful96/ShiftRAG.

How it works

Encoder mode. The causal mask is relaxed to bidirectional attention, and LoRA_enc is trained for retrieval. It starts with a short MNTP warm-up, following LLM2Vec, and then uses two losses.

  • Layer-wise gap-aware contrastive (LGC) loss. The first half of the layers is split into five buckets, and one layer is sampled from each. The relevance score is the largest gap between two values: the final-layer question/document cosine and the question's cosine with each candidate layer's document embedding. This score goes into an InfoNCE loss with positives and hard negatives. At inference, plain final-layer cosine is used, which avoids multi-layer compute.
  • Soft orthogonal alignment (SOA) loss. For each sequence, token hidden states are L2-normalized and their Gram matrix is pushed toward the identity, using squared Frobenius distance divided by |S|². The loss is applied to the last n layers on both the question and document sides and weighted by λ_SOA.

Semantic Signal Translator (SST). For each top-K document, the question and document embeddings are concatenated and encoded. The result is broadcast to M token slots and passed through a low-rank bottleneck (rank r) to give offset embeddings. The offsets are scaled by a normalized relevance weight from the top-K similarities. They are then added to learnable anchor embeddings, initialized from pretrained latent vectors, and to learnable positional embeddings. The K resulting blocks are concatenated.

Decoder mode. The model switches back to causal decoding. It generates from the question, the retrieved document text and the soft tokens.

flowchart LR
  Q[Question] --> ENC[Encoder mode LoRA]
  D[(Corpus, indexed offline by encoder)] --> R[Top-K retrieval]
  ENC --> R
  R -->|embeddings + scores| SST[Semantic Signal Translator]
  R -->|raw text| DEC[Decoder mode LoRA]
  SST -->|soft tokens| DEC
  Q --> DEC
  DEC --> A[Answer]

Evidence

The fetched text gives evidence only in qualitative form:

  • Figure 1 reports that unified LLMs show higher token covariance and lower effective rank than base Gemma 2.
  • Figure 1 also reports that existing baselines show small entropy gaps between correct and incorrect answers on 2WikiMultihopQA.
  • The paper also claims preserved text-encoding and general generation ability and a "favorable efficiency-performance trade-off".

Result tables, baselines and backbone sizes are not in the excerpt.

Limits

The page text ends partway through the method section. No result numbers, hyperparameters (M, r, n, τ, λ_SOA, K) or training costs are visible. The LGC training score and the cosine score used at inference differ. The authors analyze that trade-off in Appendix D.3, which the excerpt does not include. The design needs two adapters, plus an SST module trained on top of a frozen retriever interface. Mode switching per query means a single model serves both stages, so retrieval and generation run in sequence on the same backbone.

Organisations named here

  • arXiv

Technologies named here

  • LLM
  • ShiftRAG

Techniques named here

  • RAG
  • Continuous soft tokens
  • Layer-wise relevance dynamics
  • Soft orthogonality objective

Concepts named here

  • Hallucination
  • Objective conflict
  • Rank collapse
  • Textual bottleneck

Also named here

  • Transformer