The paper derives a theory that the effect of retrieval-augmented generation (RAG) can be predicted at token level by comparing the similarity of distributions from the RAG model, the pure LLM, and the retrieved text distribution. This theory was validated on Llama-2-7b, OPT-6.7b, and Mistral-7b across several QA and generation tasks. In practice, the authors implement this theory as Tok-RAG, which performs collaborative generation between the pure LLM and the RAG system at each token. The method is shown to achieve better accuracy and robustness than baselines, especially when retrieved passages contain noisy or irrelevant information.
Evidence
correlational
Key metric
Accuracy on TriviaQA: 85.7 (0% hard negatives), 53.5 (100% hard negatives) with llama-2-7b; AUC on benefit-detriment classification: 81.89 for llama-2-7b, 68.64 for opt-6.7b, 75.85 for mistral-7b on Wikitext
Caveat
The method's effectiveness depends on a hyperparameter for identifying the distribution fusion layer; the theory relies on approximations using KL divergence and may have maximum error that decreases as detriment increases