# A Theory for Token-Level Harmonization in Retrieval-Augmented Generation

arXiv: https://arxiv.org/abs/2406.00944 · HTML: https://mlsanity.com/paper/2406.00944/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in RAG pipeline design and evaluation, one of 15 directions in Retrieval-augmented generation.
- Builds on 46 of its 59 references that are in the map.
- First posted on arXiv 3 Jun 2024.

## Placed in

- Retrieval-augmented generation / [RAG pipeline design and evaluation](https://mlsanity.com/rag/rag-pipeline-design-and-evaluation/)

## Builds on (most cited first, 30 of 46)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [FActScore: Fine-grained Atomic Evaluation of Factual Precision in Long Form Text Generation](https://mlsanity.com/paper/2305.14251/) (2023)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [BERT Rediscovers the Classical NLP Pipeline](https://mlsanity.com/paper/1905.05950/) (2019)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation](https://mlsanity.com/paper/2302.09664/) (2023)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [In-Context Retrieval-Augmented Language Models](https://mlsanity.com/paper/2302.00083/) (2023)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [ELI5: Long Form Question Answering](https://mlsanity.com/paper/1907.09190/) (2019)
- [KILT: a Benchmark for Knowledge Intensive Language Tasks](https://mlsanity.com/paper/2009.02252/) (2020)
- [Wizard of Wikipedia: Knowledge-Powered Conversational agents](https://mlsanity.com/paper/1811.01241/) (2018)
- [Active Retrieval Augmented Generation](https://mlsanity.com/paper/2305.06983/) (2023)
- [CodeBLEU: a Method for Automatic Evaluation of Code Synthesis](https://mlsanity.com/paper/2009.10297/) (2020)
- [ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction](https://mlsanity.com/paper/2112.01488/) (2021)
- [DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models](https://mlsanity.com/paper/2309.03883/) (2023)
- [Making Retrieval-Augmented Language Models Robust to Irrelevant Context](https://mlsanity.com/paper/2310.01558/) (2023)
- [Adaptive Chameleon or Stubborn Sloth: Revealing the Behavior of Large Language Models in Knowledge Conflicts](https://mlsanity.com/paper/2305.13300/) (2023)
- [INSIDE: LLMs' Internal States Retain the Power of Hallucination Detection](https://mlsanity.com/paper/2402.03744/) (2024)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)

## Built on by (most cited first, 7 of 7)

- [Retrieval-Augmented Generation with Graphs (GraphRAG)](https://mlsanity.com/paper/2501.00309/) (2024)
- [How Much Can RAG Help the Reasoning of LLM?](https://mlsanity.com/paper/2410.02338/) (2024)
- [Cross-Modal Safety Mechanism Transfer in Large Vision-Language Models](https://mlsanity.com/paper/2410.12662/) (2024)
- [RLKD: Distilling LLMs' Reasoning via Reinforcement Learning](https://mlsanity.com/paper/2505.16142/) (2025)
- [Predicting Retrieval Utility and Answer Quality in Retrieval-Augmented Generation](https://mlsanity.com/paper/2601.14546/) (2026)
- [Reverse Physician-AI Relationship: Full-process Clinical Diagnosis Driven by a Large Language Model](https://mlsanity.com/paper/2508.10492/) (2025)
- [When Does Context Help? Error Dynamics of Contextual Information in Large Language Models](https://mlsanity.com/paper/2602.08294/) (2026)

## Abstract (arXiv)

Retrieval-augmented generation (RAG) utilizes retrieved texts to enhance large language models (LLMs). Studies show that while RAG provides valuable external information (benefit), it may also mislead LLMs (detriment) with noisy or incorrect retrieved texts. Although many existing methods attempt to preserve benefit and avoid detriment, they lack a theoretical explanation for RAG. The benefit and detriment in the next token prediction of RAG remain a black box that cannot be quantified or compared in an explainable manner, so existing methods are data-driven, need additional utility evaluators or post-hoc. This paper takes the first step towards providing a theory to explain and trade off the benefit and detriment in RAG. First, we model RAG as the fusion between distribution of LLMs knowledge and distribution of retrieved texts. Then, we formalize the trade-off between the value of external knowledge (benefit) and its potential risk of misleading LLMs (detriment) in next token prediction of RAG by distribution difference in this fusion. Finally, we prove that the actual effect of RAG on the token, which is the comparison between benefit and detriment, can be predicted without any training or accessing the utility of retrieval. Based on our theory, we propose a practical novel method, Tok-RAG, which achieves collaborative generation between the pure LLM and RAG at token level to preserve benefit and avoid detriment. Experiments in real-world tasks using LLMs such as OPT, LLaMA-2, and Mistral show the effectiveness of our method and support our theoretical findings.