Research map / Large language models / Retrieval and knowledge
A Theory for Token-Level Harmonization in Retrieval-Augmented Generation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Models RAG next-token prediction as a fusion of LLM knowledge and retrieved-text distributions, derives a benefit-minus-detriment theory via KL divergences, and proposes Tok-RAG, a training-free token-level collaborative decoding between pure LLM and RAG.
Problem
RAG can both help (benefit) and mislead (detriment) LLMs, but the token-level trade-off between these effects is not theoretically explained or quantified, and existing fixes need training, utility evaluators, or post-hoc processing.
Approach
Treats RAG next-token prediction as latent-concept distribution fusion, decomposes it into KL-based benefit (distribution completion) and detriment (distribution contradiction), and shows the sign of their difference can be predicted by comparing representation similarities; Tok-RAG uses this to pick between RAG and pure-LLM tokens at each step without training.
Key results
- On short-form QA with 100% hard negative passages, Tok-RAG reaches 53.5 on TriviaQA versus 43.8 for Standard RAG with LLaMA-2-7B, and it leads all baselines at every noise ratio.
- On benefit-detriment comparison with LLaMA-2-7B on Wikitext, Tok-RAG attains AUC 81.89 versus 74.96 for the best baseline (Consistency-Semantic).
- Tok-RAG costs about 6.5 GB more GPU memory than standard RAG (23.98 vs 17.50 GB) and about 0.3 s more per run, using a batch size of 2.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Tok-RAG token-level selection instead of standard RAG | LLaMA-2-7B RAG · TriviaQA, 100% hard negatives | accuracy: 43.8 -> 53.5 | Helped |
| Use Tok-RAG token-level selection instead of the strongest listed baseline | LLaMA-2-7B RAG · TriviaQA, 0% hard negatives | accuracy: 82.5 -> 85.7 | Helped |
| Use Tok-RAG token-level selection instead of the strongest listed baseline | LLaMA-2-7B RAG · WebQ, 0% hard negatives | accuracy: 57.0 -> 57.3 | Helped |
| Use Tok-RAG token-level selection instead of the strongest listed baseline | LLaMA-2-7B RAG · Squad, 0% hard negatives | accuracy: 67.5 -> 68.1 | Helped |
| Use Tok-RAG instead of INFO-RAG | LLaMA-2-7B RAG · Long-form NLP tasks (WoW, Python, Java, T-REx, WikiText-103, ELI5) | task metric reported per dataset: 9.09 -> 9.68; 26.75 -> 27.44; 32.06 -> 32.59; 65.91 -> 68.70; 62.91 -> 64.28; 17.18 -> 17.59 | Helped |
| Use Tok-RAG benefit-detriment comparison instead of the strongest listed baseline | OPT-6.7B · Wikitext token-level benefit-detriment comparison | AUC: 65.93 -> 68.64$^{+}$ | Helped |
| Use Tok-RAG benefit-detriment comparison instead of the strongest listed baseline | Mistral-7B · ASQA token-level benefit-detriment comparison | AUC: 70.05 -> 71.51$^{+}$ | Helped |
| Use Tok-RAG benefit-detriment comparison instead of the strongest listed baseline | OPT-6.7B · Bio token-level benefit-detriment comparison | AUC: 65.76 -> 66.27$^{+}$ | Helped |
Surprising findings
- Standard RAG can underperform the pure LLM at high hard-negative ratios, e.g. 8.6 vs. Tok-RAG 12.8 on SQuAD at 100% noise.
- Tok-RAG beats fine-tuned robust-RAG baselines (RetRobust, INFO-RAG, Self-RAG) without any training or utility evaluator.
Limitations
- The theory relies on approximations (O(1) bounds, assumptions about HMM transition probabilities) and the paper's derivations contain loose steps.
- Tok-RAG requires a hyperparameter threshold a (set to 5e-7) and heuristic layer selection to locate the distribution-fusion layer.
- Experiments are limited to 7B-scale models and a small set of tasks.
Experimental setup
Datasets: Natural Questions TriviaQA WebQ SQuAD Wikitext ASQA Bio Wow T-REx WikiText-103 ELI5 Python code generation
Benchmarks: TriviaQA WebQ SQuAD Wow T-REx WikiText-103 ELI5 ASQA Bio Natural Questions
Metrics: accuracy F1-score AUC CodeBLEU ROUGE GPU memory time
Baselines: Standard RAG NLI+RAG CRAG RetRobust Self-RAG INFO-RAG Logprobs Uncertainty Consistency-Lexical Consistency-Semantic
Builds on: in-context learning as latent variable inference hidden Markov model Self-RAG CRAG
Closest papers by meaning
- One Token Can Help! Learning Scalable and Pluggable Virtual Tokens for Retrieval-Augmented Large Language Models (AAAI 2025 · similarity 0.92)
- RAGGED: Towards Informed Design of Scalable and Stable RAG Systems (ICML 2025 · similarity 0.91)
- Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG (ICLR 2025 · similarity 0.91)
- Speculative RAG: Enhancing Retrieval Augmented Generation through Drafting (ICLR 2025 · similarity 0.91)
- Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection (ICLR 2024 · similarity 0.90)
- SePer: Measure Retrieval Utility Through The Lens Of Semantic Perplexity Reduction (ICLR 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.