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

2024-06-03

AI Summary

Retrieval-augmented generation (RAG) faces a fundamental trade-off: external knowledge improves token-level predictions yet risks misleading large language models (LLMs) due to distributional misalignment between the LLM’s prior and retrieved text. Method: We propose the first theoretically grounded, quantifiable framework that models the fusion of these two distributions, proving that distributional divergence is the decisive factor governing RAG effectiveness. Based on this insight, we introduce Tok-RAG—a training-free, token-level dynamic coordination mechanism between pure LLM and RAG outputs—leveraging unsupervised distribution harmonization and information-theoretic dynamic weighting to estimate per-token benefit–risk trade-offs without discriminators, fine-tuning, or post-processing. Contribution/Results: Evaluated on OPT, LLaMA-2, and Mistral, Tok-RAG significantly enhances RAG robustness and accuracy across diverse benchmarks, empirically validating the framework’s interpretability, theoretical soundness, and cross-model generalizability.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Safety and Robustness
• Search and Optimization: Distributed Search

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

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.

Problem

• Explains benefit-detriment trade-off in RAG token prediction.
• Models RAG as fusion of LLM and retrieved text distributions.
• Proposes Tok-RAG for token-level collaborative generation.

Innovation

• Model RAG as knowledge-retrieval fusion
• Formalize benefit-detriment trade-off in RAG
• Propose Tok-RAG for token-level generation