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

Youmin Ko; Jihong Jeong; Hyunjoon Kim

cs.CLOct 31, 2026

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

Authors: Youmin Ko, Jihong Jeong, Hyunjoon Kim

Organizations: Department of Artificial Intelligence, Hanyang University · Department of Data Science, Hanyang University

Abstract

Despite advances in retrieval-augmented generation (RAG), suppressing hallucinations during complex reasoning remains a persistent challenge. We identify that existing RAG pipelines suffer from a textual bottleneck: by passing only raw text to the generator, they discard discriminative signals of the retriever. Furthermore, recent attempts to unify retrieval and generation within a single LLM induce severe objective conflict and exacerbate the rank collapse of decoder-only architectures, degrading performance on both tasks. To address these limitations, we introduce ShiftRAG, a decoupled RAG framework that employs a single LLM backbone but isolates discriminative retrieval learning from autoregressive generation. To bridge these decoupled modes, we project retrieval alignments into continuous soft tokens, establishing a high-bandwidth interface that infuses the generator with retrieval-side signals. Concurrently, our soft orthogonality objective mitigates rank collapse, while layer-wise relevance dynamics preserve the fine-grained latent space required for highly discriminative evidence separation. Extensive evaluations demonstrate that ShiftRAG significantly improves retrieval and question answering performance over state-of-the-art baselines. ShiftRAG preserves foundational capabilities of the backbone while maintaining a favorable efficiency-performance trade-off on general text encoding and generative tasks.

Figures & tables

Figure 1: Motivating observations on layer-wise rank collapse and evidence-sensitive generation. (a) Layer-wise token covariance. (b) Effective rank (eRank) of token representations. (c) Average token entropy on 2WikiMultihopQA by answer correctness.

Figure 2: Illustration of the inference pipeline of ShiftRAG.

Figure 3: Overall framework of ShiftRAG which separately trains the encoder and decoder modes.

Figure 4: Illustration of Semantic Signal Translator (SST).

Methods	HotpotQA			2WikiMultihopQA			MuSiQue
	R@2	R@5	R@10	R@2	R@5	R@10	R@2	R@5	R@10
LLM2Vec	0.6130	0.7854	0.8612	0.5845	0.6451	0.6829	0.3598	0.5126	0.5977
NV-Embed-v2	0.6318	0.7982	0.8736	0.5956	0.6632	0.7202	0.3781	0.5290	0.6236
OneGen	0.5792	0.7450	0.8402	0.5584	0.6077	0.6392	0.3320	0.4777	0.5658
GritLM	0.5890	0.7487	0.8085	0.5558	0.6311	0.6977	0.3760	0.4990	0.5983
ShiftRAG	0.6598	0.8214	0.8806	0.6080	0.7109	0.7530	0.4186	0.5634	0.6516

Table 1: Retrieval performance. The best and second-best performances are presented in bold and underlined, respectively. LLM2Vec, OneGen, and ShiftRAG use Gemma2-9B as the backbone.

Methods	HotpotQA		2WikiMultihopQA		MuSiQue
	EM	F1	EM	F1	EM	F1
LLM2Vec	0.5030	0.6106	0.4400	0.4914	0.2430	0.3670
NV-Embed-v2	0.5140	0.6220	0.4580	0.5064	0.2480	0.3783
GritLM	0.4430	0.5258	0.3020	0.3665	0.1870	0.2872
OneGen	0.4670	0.5869	0.4220	0.4785	0.2200	0.3342
ShiftRAG	0.5610	0.6790	0.5080	0.5528	0.3050	0.4129

Table 2: Question answering performance of ShiftRAG and baselines. The best and second-best performances are presented in bold and underlined, respectively.

Category	HotpotQA			2WikiMultihopQA			MuSiQue
	R@2	R@5	R@10	R@2	R@5	R@10	R@2	R@5	R@10
ShiftRAG	0.6598	0.8214	0.8806	0.6080	0.7109	0.7530	0.4186	0.5634	0.6516
w/o SOA	0.6535	0.8021	0.8678	0.5998	0.6984	0.7425	0.4045	0.5511	0.6493
w/o LGC	0.6360	0.7892	0.8445	0.5816	0.6917	0.7084	0.4008	0.5457	0.6140
w/o Both	0.6120	0.7800	0.8352	0.5711	0.6824	0.7005	0.3916	0.5382	0.6051

Table 3: Ablation study on encoder-mode training losses. The best and second-best performances are presented in bold and underlined, respectively.

Category	HotpotQA		2WikiMultihopQA		MuSiQue
	EM	F1	EM	F1	EM	F1
ShiftRAG	0.5610	0.6790	0.5080	0.5528	0.3050	0.4129
w/o SST	0.5230	0.6468	0.4630	0.5194	0.2450	0.3639
w/o $\text{LoRA}_{\text{dec}}$	0.5300	0.6570	0.4720	0.5212	0.2650	0.3883
w/o Both	0.5020	0.6253	0.4400	0.4980	0.1970	0.2979

Table 4: Ablation study on decoder-mode training components. The best and second-best performances are presented in bold and underlined, respectively.

		HotpotQA				2WikiMultihopQA
Training method	Parameter	Retrieval		QA		Retrieval		QA		Time (h)
		R@2	R@5	EM	F1	R@2	R@5	EM	F1
Joint learning	shared	0.4896	0.6298	0.1300	0.3430	0.3230	0.4666	0.1270	0.3432	21
Joint learning	separate	0.5264	0.6865	0.4730	0.5885	0.5159	0.6184	0.4360	0.5115	23
Separate learning	separate	0.6598	0.8214	0.5610	0.6790	0.6080	0.7109	0.5080	0.5528	10

Table 5: Performance and training time comparisons across different learning strategies and parameter sharing configurations. The best and second-best scores are bolded and underlined, respectively.

Category	HotpotQA		2WikiMultihopQA		MuSiQue
	EM	F1	EM	F1	EM	F1
ShiftRAG	0.5610	0.6790	0.5080	0.5528	0.3050	0.4129
Textual Evidence	0.5330	0.6574	0.4790	0.5226	0.2650	0.3814
w/o SST	0.5230	0.6468	0.4630	0.5194	0.2450	0.3639

Table 6: Impact of SST-generated soft knowledge tokens on QA performance. The best and second-best performances are presented in bold and underlined, respectively.

Figure 5: $t$ -SNE visualization of anchor token embeddings ( red ), anchor token embeddings + offset embeddings ( green ), and original document token embeddings ( orange ) for each benchmark.

Methods	Inference Latency
	Time (s)	VRAM (GiB)	EM
NV-Embed-v2	8.72	38.82	0.514
GritLM	2.42	20.25	0.443
ShiftRAG	2.79	24.14	0.561

Table 7: End-to-end inference latency, peak VRAM usage, and QA performance on HotpotQA.

Figure 6: Impact of the number of soft knowledge tokens on QA performance.

Appendix figures & tables31 assets

Supplementary material from the paper’s appendix.

Appendix

	1-hop	2-hop	3-hop	4-hop	Total
Examples	2,000	6,000	500	1,500	10,000

Appendix

Table 8: Training-set statistics for decoder mode.

Category	Single-hop QA	Multi-hop QA			Factual QA
	NQ	HotpotQA	2WikiMultihopQA	MuSiQue	SimpleQA
Questions	1,000	1,000	1,000	1,000	1,000
Documents	9,633	9,811	6,119	11,656	20,686

Appendix

Table 9: Dataset statistics (test sets)

Methods	HotpotQA			2WikiMultihopQA			MuSiQue
Methods	R@2	R@5	R@10	R@2	R@5	R@10	R@2	R@5	R@10
Encoder-based Retrieval
Contriever	0.4789	0.6185	0.7214	0.5060	0.6349	0.6386	0.3319	0.4646	0.5648
GTR (T5-base)	0.4913	0.6409	0.7379	0.5154	0.6256	0.6487	0.3314	0.4799	0.5754
ColBERTv2	0.5034	0.6971	0.7561	0.5460	0.6355	0.6565	0.3392	0.4823	0.5859
RaLa	0.5485	0.6800	0.7415	0.5510	0.6563	0.7007	0.3530	0.5047	0.5882

Appendix

Table 10: Complete retrieval performance of ShiftRAG and baselines. The best and second-best performances are presented in bold and underlined, respectively. We conduct five runs with different random seeds and report the mean; the maximum standard deviation is ±0.05.

Methods	HotpotQA		2WikiMultihopQA		MuSiQue
Methods	EM	F1	EM	F1	EM	F1
Encoder-based Retrieval
Contriever	0.3660	0.4436	0.3310	0.3611	0.1990	0.2697
GTR (T5-base)	0.3810	0.4611	0.3450	0.3705	0.2070	0.2804
ColBERTv2	0.4030	0.4829	0.3650	0.3965	0.2190	0.2968
RaLa	0.4460	0.5400	0.4040	0.4397	0.2430	0.3284

Appendix

Table 11: Complete QA performance of ShiftRAG and baselines. The best and second-best performances are presented in bold and underlined, respectively. We conduct five runs with different random seeds and report the mean; the maximum standard deviation is ±0.05.

Methods	Retrieval			QA
	R@2	R@5	R@10	EM	F1
RaLa	0.5923	0.8180	0.9145	0.3010	0.5036
LLM2Vec	0.6006	0.8217	0.9260	0.3150	0.5108
NV-Embed-v2	0.6078	0.8144	0.9318	0.3220	0.5207
OneGen	0.5648	0.7768	0.8746	0.2830	0.4983
GritLM	0.5776	0.7789	0.9090	0.2900	0.5062

Appendix

Table 12: Retrieval and QA performance on NaturalQuestions. The best and second-best performances are presented in bold and underlined, respectively.

Methods	SimpleQA
	EM	F1
RaLa	0.5450	0.6371
OneGen	0.4970	0.5826
GritLM	0.5860	0.7182
ShiftRAG	0.6740	0.8005

Appendix

Table 13: QA performance on the recent factual QA benchmark SimpleQA. The best and second-best performances are presented in bold and underlined, respectively.

MTEB Performance
Task $\rightarrow$	Clust.	STS	Summ.
Metric $\rightarrow$	V-Meas.	Spear.	Spear.
# Dataset $\rightarrow$	11	10	1
LLM2Vec	0.4410	0.8370	0.3105
GritLM	0.5082	0.8303	0.3565
ShiftRAG	0.5110	0.8414	0.3759

Appendix

Table 14: Embedding and generative generalization under the common Mistral-7B backbone. Left: performance of Clustering, STS, and Summarization on the MTEB benchmark. Right: generative performance on MMLU, GSM8K, and TyDiQA. The best and second-best performances are presented in bold and underlined, respectively.

Methods	HotpotQA		2WikiMultihopQA		MuSiQue
	EM	F1	EM	F1	EM	F1
Coconut	0.5160	0.6185	0.4120	0.4404	0.1895	0.2290
CODI	0.5270	0.6254	0.4210	0.4484	0.1970	0.2413
Sim-CoT	0.5360	0.6416	0.4420	0.5146	0.2070	0.2659
ShiftRAG	0.5610	0.6790	0.5080	0.5528	0.3050	0.4129

Appendix

Table 15: QA performance comparison with retrieval-adapted Latent-CoT baselines. The best and second-best performances are presented in bold and underlined, respectively.

Category	HotpotQA		2WikiMultihopQA		MuSiQue
	EM	F1	EM	F1	EM	F1
Ours	0.5610	0.6790	0.5080	0.5506	0.2960	0.4067
w/o relevance weight	0.5570	0.6765	0.5060	0.5480	0.2910	0.4025
w/o low-rank projection layer	0.5520	0.6754	0.4970	0.5421	0.2880	0.3976
w/o anchor tokens	0.5400	0.6533	0.4890	0.5306	0.2800	0.3705
w/o SST	0.5230	0.6468	0.4630	0.5194	0.2450	0.3639

Appendix

Table 16: Ablation study on SST components. The best and second-best performances are presented in bold and underlined, respectively.

Methods	HotpotQA		2WikiMultihopQA		MuSiQue
	EM	F1	EM	F1	EM	F1
1-layer MLP	0.4505	0.5704	0.4008	0.4572	0.2376	0.3339
2-layer MLP	0.4651	0.5894	0.4125	0.4595	0.2446	0.3355
ShiftRAG	0.5610	0.6790	0.5080	0.5528	0.3050	0.4129

Appendix

Table 17: QA performance comparison between SST and simple embedding-to-soft-token projection baselines. All methods use the same question–document representations and number of soft knowledge tokens, with other training and inference settings fixed.

Methods	HotpotQA			2WikiMultihopQA			MuSiQue
	R@2	R@5	R@10	R@2	R@5	R@10	R@2	R@5	R@10
RaLa	0.5485	0.6800	0.7415	0.5510	0.6563	0.7007	0.3530	0.5047	0.5882
ShiftRAG (Uni.)	0.5938	0.7475	0.8190	0.5506	0.6440	0.6944	0.3818	0.5127	0.5909
ShiftRAG (Bi.)	0.6598	0.8214	0.8806	0.6080	0.7109	0.7530	0.4186	0.5634	0.6516

Appendix

Table 18: Effect of causal-mask relaxation on retrieval performance. Uni. and Bi. denote unidirectional and bidirectional attention in encoder mode, respectively. The best and second-best performances are presented in bold and underlined, respectively.

Methods	Time	Data Size
LLM2Vec	5h 30m	96,000
OneGen	14h	174,500
GritLM	48h	2,886,912
ShiftRAG	10h 10m	106,000

Appendix

Table 19: Training latency and data size comparison between ShiftRAG and baselines. Time denotes end-to-end training time measured.

Scoring Methods	HotpotQA				2WikiMultihopQA				MuSiQue
Scoring Methods	R@2	R@5	R@10	Time (s)	R@2	R@5	R@10	Time (s)	R@2	R@5	R@10	Time (s)
cosine	0.6598	0.8214	0.8806	0.93	0.6080	0.7109	0.7530	0.96	0.4186	0.5634	0.6516	0.87
gap-aware	0.6621	0.8275	0.8869	3.17	0.6099	0.7126	0.7571	3.73	0.4188	0.5632	0.6527	3.12

Appendix

Table 20: Effectiveness-efficiency trade-off between cosine similarity and gap-aware relevance scoring. The best retrieval performances are presented in bold.

Methods	Candidate Pool Size
	10,000	20,000	30,000	40,000	50,000	60,000
RaLa	0.4789	0.4574	0.4207	0.3781	0.3351	0.3235
GritLM	0.5890	0.5708	0.5307	0.4824	0.4341	0.3901
NV-Embed-v2	0.6318	0.6165	0.5787	0.5339	0.4904	0.4494
ShiftRAG	0.7530	0.7410	0.6977	0.6516	0.6085	0.5654
Relative Gain over RaLa (%)	57.24	62.00	65.84	72.34	81.59	74.78

Appendix

Table 21: Scalability analysis on R@10 under increasing candidate pool sizes. Relative gain is computed against each baseline at the same candidate size. The best and second-best performances are presented in bold and underlined, respectively.

Category	HotpotQA		2WikiMultihopQA		MuSiQue
Category	EM	F1	EM	F1	EM	F1
Random initialization	0.5530	0.6635	0.4960	0.5412	0.2990	0.4020
Anchor token initialization	0.5610	0.6790	0.5080	0.5528	0.3050	0.4129

Appendix

Table 22: Impact of anchor token initialization on QA performance. The best performances are presented in bold.

Methods	HotpotQA	2WikiMultihopQA	MuSiQue	Correct		Incorrect
Methods	HotpotQA	2WikiMultihopQA	MuSiQue	Entropy	Top-1 Prob.	Entropy	Top-1 Prob.
ShiftRAG	0.5610	0.5080	0.3050	0.4187	0.8343	1.3430	0.5091
Generate w/o SST	0.5490	0.4980	0.2890	0.3869	0.8804	1.0063	0.6762
Train w/o SST	0.5230	0.4630	0.2450	0.2204	0.8831	0.4102	0.8455
Vanilla Backbone	0.5020	0.4400	0.1970	0.1492	0.9186	0.1899	0.8957

Appendix

Table 23: QA performance and confidence statistics under different uses of SST. “Generate w/o SST” removes SST only at inference after SST-based training, whereas “Train w/o SST” trains the generator without SST. The best and second-best performances are presented in bold and underlined, respectively.

Retriever	Inference	HotpotQA		2WikiMultihopQA		MuSiQue
		EM	F1	EM	F1	EM	F1
LLM2Vec	w/o SST	0.5030	0.6106	0.4400	0.4914	0.2430	0.3670
LLM2Vec	w/ SST	0.5220	0.6281	0.4630	0.5202	0.2610	0.3902
NV-Embed-v2	w/o SST	0.5140	0.6220	0.4580	0.5064	0.2480	0.3783
NV-Embed-v2	w/ SST	0.5300	0.6392	0.4690	0.5251	0.2630	0.3913

Appendix

Table 24: QA performance with fixed external retrievers. Both inference settings use decoder-mode LoRA trained with SST and KD, while w/ SST additionally provides soft knowledge tokens at inference. The best performance within each retriever is presented in bold.

		HotpotQA				2WikiMultihopQA
Training Methods	Parameter	Retrieval		Generation		Retrieval		Generation		Time (h)
		R@2	R@5	EM	F1	R@2	R@5	EM	F1
Joint Learning	shared	0.4896	0.6298	0.1300	0.3430	0.3230	0.4666	0.1270	0.3432	21 h
Joint Learning	separate	0.5264	0.6865	0.4730	0.5885	0.5159	0.6184	0.4360	0.5115	23 h
Alternating Learning	shared	0.5166	0.6634	0.2200	0.4036	0.4923	0.5012	0.1860	0.4172	20 h
Alternating Learning	separate	0.5938	0.7665	0.5220	0.6451	0.5654	0.6682	0.4780	0.5409	10 h

Appendix

Table 25: Impact of learning strategies and parameter sharing on retrieval, QA, and training time. The best and second-best performances are presented in bold and underlined, respectively. The Parameter column indicates whether retrieval and generation adaptation parameters are shared or separate.

Figure 7: Average token entropy over the generated answer on MHQA benchmarks, grouped by answer correctness. Red values indicate the entropy gap between incorrect and correct cases.

Category	HotpotQA			2WikiMultihopQA			MuSiQue
Category	EM	F1	Time (ms)	EM	F1	Time (ms)	EM	F1	Time (ms)
Pairwise	0.5630	0.6827	1.65	0.5060	0.5736	1.75	0.2910	0.4032	1.35
Between	0.5580	0.6779	0.87	0.4980	0.5470	1.01	0.2880	0.4008	0.92
Behind	0.5610	0.6790	0.86	0.5080	0.5528	1.02	0.3050	0.4129	0.91

Appendix

Table 26: Effect of different placements of soft knowledge tokens on QA performance and soft-token injection latency in ShiftRAG.

Figure 8: $t$ -SNE visualization of document embeddings in the 2WikiMultihopQA dataset, comparing distributions before and after training under bi-directional and uni-directional attention mechanisms.

Figure 9: Layer-wise covariance between hidden representations of different tokens and effective rank (eRank) of token representations on Gemma2-9B (left) and Mistral-7B (right), respectively.

Figure 10: Layer-wise covariance between hidden representations of different tokens and effective rank (eRank) of token representations on Qwen2.5-7B (left) and LlaMA3.1-8B (right), respectively.

Generator	Inference	HotpotQA		2WikiMultihopQA		MuSiQue
		EM	F1	EM	F1	EM	F1
Gemma4-31B	w/o SST	0.5920	0.7012	0.4960	0.5976	0.3120	0.4358
Gemma4-31B	w/ SST	0.6140	0.7116	0.5280	0.6142	0.3410	0.4873
Qwen3.6-35B	w/o SST	0.6050	0.7293	0.5070	0.5624	0.3290	0.4437
Qwen3.6-35B	w/ SST	0.6310	0.7548	0.5330	0.5788	0.3660	0.5026

Appendix

Table 27: QA performance of large LLM generators with a fixed NV-Embed-v2 retriever. For each generator, both settings use the same SST-trained checkpoint and differ only in whether SST-generated soft knowledge tokens are provided at inference. The best performance within each generator is presented in bold.

Figure 11: Hyperparameter sensitivity analysis on (a) batch size in encoder mode, (b) LoRA rank in encoder mode, (c) $\lambda_{\text{SOA}}$ in encoder mode training, (d) batch size in decoder mode, (e) LoRA rank in decoder mode, (f) $\lambda_{\text{KD}}$ in decoder mode training, (g) temperature $\tau$ in $\mathcal{L}_{\text{LGC}}$ , (h) bucket size in $\mathcal{L}_{\text{LGC}}$ . (i) number of candidate layers in $\mathcal{L}_{\text{SOA}}$ .

Figure 12: Retrieval performance of ShiftRAG and baselines across different LLM backbones.

Figure 13: QA performance of ShiftRAG and baselines across different LLM backbones.

Panel of Figure 14 — Figure 14: A qualitative case study showing that bidirectionally trained retrieval preserves the full evidence chain needed for the answer, including the surface entity cue ( orange ), the father relation ( blue ), and the death-related clue ( red ), whereas unidirectionally trained retrieval overemphasizes surface lexical overlap ( orange ) and retrieves non-supporting documents that miss the relation and death evidence ( blue , red ) required to infer the correct answer.

Panel of Figure 15 — Figure 15: Prompt for generating pair-specific textual rationales.

Panel of Figure 16 — Figure 16: Prompt used for anchor token initialization with textual rationale supervision.

Panel of Figure 17 — Figure 17: Prompt used for generation-mode training with retrieved documents and contextualized tokens.

Retrieval-Augmented Generation

Retrieval-augmented generation connects language models to external knowledge so answers can use current, domain-specific evidence.

See all 817 papers

Latent Abstraction for Retrieval-Augmented Generation
Apr 20, 2026Ha Lan N. T, Minh-Anh Nguyen, Dung D. LeRetrieval-Augmented Generation Adaptive RAG
Center for AI Research, VinUniversity, Vietnam
Guaranteeing Knowledge Integration with Joint Decoding for Retrieval-Augmented Generation
Apr 9, 2026Zhengyi Zhao, Shubo Zhang, Zezhong Wang +7Hallucination in RAG Retrieval-Augmented Generation
The Chinese University of Hong Kong · Ministry of Education Key Laboratory of High Confidence Software Technologies, CUHK · University of International Relations +2
Compression Beyond the Uncompressed: A Two-Stage Training Recipe for Soft Context Compression in RAG
Sep 4, 2026Shuyu Guo, Shuo Zhang, Zhaochun RenRL for RAG LLM Inference Acceleration
Shandong University Qingdao, China · Bloomberg London, United Kingdom · Leiden University Leiden, The Netherlands

Apr 20, 2026J/K move · Enter open · S save

Latent Abstraction for Retrieval-Augmented Generation

Ha Lan N. T, Minh-Anh Nguyen, Dung D. Le

Center for AI Research, VinUniversity, Vietnam

Retrieval-Augmented Generation Adaptive RAG Adaptive Retrieval Dense Retrieval

Open paper AskPDF arXiv

Retrieval-Augmented Generation (RAG) has become a standard approach for enhancing large language models (LLMs) with external knowledge, mitigating hallucinations, and improving factuality. However, existing systems rely on generating natural language queries at each hop and maintaining a strict architectural separation between retriever and generator, preventing them from leveraging the full representational capacity of the LLM. We propose LAnR (Latent Abstraction for RAG), a unified framework in which a single LLM jointly performs encoding, retrieval, and generation entirely within its own latent space. Rather than generating textual queries, LAnR produces dense retrieval vectors from the hidden states of a designated [PRED] token and uses them to match against encoded document representations from the same model. Furthermore, LAnR adaptively decides when sufficient evidence has been retrieved using a lightweight MLP control head over those same hidden states, eliminating both the separate retriever and explicit token-level stopping reasoning. This design is motivated by our empirical observation that answer token entropy reliably signals retrieval sufficiency. Extensive experiments on six QA benchmarks spanning single-hop and multi-hop settings demonstrate that LAnR outperforms existing RAG methods, while achieving improved inference efficiency through reduced number of retrieval calls and tighter model integration.

Abstract

Retrieval-Augmented Generation

Explore similar work

Latent Abstraction for Retrieval-Augmented Generation

Guaranteeing Knowledge Integration with Joint Decoding for Retrieval-Augmented Generation

Compression Beyond the Uncompressed: A Two-Stage Training Recipe for Soft Context Compression in RAG