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

TL;DR

This paper introduces a theoretical framework for understanding and balancing the benefit and detriment of retrieval-augmented generation (RAG) in language models, enabling prediction and optimization without additional training.

Contribution

It provides the first formal theory modeling RAG as a fusion of distributions and predicts RAG effects without training, leading to a novel token-level harmonization method called Tok-RAG.

Findings

Tok-RAG effectively balances benefit and detriment in experiments.

Theoretical predictions align with empirical results.

Method improves generation quality in real-world tasks.

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.