A Theory for Token-Level Harmonization in Retrieval-Augmented Generation
Authors: Shicheng Xu
Liang Pang
Huawei Shen
Xueqi Cheng
Venue: ICLR 2025 Poster (2025)
Keywords: Retrieval-augmented generation
Duality
Large Language Models
Primary area: applications to computer vision, audio, language, and other modalities
TL;DR: This paper provides the essential understanding of benefit and detriment in RAG to make them interpretable, quantifiable, and comparable.

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. We model RAG as the fusion between distributions of LLMs’ knowledge and distributions 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. Code is in supplemental material and will be released on GitHub after acceptance.

Published as a conference paper at ICLR 2025
A THEORY FOR TOKEN-LEVEL HARMONIZATION IN RETRIEVAL-AUGMENTER GENERATION
Shicheng Xu Liang Pang'; Huawei Shen', Xueqi Cheng'
ICAS Key Laboratory of Al Safety, Institute of Computing Technology, Chinese Academy of Sciences University of Chinese Academy of Sciences (xushi cheng2s pang i ang, shenhuawei, cxq) eict.ac...
ABSTRACT
Retrieval-augmenned generation (RAG) utilizes retrieved texts enhance large language models (L..s). Studies show that while RAG provides valuable external information (benefit), it may also mislead LLMs (detriment) with noisy 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 an explainable manner, so existing methods are data- driven, need additional utility evaluators or post-hoc. This paper takes the first step towards providing theory to explain and trade off the benefit and detriment RAG. First, we model RAG as the fusion between distribution of LLM's 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, 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 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 L.MSs such as OPT, LLaMA-2, and Mistral show the effectiveness of our method and support our theoretical findings. Code is available")
INTRODUCTION
Retrieval-augmented generation (RAG) has shown promising performance enhancing Large Language Models (LLMs) by integrating retrieved texts (Xu et al. [023] Shhi et al., 2023 Asai et al.)
2023 Ram et al. 2023). Studies indicate that while RAG provides LLMs with valuable additional
knowledge (benefit), also poses risk of misleading them (detriment) due noisy incorrect retrieved texts (Ram et al. 2023 Xu et al. 2024b a Jin et al., 2024 Xie et al. 2023). Existing methods attempt preserve benefit and avoid detriment by adding utility evaluators for retrieval, prompt enginecring, or fine-tuning LLMs (Asai et al. [223 Ding et al. 2024 [uu et al. 2024b Yoran al. 2024, Ren al., 2023 Mallen et al. [022]] Jiang et 2023). However, existing methods are data-driven, need evaluator for utility of retrieved texts post-hoc. A theory-based method, focusing on core principles of RAG is urgently needed, which is crucial for reliable improvements without relying on additional training utility evaluators and improving understanding for RAG.
This paper takes the first step in providing a theoretical framework to explain and trade off the benefit and detriment at token level in RAG and proposes novel method to preserve benefit and avoid detriment based on our theoretical findings. Specifically,, this paper pioneers in modeling next token prediction in RAG the fusion between the distribution of LLM's knowledge and the distribution of retrieved texts as shown in Figure [ Our theoretical derivation based on this formalizes the core of this fusion the subtraction between two terms measured by the distribution difference: one is distribution completion and the other is distribution contradiction. Further analysis indicates that
"Corresooningg Author https: thub. com http///cl1233 //334//ckk-aa

---

Published as a conference paper at ICLR 2025
Pure LLM:
Query: Who the first Nigerian win the Nobel Prize, in which year?
Query
Emst
Clanipdcun
1986
LLM
LLMY Dianibtion
Datribatio
Ditritation
Datrimment Win
RAG: Retrieved Texts
Blaribatien
X
Mierhariio
X
Distribution Differencs
Rontrval Douittttan
Differcney
Bonelli Wila
Botrieved Datrization
Differenes
Henefir Win
Retreved Diantuttio
Query
Welc
Soylnka
1950
Dianibztion
Disuiration
LLM
Figure I: Framework of our Tok-RAG. It performs collaborative generation between pure LLM and RAG at the token-level by comparing benefit and detriment based on theoretical findings about distribution difference. The selected tokens at each step are used as the prefix for both pure LLM and RAG. Tok-RAG preserves benefit and avoids detriment without any training utility evaluators.
the distribution completion measures how much outto--distribution knowledge that retrieved texts provide to LLM in the prediction of the next token, representing the benefit of RAG for LLM at the token level. However, since the retrieved texts may contain noisy or incorrect information, posing the risk of misleading the LLM, distribution contradiction captures this risk by measuring the degree of conflict between the LLM's knowledge and the external knowledge from the retrieved texts. This represents the detriment of RAG for LLM at token level. Thus, we prove that the fusion between distribution of L.M's knowledge and the retrieved texts in next token prediction of RAG is governed by the combined effect of benefit and detriment, whose relationship is described subtraction.
In this way, we successfully decouple benefit and detriment from next token prediction of RAG and formalize them the subtraction between terms. This subtraction describes the trade-off between the value of external knowledge and its potential risk of misleading LLM without accessing the utility of the retrieved texts. We then prove that value of this subtraction is approximately positively correlated with the similarity between the representation of RAG's 's output and the representation of retrieval texts. Finally, we establish a theory to predict the actual effect of RAG token level. Specifically, in next token prediction given prefix and retrieved texts, this theory determines whether the benefit brought by retrieved texts on the prediction the token outweighs the detriment.
Based on our theoretical results, we propose practical novel method called Tok-RAG that can achieve collaborative generation between pure LLM and RAG token level to preserve benefit and avoid detriment without any training or additional modules for utility evaluation of retrieved texts. As shown in Figure pure LLM and RAG generate the texts in parallel. At the generation step where LLM and RAG generate the different tokens, Tok-RAG uses our theorctical results to determine which token will be selected by comparing the values of benefit and detriment brought by RAG to the token. Experimental results in real-world tasks such as Q&A and Long-Form Q&A based on LLMs including OPT, LLaMA-2, and Mistral show the effectiveness of our method and support our theoretical results. Our method does not need any additional modules or training but outperforms baselines that need additional modules and fine-tuning LLMs by just accessing the layers and logits in inference, which indicates that our theoretical results are essential and fundamental for RAG. The main contributions of this paper are:
This paper takes the first step in theoretically giving the essential explanation of benefit and detriment in RAG to make them explainable and comparable at teken-level.
We provide a theory that actual effect of RAG (i.e., the comparison between benefit and detriment) can be predicted at token level by representation similarity without any training or accessing the utility of the retrieved texts, which significant for ene-ggaieed preserving benefit and avoiding detriment in practical applications of RAG.
Based on the theoretical results, propose practical novel method to enable pure LLM and RAG to collaboratively generate at token level. Experimental results on real-world tasks across different I L.Mss show the effectiveness of our method and support our theoretical results.
2

---

Published as a conference paper at ICLR 2025
Equation 9
CaCalleye))) > (( X())
Cacollary Actool of RAG can proslicied idlenséesr
Conalirr::: Divaribation lifference
Conatiary Blogght
formralizes beweltt and deriissee R46G
relansffoo epresisiation similariy
Actual Effect
Actual Effect Benefit
Decrimeut
Rotiesal
Differeae
Qusppiiee: Beacil
@eBnncli, Detriment)
Upper Bound
Theoren 2,3
Sind OrSiiiO
Beseff - Detiitent
Hunofii Datnmani
1.2
Theoren Bound
Sinilacie
②
Lower Bound
LLM's Duniinatia
Fraice
Dasibyiic Gommddttoe:
Theeeum 2: Punitively corritated
Dariment Similanty
O(Benifi Deriment)
Simi S
in Benelli Dovimens
Sindl inalarty repoctstare perilice Suar)
Figure 2: Derivation path of our theory. Reference: Equation Theorem
I
annd Corollary 3
2 THEORY TO UNDERSTAND BENEFIT AND DETRIMENT AT TOKEN LEVEL
Although retrieved texts can provide LLM with external knowledge (benefit), they also pose a risk of misleading LLM due to the noise within the retrieved content (detriment). This section provides the theoretical framework to explain and trade off benefit and detriment in RAG at token level. First, we propose use latent variable model to model the next token prediction in RAG the fusion between the distribution of L.M''s knowledge and the distribution of retrieved texts. Second, our theoretical derivation formalizes the core of this fusion the subtraction between two terms measured by the distribution difference: one is distribution completion, and the other is distribution contradiction. Our further discussion shows that distribution completion is actually the benefit and distribution contradiction is detriment in the next token prediction of RAG. Last but not least, we provide a theory to predict the actual effect of RAG at token level. Figure shows the derivation path of our theory.
2.1 PRIMARY DEFINITIONS AND ASSUMPTIONS
Retrieved Texts List R. R = (ri, T3...... I' 3 is the list that contains retrieved passages for RAG, in which r; is passage. In textual format of R, T:-1 and r; is separated by a delimiter such as "[Retrieved Passage]". denoted as d.
The Distribution of Retrieved Texts List. pn(-) ) is the distribution of retrieved texts list R. It the language modeling distribution just like the distribution p(-) of the LLM. This distribution p(.) is learned through training LLM on large corpus with next-token prediction paradigm. p(2;|zi5--) represents the probability distribution predicted by the LLM for æ;, given the prefix F1--1-- The LLM capable of generating reasonable and coherent text based on the distribution patterns of vocabulary and context learned from the large training data. Here, p+(+) refers to the distribution p(-) learned when the training data that is limited to the retrieved passages list It and it represents the vocabulary and contextual distribution patterns of the retrieved passages list R. Define r; the i-th passage in the retrieved passages list R, so PR(r;) represents the joint probability of the natural language sequence r; = [r! r2, r.? ..., r under the distribution (-()). in which rf is the æ-th token in passage r. Specifically, according to the chain rule, pn(r;) can be decomposed into the product of a series of conditional probabilities:
(R( = -P(T!) pR(F?(r!) pR(rir!..))) PR(rrir
Latent Variable Inference as Hidden Markov Model. Inspired by previous studies that prove LLMs implicitly perform latent variable inference (Zhang et al. 2023 Wang et al. 2024), we first propose to analyze RAG by latent variable inference like Hidden Markov Model (HMM) Xie et al.) (202I), where the latent concept = determines the transition probability matrix between hidden states h, and both hidden states and latent concepts together determine the prediction of the token.
Assumptions. Based on previous work (Xie et al. 2021 Zhang et al. 2023), we make the following assumptions that:
Assumption 1.
All tokens can be predicted, which means that for every token x, there is some hidden state h lower bounds it that p( > 01 > O.
Assumption 2. Delimiter d is an important distinguishing signal between each passage in the retrieved texts R. The delimiter hidden state hd implies that p(d h² ) = 1, meaning that the hidden state he uniquely determines the output d when 2 is given, making the probability of observing equal to 1. For any delimiter hidden state h and other hidden state h, there are upper and lower bounds on the transition probability from h to hd: c2 S 1h, S cg.
3

---

Published as a conference paper at ICLR 2025
2.2 MODEL NEXT TOKEN PREDICTION RAG AS DISTRIBUTION FUSION
We propose to analyze RAG by latent variable inference to model the next token prediction in RAG the fusion between the distribution of LLM's knowledge and the distribution of retrieved texts. Given the prefix T1:1-1 = "01,42,...53
from the perspective of the latent variable inference, the probability distribution of the token 24 at the i-th step described as this:
p(Z)(T1--1) / [21]/111-11 2)p(22/11111))) dz,
(I)
in which 2 is the space of high dimensional concept variable, p(=(a11--1) is the probability that the model samples latent concept from 2 given the prefix 11 and ( means the probability for token I; conditioned on the prefix I1:i-1 and the sampled latent concept z. p(x)(221-)1) can be obtained by integrating over all latent concepts from the space Z. Latent variable model has been applied in many methods such as LDA (Blei et al., 2003). Recent studies prove that in-context learning of LLMs can also be seen as the latent variable model, in which the LLMs sample the concept across the input examples (Xie et al.) 2021 Zhang et al. (023). Inspired by this, we analyze the next token prediction of RAG given prefix T1si-1 as sampling the shared Retrieved Concept 2" from the input retrieved texts list R = (f..72....,7)) Tn) (r, is a retrieved passage), and then predicting p(a; RR x1s-1), which can be formalized as:
p(z, RR 11 p(z, 111- R1. dz
(2)
= =( R, (1: ",(=
21,1,,
x1--1 dz.
2-1=*)
Equation describes RAG as distribution fusion. The first term is the prediction that is only conditioned on 2", which is the distribution from retrieved texts. The second term is the prediction that marginalizes out all latent concepts except 2", which is the distribution in LLMs.
2.3 FORMALIZE AND EXPLAIN BENEFIT AND DETRIMENT BY DISTRIBUTION DIFFERENCE
In this section, we make derivation on the distribution fusion in Equation 2 to formalize the core of this fusion as the subtraction between two terms measured by distribution difference: one is distribution completion, and the other is distribution contradiction. We reveal that in the next token prediction of RAG, distribution completion is the benefit and distribution contradiction is the detriment. Specifically, inspired by (Xie et al. 2021), Equation 2 is transformed as (see Appendix
p( (, = p(r, R, ()) (2) x11-1) dz
(3)
() (R, dz
(4)
dz, v(z) = log P(R, x F1:1 11:1-112)
1|2)
(5)
p(R,, x 2
Define r; is passage in the retrieved texts list R. we get (see detailed proof in Appendix p( R2 12))
B
~ log I - ((1)p(r(12)
v(2) log p(R, 31 2*)
(6)
(()pp(r,(**)
log p(ra(2) = * BA-PA
log p(r[z)
p(ri(2+)
p(r(2*)
x pn(r))og p(rlz) = pre(r)log PR(F) p(rl2*) = =-( (I((p(()))((ree)
pre(r log Pr((r)
(8)
p(r(2*)
p(r(2)
Distribution Completion: Benefit
Distribution Contradiction: Detriment
pr(()) is the distribution of the retrieved texts, p(-) is the distribution of the LLM's knowledge. v(2) is an important term in distribution fusion because reflects the proportion between the latent concept from the space of LLMs and from the retrieved texts. Details are in Appendix [
Discuss the Benefit and Detriment Based on the Theoretical Results. In Equation 9 the first term represents the distribution difference (KI. divergence) between retrieved texts (p)(r) and LLM's
4

---

Published as a conference paper at ICLR 2025
knowledge (p(rlz)) given the concept 2, which is sampled from 2 (latent variables in LLM). This can be defined as distribution completion that measures how much of---iiitiiuution knowledge that retrieved texts provide to LLM in the prediction of the token 2ę. This term is actually the benefit for the prediction of token x; in RAG. This is because that if the retrieved text r is perfect, pr((r) would be infinitely close to the ground-truth distribution. The larger the difference between Prr(r) and p(r[2). the more the knowledge distribution about z; in the LLM's knowledge deviates from the ground truth. This means the retrieved texts can provide more valuable out-of-distribttion knowledge to the LLM to predict 2;, resulting in greater benefit.
Considering that the retrieved texts are not always perfect and may contain incorrect information and noise that contradict the correct knowledge of LLM, the second term corresponds to the dis- tribution contradiction can used to measure this risk. represents the distribution difference (KI. divergence) between the retrieved texts (pn(r)) and LLM's s knowledge given the concept =", which is sampled from retrieved texts (p(rl=")) )). p(r(=") is the prediction made by L.LM conditioned on the concept 2" sampled from the retrieved texts. If the external knowledge in the retrieved texts contradicts LLM's knowledge, p(r[=") will deviate from Pa(r) (the actual distribution of the re- trieved texts). Therefore, the difference between p(r)2") and pa(r) primarily stems from the LLM's resistance to any external knowledge in the retrieved texts that conflicts with LLM's knowledge. The larger difference indicates the stronger resistance from LLM,, and the more confident the LLM is in its pre-trained knowledge, which means the greater the potential detriment caused by the retrieved texts. So this term is actually the detriment for the prediction of token I, in RAG.
Two terms about distribution difference in Equation Inmatiiio@@meee measure the benefit and detriment respectively. The subtraction between benefit and detriment describes the trade-off relationship between the value of external knowledge and its potential risk of misleading LLM in the next token prediction without accessing the utility of retrieved texts.
Corollary 1.
2.4 ACTUAL EFFECT OF RAG CAN BE PREDICTED AT TOKEN LEVEL
Based above analysis, we successfully formalize the benefit and detriment in next token prediction of RAG by measuring distribution difference. Next, we further explore the method to compare the values of benefit and detriment. Specifically, we derive Theorem and Theorem [2] from Equation
Theorem 1. Define D = ||p(2, [R, 214-1) PP(x,,
la to measure the difference between the distribution of RAG xp(xz: R. x1s-1)) and the distribution of retrieved texts (pR(2) (1a-11) in token prediction of x; conditioned on prefix F1:--1- Both benefit and detriment are important terms of the upper and lower bounds of D, which can be described as:
| V2KL(pa(r)ll(r(z>) S D < |1 + (2KL(pa(7)ll((rz>)
(10)
R ① x a 2-(2-
(KL(pR(F)lP((Ia) KL(pR(r)
p(=) dz,
benefit
detriment
in which is a constant. Our detailed proof of Theorem can be found in Appendix D
Theorem 2. D is the difference, so to can be treated as similarity between p(a,) 11a-1) and (1)-)) The result of benefüt minus detriment is approximately positively correlated with D:
(lppa())ll((rat) (L(pz(()ll((((=)) & D
(11)
benefit
detriment
Our proof of Theorem 2 and the maximum error analysis of this approximation is in Appendix D
Corollary 2. The difference between values of benefit and detriment in Equation 17 indicates the extent to which the benefit (value of external knowledge) outweighs the detriment (potential risk of misleading LLM) in the prediction of token I;, This difference is approximately positively correlated with the representation similarity, which is the value that can be predicted.
Recapping our motivation that aims to build a theory to predict whether the positive impact of the retrieved texts R on (benefit) outweighs the potential risk of misleading the LLM (detriment). The key for this is comparing the values of benefit and detriment. As stated in Theorem 2 the result of benefit minus detriment is approximately positively correlated with 1. Therefore, the value of /) at which the benefit minus detriment equals zero is important threshold, A D value greater than
5

---

Published as a conference paper at ICLR 2025
Lig
Figure 3: Attention score for z; (blue line) and difference of word distribution change (yellow line) vary with layers. stage I: Lexical and Syntactic. stage 2: Text Matching, stage 3: Distribution Fusion.
this threshold indicates that benefit exceeds detriment, while a value less than this threshold indicates that detriment outweighs benefit. We derive Theorem33to identify this dividing point and map the value order between benefit and detriment of token 2; to the relationship between representation similarity, which can be calculated in practical applications:
Theorem 3. Define M [|p(z,) æ1á-1) (102:/1aa--il to measure the difference between distribution of RAG (p(z:) R, 215-1)) and pure LLM (p(=)( so M can be treated as the similarity between them. 1
D and 1 M is the dividing point in which benefit is equal to detriment, and the value order between 1
D
. can indicate the value order between benefit and detriment as:
(L(Pa(r))lp(rt) < KL(pa(r)lle(rl=")) detriment outweighs benefit. = (L(pn(r)llp((ra)) KL(pn(r) (((p(())lp(((22")) (p((((=)))))) detriment is equal to benefit. (L(pa(r))lp(rt > (r(p)(r)ll((re>")) benefit outweighs detriment.
if if D M
(12)
if D
M
Our detailed proof of Theorem 3 can be found in Appendix [F] Equation 12 is novel principle that can compare the values of benefit and detriment RAG at token level. It does not rely on additional modules access the utility of retrieved texts training but simply compares D and x
The actual effect of RAG, i.e., the comparison between benefit and detriment, can be predicted at token level by the similarity relationships between p(r: R, æra-1), ((1s())) and pre (21:--1), without additional modules access the utility of retrieved texts.
Corollary 3.
The next section introduces how to apply above theoretical findings to improve RAG practice.
3 TOK-RAG: IMPROVE RAG BASED ON TOKEN-LEVEL THEORY
Tok-RAG is a novel method that enables the LLM and RAG to collaborate at token level for generation to preserve benefit and avoid detriment based on our token-level theory. Tok-RAG makes pure LLM and RAG generate in parallel at token level as shown in Figure T
It determines which token will be selected by comparing the values of benefit and detriment brought by RAG to the token according to the size relationship between 1 and in Equation 12 The terms related to the comparison between I and 1 consist of three parts: (1) p(x;)/t #1:i-1 1) can be directly obtained from the prediction of RAG; (2) p(r; (x i can be directly obtained from the prediction of pure LLM; (3) however, the distribution of retrieved texts conditioned on the prefix T1:i-1- (R(x: (11511)))) is hard directly obtained, which is the main challenge that the following Section 3.1 aims to solve.
3.1 DISTRIBUTION PREDICTION FOR RETRIEVED TEXTS
Our theoretical analysis in Section 2.2 shows that RAG can be modeled as fusing the distribution from retrieved texts with LLMs' distribution. Therefore, an intuitive idea is that the retrieved distribution (11 can be approximately predicted by capturing the signal from the retrieved texts in distribution fusion. The main challenges in achieving this are: (1) determining where distribution fusion occurs, and (2) capturing the signal fused from retrieved texts and transforming it to distribution pr( (z, (1)- To address these challenges, in the following parts, we first explore the operating mechanism of RAG and then propose novel method that dynamically determines the layers where distribution fusion occurs and use the signal from retrieved texts in these layers as Pa(#; pr(2,/220 (1)
Exploring the Mechanism of RAG. This part finds that the mechanism of RAG can be decomposed into two parts. The first is text matching, which means extracting information relevant to the generation of z; from the retrieved texts R. The second is distribution fusion, which means fusing

---

Published as a conference paper at ICLR 2025
the distribution from the retrieved texts with the distribution in LLM's knowledge. When performing RAG, LLMs first perform text matching in the middle layers, extracting relevant knowledge from the retrieved texts. As the depth increases, the matching becomes increasingly accurate, reaching a turning point. In the deep layers after this turning point, LLMs instead carry out distribution fusion, and the attention shifts from R to T1:4-1- Distribution of R used for fusion comes from the matching information around the turning point (because matching decreases after the turning point). Recapping the two challenges introduced the beginning of Section 3.1 for the first challenge, we identify the layer where distribution fusion starts by detecting the turning point in Figure [] For the second challenge, we use the matching information in the layer where distribution fusion starts to approximate the distribution pn (n(#:: [1si-1). Our experiments about these findings are introduced in Exp 1. and Exp 2.. Experiments are conducted based on LLaMA-2-77 on Natural Question dataset. We also perform these experiments using more LLMs on more datasets. The conclusions of them are consistent with LLaMA-2-77 - on Natural Question, and detailed results can be found in Appendix H
Exp 1. For text matching, we quantify the relevance of the information in the retrieved texts to the generation of token T; by the attention score between token d; and the tokens in the retrieved texts R. We analyze how the sum of attention scores from token z; to tokens in R varies across layer. As shown by the blue line in Figure 3 (1) The value increases sharply to a peak in shallow layers (0-5), which is mainly because LLMs capture the low-level lexical and syntactic information on the entire input (Tenney et al. 2019). (2) The value first decreases, then increases to a maximum point in the middle layers (5-23), which mainly because LLMs select the relevant semantics that can be used to generate x; from Ht and complete this selection at the maximum point. (3) The value decreases after the maximum point in deep layers (24-32). This is because that LLMs use the selected knowledge at the maximum point for distribution fusion to predict F;, the attention shifts from R to prefix 21:i-1-
Exp For distribution fusion, since distribution fusion is often marked by change in word distribution (Bengio et al. (2000), we identify the occurrence of distribution fusion in RAG by comparing the change in word distribution between pure L.LM and RAG. (Chuang et al. 2023 Schuster et al.] 2022) prove the language heads can be directly applied to the hidden states, so we propose to obtain the word distribution of hidden states in each layer by language heads o as (lh)), in which hy is the hidden states for token v; in the 1-th layer. We then measure the word distribution change in the l-th layer by Jensen- Shasnonn Divergence (ISD) between d(h(-1) and ((h)) as: C = (SD(o(h) (S)( (k()). The difference of word distribution change between pure LLM and RAG in the 1-th layer is described as:
D' = JJSD ((())))l((())))) - (()(((())))))((())))
(13)
in which 8/-1 and ht are from RAG, h!-1 and h² are from pure LLM. The yellow line in Figure shows D is very small in the shallow and middle layers (0-23) and rises sharply in the deep layers (24-32). This suggests that distribution fusion occurs in deep layers.
Dynamically Identify the Layer Where Distribution Fusion Starts. For p(z; R, -1.--1,, the layer where distribution fusion starts can be located by detecting the turning point in Figure 3 Specifically, we use (() to denote the attention score for I; and g() to denote the difference of word distribution change in Equation 13 varies with layer 1. The layer where distribution fusion starts is:
1' 11=1, (arg max f() + min(l: g()) > a))]-
(14)
The first term is the that maximizes f(), which is the third turning point in the blue line of Figure The second term means the minimum / value for which g() is greater than a (hyperparameter, can be set to 5e-7 according our statistics), which is the turning point in the yellow line of Figure We take the average of the two values and round down as the layer 1" where distribution fusion starts.
3
Matching Distribution. The matching information between R [rti, ..., rtm] (rt is the token in R2) and token x; at turning point (the -" layer) can be used to approximate the distribution (R 11) of the retrieved texts R conditioned on x1 The matching information consists of two parts, one is the attention score, which can measure the matching between retrieved tokens and current token 2y at the hidden state level. The other is the similarity of word embeddings, which can measure the matching between retrieved tokens and current token x; at the word distribution level:
(nf W (R)( Wx)
WordSim softmax (a=p-"" A)(ttfimm))))) (15)
Att softmax
/ /dp
7

---

Published as a conference paper at ICLR 2025
W3 and Wi are matrices in attention Vaswani et al. 2017), Kf is the hidden state of token 2; and fr.m are hidden states of R. A is word embedding matrix, x is the token with the largest logits increase in word distribution from layer to the final layer 1", rtt... are tokens in R. Then:
PR( (1:-1) = softmax (Att @ WordSim). @ is element-wise multiplication.
(16)
3.2 TOKEN-LEVEL COMPARISON BETWEEN BENEFIT AND DETRIMENT IN PRACTICE
Equation 12 shows that the relationship between Sim(p( R, (P.--1),PR((0) (1111)))111))))) and Sim(p(r, R, 11:;-1); ((=r)()) indicates the value order between benefit and detriment (Sim(..-) the similarity). We propose use the token semantics as the representation for (1s-) ( (1::-1) and p(z; (1:-)) and use cosine to compute the similarity. It not only follows the principle of Equation| but also takes into account the semantic similarity, which is more robust in practical applications. Specifically, we use word embedding matrix of LLMS to calculate the weighted average word embedding for p(z; R, 315-1) as W WRAG = Er E(p .weV p' /w, for each token in E(pr.)ev
vocabulary V, p' is its logits from p( RR 11) and w is its word embedding. We can also use this to get the weighted average word embedding WALM for (( and WIR for The similarity between them can be calculated via cosine similarity as:
Sim(p(x;) RR -1 ( ()) = cos( w WRAG, WTR)
(17)
Sim(p(r;) R2 -1. ,) ((.-)) = cos cos(WRAG. WLLM)
(18)
Combining our theoretical analysis of Theorem T 2 and B we can derive this principle to compare the values of benefit and detriment brought by RAG to the token x; in practical applications:
benefit win if cos(W RAG, WIR) coos CAAG, WLLM)),
RAG,
(19)
detriment win if cos(W WIR) < cos RAGG WLLM),
Given the prefix -1:1-1. 1. RAG generates the token 2; and LLM w/o RAG generates x/. We compare the value of benefit and detriment by Eqn [9. If benefit wins, 2; is selected, otherwise, x is selected. The selected token will be concatenated with -1ai 1 as the new prefix for next step generation.
4 EXPERIMENTS
4.1 EXPERIMENTAL DETAILS
One experiment is in RAG setting for short-form Q&A given retrieved texts with different qualities, evaluates the robustness and performance of RAG. The other is RAG for many long- long-form text generation tasks including dialogue, code generation, slot filling, language modeling and long-form Q&A. Baselines include the methods that use additional modules to filter irrelevant texts (NLI+RAG (Yoran et al. (024))
Practical RAG Performance.
action triggers (CRAG (Yan et al. 2024
fine-tune LLMs for robust RAG (RetRobust (Yoran et al., 2024 and INFO-RAG (Xu et al.
2024b
and fine-tune LLMs dynamically retrieve and critique retrieved texts (Self-RAG (Asai
2023
Setup for Genefit-Detrimen Comparison Experiment. Given prefix T1:1-1 and retrieved texts R. our motivation aims to build a theory to predict whether the positive impact of the retrieved texts R on 2; (benefit) outweighs the potential risk of misleading LLM (detriment). This is binary classification task at leven- To evaluate this, we construct test data and ground-truth as:
(1) For a sentence x, we truncate it at the i-th token to obtain the prefix I1 and the next token i;, (2) Input prefix 21:-- to LLM w/o RAG and LLM w/ RAG to get the predicted token a and b.
(3) If b is x; but not, it means LLM w/ RAG performs better than LLM w/o RAG, so the benefit is
greater than detriment, the ground -ruth label is 1.
(4) If a is I2 but not, it means LLM w/o RAG performs better than LLM w/ RAG, so the benefit is lower than detriment, the ground ground-truth label is 0.
We use the above method to traverse all sentences in the datasets to obtain prefix and its ground-truth as samples. In evaluation, we input prefix to LLM w/ RAG and use our theoretical findings to judge whether the benefit of the next predicted token is greater than detriment. We use AUC and F1-score as the evaluation metrics for this binary classification task.
Baselines for Benefit-Detriment Comparison Experiment. Task in benefit-detriment comparison experiment can also be viewed as predicting the correctness of the generated tokens. Therefore, baselines for this are the methods that detect the LLMs' hallucination. We use these baselines to
8

---

Published as a conference paper at ICLR 2025
Table I: Accuracy on short-form open-domain Q&A given the retrieved texts containing different ratios (0% to 100%) of hard negative passages (irrelevant but are ranked in top-10 by retrieval model). Our Tok- Tok-RAG does not need any training or additional modules while baselines need.
Methods
Train
Utiliy Evaluator
Ratio
Hard Negutive Passages
TriviaQA
Ratio of Hard Negative Passages
WebQQ
Ratio
Hard Negative Passages
Squal
100%
80%
60%
10%
20%
0%
100%
80%
60%
10%
20%
0%
100%
80%
60%
10%
20%
0%
Standard RAG
71.3
76,2
78,2
40.6
43.4
43.2
58.8
67.2
NLI+AG
V aced x
50.8
61.2
68.2
73.0
76.4
79.1 30.7
40.3
44.5
47,5
50.9
52.8
9.9
21.1
33.7
43.4
51.7 60.5
CRAG
need
76.7
81.5
82.2
374
41.
61.1
RetRobust
need no ar
49.2
67.3
72.9
77.5
79.4
82.3
30,0
38.9
42.5
48.2
49.8
51,3
10.5
30.8
43.3
52.5
58.4
66.0
Self-RAG
no
80.8
82.
18.3
34.8
52.5
58.3
INFO-RAG
need X no ar
49.7
68.4
73.2
77.9
80.0
82.5
29.7
38.0
43.9
48.1
49.4
54.8
10.7
30.1
43.5
53.7
59.2
67.5
Tok-RAG (Ours)
77.6
47.3
Table 2: Accuracy on various long-form NLP tasks.
Train
Unlity Evaluator
Dialoguc Wew F1 S1-re
Code Generation
Slot Fill
Langange Model
Long-forn QA
Method
LLM
Python CadeBLEEU
CodeBLEEL
Java
Accumay
T-REx
WikiTest- 103
ROUGE
ROUGE
EL15
Standard RAG
no
no
7.85
21.44
22.99
$5.60
60.77
15:18
NLI-RAG
no of
need x
8.04
22.79
27.45
63.28
62.05
16.14
CRAG
need x
8.96
24,90
30.03
64,17
62,28
17:03
RetRobust
need
no
9.03
23.18
29:74
63.19
62,40
16:90
Self- AGG
need
no
8.55
22.15
29.60
63.24
61.22
16.47
INFO-RAG
nocd
9.00
26.75
32,06
65.91
62.91
17.18
(Ours)
Bo
32.59
compare the benefit and detriment by comparing the hallucination degree at token level between RAG and pure LLM (details in Appendix [.3). Baselines include: (1) Logprobs-based (Kuhn et al.
2023). we use the value order between
log-probabillty of the tokens output by pure LLM and RAG to determine the value order between benefit and detriment. (2) Uneertainty-basee
we use L.engthnnormaizeed Entropy (Malinin & Gales] (020) to measure the uncertainty of the tokens and compare it between pure LLM and RAG. (3) consistency-basee we run LLMs multiple times and calculate consistency scores among multiple answers using Lexical and Semantic Similarity (Lin
ct al. 2022; Chen et al. 2024). We compare these scores between pure LLM and RAG to indicate the comparison between benefit and detriment.
Implementation details. As for retrieval in RAG, we follow (Xu et al. 2023 to use ColBERT (San- Thanam al. 2021 laa the retriever, and use Wikipedia consisting of 1,1553324 passages (Karpukhin et al. 2020) as retrieval database. All baselines and Tok-RAG share the same retrieval setup and input. We use OPT-6.7B, LLaMA-2-7B, and Mistral-7B-v0.1 as LLMs in the derefit- com- parison experiment and use decodddd strategy for generation. Details of retrieval, Tok-RAG, baselines, datasets for each task and metries are Appendix[]
4.2 EXPERIMENTAL RESULTS
Experiment on Practical RAG. This experiment under the practical RAG setting for both short- form and long-form NLP tasks and the LLM is LLaMA-2-7B. We adjust the radio of irrelevant passages in the retrieved passage list from 0% to 100%, which can simulate the degree of noise in the retrieved texts. Table shows that for short-form Q&A in RAG given the retrieved texts with various qualities. our Tok-RAG shows better robustness and performance than baselines. Table 2 shows that our RAA can also perform better than RAG baselines on various long-form NLP tasks. Tok-RAG does not need any additional modules or training and outperforms the strong baselines that need additional modules to evaluate the utility of retrieved texts or training L.MMs. This means Tok-RAG achieves better trade-off between benefit and detriment in RAG, avoiding detriment while securing benefit. It is because our theoretical analysis helps us propose a more fundamental method in comparing benefit and detriment at token level, while baselines cannot achieve.
Experiment on Benefft-Detriment Comparison. Table shows that our RAA achieves better performance in comparing the values of benefit and detriment token level in RAG than baselines across different tasks and LLMs. Basclines compare the benefit and detriment by detecting the degree of hallucination, while our Tok-RAG can directly compare the benefit and detriment based our

---

Published as a conference paper at ICLR 2025
Table 3: Performance on comparing the valurs of benefit and detriment at token level. Significant test with p-value I 0.05 compared with all baselines are denoted as '+'. Setup in this experiment is Comparison Experiment in Section
introduced in Setup for
M.
Methods
Generation
Wikitexi
ASQA
Bio
Times
AUC
F1
ALC
F1
AUC
F1
AUC
F1
Logprobs
65.15
64.33
68.96
67.55
63,24
64.39
55.31
51.41
Uncertiinty
OPT-6.7B
Consistenay-Lssee Lexicaal
10
64.01
62.17
69.42
67.04
65.41
65.28
55.06
51.13
Somissency-Semmant
10
65.93
64.22
70.11
69 19:50
65.76
64.37
56,24
52.88
Tok-RAG (Ours)
72.28
72.05
57.92
52.90*
Legprobs
73.52
72.90
68.05
66.86
65.22
64.39
57.04
57.23
Unecrtainty
73.72
72.71
67.47
65 63
65,59
65.83
57.19
57.10
Mistral-78
Comiistency-Looee
64,79
Consistency- Semantic
10
73,98
72.26
70.05
69 54
65.68
65.12
57.43
56.
Tok Tok-RAG (Ours)
75.85
74.11*
71.51*
71.47
66.37'
66.04
58.52"
57.56*
Loggrobs
73:47
05.04
62.11
Uncertainty
73.98
73.01
68.72
67.66
63.67
63.50
68.03
69.11 15
LLaMA-2-78
Comsistency-Loxxee
10
73.51
71.62
70.09
68. 45.45
62.49
61,98
68.17
70.09
Convistency Seniaaii
74.96
74.23
Tok-RAG (Ours)
81.89
80.42*
76.96*
76.80
64.08*
64.19*
70.50
72.45*
theoretical analysis, which is more fundamental so it performs better. The detailed setup for this experiment can found in Section 4.1
Analysis on Computational Costs. Table shows that our Tok-RAG achieves significant performance (accuracy) improve- ment with little increase of GPU memory and running time. In
Method
GPU (GB) 1 Time (s) J Per. (Acc)
Siandard
17.50
1.95
51.96
addition, the baselines all require training L.LM or introducing additional modules, while our method is based on solid the-
NLI+RAG
20.25
2.03
49.76
CRAG
20.60
2.15
51.53
RetRobust
17.50
2.00
$2.98
oretical findings, dose not require training LLM nor additional modules. The parallel generation at the token-level can be done in a batch with 2 batch size. In actual practice, this does not bring significant computational time GPU memory overhead. Experiments are performed on three Q&A datasets (TrviaQA, WebQ, Squad) with V100 GPU, the LLM is LLaMA-2-7B.
Self-RAG
17.50
2.98
50.26
INFO-RAM RAG
17.50
2.10
52.35
Tok-RAG
23.98
2.24
56.12
Table 4: Comparison about GPU Memory, Time, and Performance.
RELATED WORK
Robust RAG. To make LLMS robust RAG to avoid the detriment caused from noisy in retrieved
texts, some methods use additional modules to filter out irrelevant documents (Yoran et al.) 2024
Yan al.) 2024, Deng et al. 2023). Some methods train L.LMS to make them robust noisy in
retrieved texts (Xu et al. 2024b; Yoran al. (0244). Some methods let LLMs dynamically determine
whether the query needs RAG (Asai et al. 2023 Xu et al.] 2023 Ren et al., 2023 Feng et al. 2023
Mallen et al. 2022 Jiang& et al. (023). All the previous works solve the contradiction between benefit and detriment in RAG from the perspective of application but lacking essential and theoretical
analysis, which limits the understanding and cannot find the fundamental method solve it. Our paper explains the benefit and detriment in RAG by theoretical analysis and proposes a novel method to preserve benefit while avoiding detriment without any additional modules or training.
Theoretical analysis of ICL. Our paper is inspired by theoretical analysis of ICL. Some works explain ICL as one-step gradient descent (Von Oswald et al. 2023) Akyürek et al.] [022] Dai et al. 2022). Besides, there are other explanations of ICL such as Bayes inference (Xie et al. 2021]. Bayes model averaging (Zhang et al.) (023). leaning topic structure (Li et al., 2023) and kernel regression (Han al., 2023). They focus on explaining why ICL occurs. Our contribution lies in analyzing the benefit and detriment in RAG and proposing practical method based on our theory.
CONCLUSIONS AND DISCUSSION
This paper takes the first step in theoretically giving the essential explanation of benefit and detriment in RAG to make them explainable and comparable at token-level. We provide a theory about describing the trade-off between the value of external knowledge and its potential risk of misleading LLMs in next token prediction of RAG. We prove that the actual effect of RAG can be predicted at token level by representation similarity. Based on our theoretical results, method that enables pure LLM and RAG to collaborate at token level, gaining benefit while avoiding detriment. Experiments show the effectiveness of our method and support our theoretical results.
propose practical novel
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the Strategic Priority Research Program of the CAS under Grants No.XD006803 the National Natural Science Foundation of China (NSFC) under Grants No. 62276248, and the Youth Innovation Promotion Association CAS under Grants No. 2023111.
REFERENCES
Mark A Aizerman, Emmanuil M Braverman, and Lev Rozonoer. Theoretical foundation of potential functions method in pattern recognition. Avlomatika Telemekhunika, 25(6):117-936, 1964.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm in-context learning? investigations with linear models. arXiv preprint arXiv:22 15666, 2022.
Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-rag: Learning to retrieve, generate, and critique through self-reflection. arXiv preprint urXiv:2310.1151/ 2023.
Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model.
Advances in neural information processing systems, 13, 2000.
David Blei, Andrew Y Ng. and Michael Jordan. Latent dirichlet allocation. Journal of machine Learning research, 3JJam):993-102) 2003.
Chao Chen, Kai Liu, Ze Chen, Yi Gu, Yue Wu, Mingyuan Tao, Zhihang Fu, and Jieping Ye. Inside: Lims' internal states retain the power of ballucination detection. arXiv preprint arXiv:2402.03744, 2024.
Yung-Sung Chuang, Yujia Xie, Hongyin Luo, Yoon Kim, James Glass, and Pengcheng He. Dola: Decoding by contrasting layers improves factuality in large language models. arXiv preprint arXiv:2309.03883 2023.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. arXiv preprint arXiv:2212. 10559, 2022.
Jingeheng Deng, Liang Pang. Huawei Shen, and Xueqi Cheng. Regavae: A retrieval-augmented gaussian mixture variational auto-encoder for language modeling. Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computational Linguistics; EMNLP 2023, Singapore, December 6- 6-10, 2023, pp. 2500-2510, Association for Computational Linguistics, 2023. doi: (0.18653W1/222FFIIDING E.0.S.-EMN...... 2023. www.diingng-mmmlpp....
Emily Dinan, Stephen Roller, Kurt Shuster, Angela Fan, Michael Auli, and Jason Weston. Wizard of wikipedia: Knoweedge-powered conversational agents. arXiv preprini 1rXiv:181 2018.
Hanxing Ding, Liang Pang. Zihao Wei, Huawei Shen, and Xueqi Cheng. Retrieve only when it needs: Adaptive retrieval augmentation for hallucination mitigation in large language models. arXiv preprint arXiv:2402. 10612, 2024.
Hady Elsahar, Pavlos Vougiouklis, Arslen Remaci, Christophe Gravier, Jonathon Hare, Frederique Laforest, and Elena Simperl. T-rex: A large scale alignment of natural language with knowledge base triples. In Proceedings of LREC 2018, 2018.
Angela Fan, Yacine Jernite, Ethan Perez, David Grangier, Jason Weston, and Michael Auli, Eli5: Long form question answering. arXiv preprint arXiv:1907.09190. 2019.
Shangbin Feng, Weijia Shi, Yuyang Bai, Vidhisha Balachandran, Tianxing He, and Yulia Tsvetkov. Knowledge card: Filling lmss knowledge gaps with plug-in specialized language models. In The Twelfth International Conference on Learning Representations, 2023.
Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. In-context learning of large language models explained as kernel regression. arXiv preprint arXiv:2305. 12766, 2023.

---

Published as a conference paper at ICLR 2025
Kazuki Irie, Röbert Csordás, and Jürgen Schmidhuber. The dual form of neural networks revisited: Connecting test time predictions to training patterns via spotlights of attention. In International Conference Machine Learning, pp. 9639-9659) PMLR, 2022.
Srinivasan Iyer, Ioannis Konstas, Alvin Cheung, and Luke Zettlemoyer. Mapping language to code in programmatic context. arXiv preprint arXiv:1008 100.008808.09888,, 2018.
Zhengbao Jiang, Frank Xu, Luyu Gao, Zhiqing Sun, Qian Liu, Jane Dwivedi-Yuu Yiming Yang. Jamie Callan, and Graham Neubig. Active retrieval augmented generation. arXiv preprini arXiv:2305.06963 2023.
Zhuoran Jin, Pengfei Cao, Yubo Chen, Kang Liu, Xiaojian Jiang, Jiexin Xu, Qiuxia Li, and Jun Zhao. Tug-of-war between knowledge: Exploring and resolving knowledge conflicts in retrieval- augmented language models. arXiv preprint arXiv:2402.1 14409, 2024.
Mandar Joshi, Funsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv: (7Xx::/705.035// 2017.
Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2004.04900 2020.
Lorenz Kuhn, Yarin Gal, and Sebastian Farquhar. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. arXiv preprint vrXiv:2302.09664 2023.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising requence-to-sequee pre-training for natural language generation, translation, and comprehension. arXiv preprint 1r11v:1910.13461, 2019.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pp. 19689-19729. PMLR, 2023.
Zi Lin, Jeremiah Zhe Liu, and Jingbo Shang. Towards collaborative neural-symbolic graph semantic parsing via uncertainty. Findings of the Association for Computational Linguistics: ACL 2022, 2022.
Andrey Malinin and Mark Gales. Uncertainty estimation in autoregressive structured prediction. arXiv preprint arXiv:2002.07650, 2020.
Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Hajishirzi. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. arXiv preprint arXiv:2212.1055 1055 2022.
Stephen Merity. The wikitext long term dependency language modeling dataset. Salesforce Metamind, 9, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick Lewis, Majid Yazdani, Nicola De Cao. James Thome, Yacine Jernite, Vladimir Karpukhin, Jean Maillard, et al. Kilt: a benchmark for knowledge intensive language tasks. arXiv preprint arXiv:2009.0223 2020.
Lawrence R Rabiner. A tutorial on hidden markov models and selected applications in speech recognition. Proceedings of the IEEE, 77(2)257-286, 1989.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv: 1606.05500 2016.
12

---

Published as a conference paper at ICLR 2025
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Brown, Yoav Shoham. In-context retrieval- wuggeente language models. arXiv preprint arXiv:2302.00083, 2023.
Brown, and
Ruiyang Ren, Yuhao Wang. Yingqi Qu, Wayne Xin Zhao, Jing Liu, Hao Tian, Hua Wu, Ji-Rong Wen, and Haifeng Wang. Investigating the factual knowledge boundary of large language models with retrieval augmentation. arXiv preprint arXiv:2307. 1019, 2023.
Shuo Ren, Daya Guo, Shuai Lu, Long Zhou, Shujie Liu, Duyu Tang, Neel Sundaresan, Ming Zhou, Ambrosio Blanco, and Shuai Ma. Codebleu: method for automatic evaluation code synthesis. arXiv preprint wxxv:2009.10299/ 2020.
Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Colbertv2: Effective and efficient retrieval via lightweight late interaction. arXiv preprint arXiv:2112.01488, 2021.
Timo Schick, Jane Dwivedi- Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. arXiv preprint 2rXiv:2302.04761, 2023.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay. and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 1741456-7747,, 2022.
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettle- moyer, and Wen-tau Yih. Replug: Retrieval-augmented black-box language models. arXiv preprint arXiv:2301.12652, 2023.
Ian Tenney, Dipanjan Das, and Ellie Pavlick. Bert rediscovers the classical nlp pipeline. arXiv preprint arXiv:/905.00 05950, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp- 35151-35174. PMLR, 2023.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are latent variable models: Explaining and finding good demonstrations for in-context learning. Advances Neural Information Processing Systems, 36, 2024.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprint 17Xiv:1704.05426, 2017
Jian Xie, Kai Zhang, Jiangjie Chen, Renze Lou, and Yu Su. Adaptive chameleon or stubborn sloth: Unraveling the behavior of large language models in knowledge conflicts. arXiv preprint arXiv:2305.13300, 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 2021.
Shicheng Xu, Liang Pang, Huawei Shen, Xueqi Cheng, and Tat-seng Chua. Gearch-in-hhe-chal Towards the accurate, credible and traceable content generation for complex knowledge-intensive tasks. arXiv preprint arXiv:2304. 14732, 2023.
Shicheng Xu, Liang Pang, Jun Xu, Huawei Shen, and Xueqi Cheng. List-aware reranking- trrnccttnt joint model for search and retrieval- auggeened generation. arXiv preprint urXiv:2402.08 2024a.
13

---

Published as a conference paper at ICLR 2025
Shicheng Xu, Liang Pang, Mo Yu, Fandong Meng, Huawei Shen, Xueqi Cheng, and Jie Zhou. Unsupervised information refinement training of large language models for eetrieval-augmenn generation. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), pp- 133-145, Bangkok, Thailand, August 2024b. As- sociation for Computational Linguistics. doi: 18555////224.4c/oooggg
URL https
//aclanthology. www.oyy.org//22....omm 1ong..n...9
Shicheng Xu, Liang Pang, Yunchang Zhu, Huawei Shen, and Xueqi Cheng. Cross-modal safety models. In The Thirteenth International Confer-
mechanism transfer in large
ence on Learning Representations, 2025. URL https::/openreeteeeeee net ftrtmfo
45rv2kJouuX
Shi-Qi Yan, Jia-Chen Gu, Yun Zhu, and Zhen-Hua Ling. Corrective retrieval augmented generation. arXiv preprint arXiv:241.15884, 2024.
Ori Yoran, Tomer Wolfson, Oni Ram, and Jonathan Berant. Making retrieval-augmented language models robust to irrelevant context. In International Conference Learning Representations, 2024.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint 1rX1v: 2023.
A PROOF FOR EQUATION
Proof. The transformation is motivated by (Xie et al. 2021) and we apply it to the analysis of RAG:
(115-1) p( R. 11 ( (12-1)
(20)
pl R, 11s- 015-112)p(2) de
p(a; R. x14-1, 2)
(21)
pl R, x1:i-1)
x
p(aae R, R, F1: 11
pl (R, a constant so we drop it
(22)
pl 1RR:I11111
1
plzz Rl
p(R, 21 1 (20) p(z) dz,
is constant add it
2
p(R, x11--122))
(23)
p(a;l R, 2xxxxx) dz, x(2) = log 11, /
(24)
pl R, 11:1-111
B PROOF FOR EQUATION
Proof. For p( R, 12)) in v(2) log P(R,2.n-1ss) pl R. ) we can make further derivation as:
P(R,I15--13 3:0= 1(0 =)(2)
(25)
According to the definition of latent variable model in the analysis of in-context learning from Xie
et al. 2021) that views the latent variable inference as Hidden Markov Model (HMM) and the latent
concept determines the transition probability matrix in HMM hidden states h, we can get the
following derivations:
p( 18, [R, = Σ 14, 2)) 2)p(()))))
(26)
p( R, 11:11))
v(2) log
(27)
p(R, 2111-112-)
Σn pl 11:1 1(111-111.2)pp(hh. + log p(R[2)
log Σn 1pp(ppppp1/11 1(
(28)
p(R]2*)
14

---

Published as a conference paper at ICLR 2025
According to our Assumption | that p =*) > c1 > 0, 2"") in the denominator of the first term in u(z) is always greater than C1:
p(a1=-1 -,,,") =²) > c1
And the numerator p( pa (h, in the numerator of the first term in u(2) is always lower than 1:
p(#1--1|h, 2) < 1. We replace p(z1 2) in the numerator with 1 and p(z1--1|/ 2*) in the denominator with e1 and we get:
Σn 1 2)))
log p(
o(z) S log
(29)
Σn C1 (R R²,") 2*)
p(((==))
p(h] 2)
log p(R]2)
log
(30)
Σn C1 p( 2*)
p((((2))
p(R]z)
logec log
(31)
p(
loger + log 11.h p(ri)(usé-1,
(32)
p(r:/ria-112"")
According to the chain rule, we can transform p( R, 21 as:
pl R, - p R, 14
Eqn 26 32 shows:
(22) = log p(R, 1: 112)
Eqn 5
p(R, 11: |2*)
p( R, 2)p(
log
p( 1 R2 2² )p( 1**)
= log p( R, 2)
+ loo ((R]2)
p 22,2*) 2*)
p(R]22)
p(R(2)
- logcc + log p( ((RR (R(]2*) 2*
So we can get:
p R2
log
S lloge
p 0(01--11 R,
p RR,2)
p RR,
C1
(R,
((51:-1|RR22)
C1
So ( is bound by O(1) and then we can get:
p R, 11:11=11/3)) = p( R, 2 (R|22)
= ()))((R(2)
Since is the passage in R, so we can get according to the chain rule:
p(R(2) = II p(../1a-1.2))
So we can get:
= 2) ([(2) ~ 00)p())mass11)))
(33)
II o0
= II Σ
p(r.ln 2p 11
(34)
A ED
15

---

Published as a conference paper at ICLR 2025
is passage in the retrieved texts list R, h is the hidden state for the delimiter between 7:-1 and r, in R. Then, decompose p(h 11; with chain rule:
1 = Σ _
According to the standard assumptions of HMM Rabiner 1989), h and r1s-1 are conditionally independent, so we can get:
p(h 10--111--112/) = Σ p(hh 2((1)))))
Assumption 2 assumes that for any delimiter hidden state ha and any other hidden state h, the transition probability from h ha has upper and lower bound 0 02 p(h4th,)) =) S cg. Then we can get:
Epthi (00--111)))0(Amaa--11))))
= = Σ
p(hlr1s-12)= =
Therefore, p(h (1 0()) Then Equation 34 is approximately equal to IT O(1)p(rs(z)
which means that p( R, 11)) ~ IT (1)pp(r, so we can get that:
PRR, 21
O(1)p(r.1=)
v(2) log p(R, -1:11112))
2 log
(35)
O(1)p(ri(=*)
log p(ri(z) = 0(1) + n F.*EE,. EnnPn
log p(r:(2)
0(1) + n
p(r,|2*)
p(ra(2*)
(36)
x pn(r)log p(rlz)
pn (n(r)log pr(r)
p(rl2*) pn(r)log pn(r)
(37)
p(rl2*) -((KL(pp(r) (p((l)) (Ilp(()))l(((ls)))))
p(r(2)
(38)
benefit, as f2
detriment denote as
(R(()) is the distribution of the retrieved texts, p(-) is the distribution of the LLMs' pre-trianed knowledge.
C EFFECT OF v(z) IN DISTRIBUTION FUSION
Recapping Equation 38 we find r(=) actually regulates the proportion between LLMS' pre-trained knowledge and retrieved knowledge in distribution fusion of RAG prediction:
The more benefit outweigh detriment, e(z) + oo and exp(u(=)) + for all = of z², this indicates that concepts sampled from LLMs' space contribute little to p(z; R, 11:--1). When 2", exp(r(2*)) )) = 1, which means that latent variable model concentrates more on 2" sampled from retrieved texts. As v(z) decreases, the proportion of retrieved knowledge in becomes larger and larger in fusion.
The more detriment outweigh benefit, x(2) too and exp(u(2)) - +oo for all * x* and when z = z², exp(r(2*)) = This indicates that concepts & sampled from LLMs' space contribute more and more than =" sampled from retrieved texts u(=) increases.
16

---

Published as a conference paper at ICLR 2025
D PROOF FOR THEOREM 1
Proof. Recapping the Equation that describes the distribution fusion in RAG via latent variable
model:
/ 3 1--
p(a, R, 11s-1)
p( R, F1.- 11.-11 2))== 215-1) p(z; 319-1, )p(2 (1222|RR RR 21
denote A
denote +
(39)
Since latent concept x" determines the hidden states h, A can be transformed as: p(z, R, F1.-1.,2 )p(p" R. 215-1)) = Σ plzjlmnn-1.h, (p())) R, T F15-1.22 p.=)== RR, R1:-11).
h
(40)
Let p(=" RR, (1111) + + 3 Σ p(a (arss h1 p)((p) h xmn-112")
(41)
pR(D.(21A-1) Σ (1:0),111111111 h ()
(42)
pn(h((1ss1) p(1111111111 (b))
(43)
2( x xp(x (
(44)
let probabilities p(z,|21s-11,h, (p11=)))p(11111.2220 in Equation 40 is represented as matrix W
g|*|*|||| for all possible T; E X and h E H. p(h| =*) in Equation 44 is matrix B, pR(h) in
Equation 43 is E IR X41, We use 1-norm to calculate the difference between p(z;//R, (1a--1 and pR( [1. ) which can be formalized as:
[|p(x; [R, -1 ) - (R(x,/=1=__1) 1ll = 11 + 3W B - Wulli. Then, according to the triangle inequality of 1-norm, the difference between p(r; (1, and (2 (11:-11 is bouned by:
(45)
| |||
|@WB B - Wull < 0 + 3W - Wull < |||]| + | B - Wull
(46)
We consider to further analyze |/W B Wall inspired by Xie et al. (002))
|@W B - Wul
WF (BB w)lF
(47)
H
WBy u
(48)
Wo n)yl
(49)
X
(3B - w)pl
(50)
P
Σ j=1
KBB KBB-
(51)
||3B - u|11
(52)
Then:
|| - u/I1 = (p(111)
TV is Total Variation Distance.
(53)
< 23TVV(p(() ((,(
(54)
$ /2KL(pp() 1ppp R),+)) Pinsker's Inequality.
(55)
I /2KL(pR( (()((P(())1p(122))))
(56)
(57)
17

---

Published as a conference paper at ICLR 2025
in which is the passage in R, (L(pp(())lp(r)a") is actually the detrimentin Equation Equation 46 we can get:
Recapping
I + 3WB Wull < |+|| /2Xll/pp((()lll((]==))
(58)
Since 0 < Wulli, then the lower bound for | + 3W Wulli is included in:
- 2v (KL [1())=))) and B- 2 ||||| - @|
+ 11 v ()) 1|
(59)
we take the minimum value as the lower bound. Define D |p(zz R, T15-1) p ||
the difference between p(z; R, F15-1) and pn(2s(21=-1), according to Equation 45 and 46 the
lower and upper bound for D is:
|
/2 LLLL (L(pp())l)(())" $
D
| || +
// (L(pp(()ll(raa")
(60)
detriment
detriment
For ease of description, denote benefit (L(pR(r)ll(rla) D and denote detriment KL(pn(r) lp((1=²) as Y. Recapping Equation and
p(a, (14-1)
21, x14-1) da + tp(,,|R, x0 )p("") RR, x1:1-1)
2-
denote
denote +
p(a; R, (1 (1- dz
p(a; R. 111- 2) p(R, 31-- 115112)p(22 dz
pl x1:5-11)
p(ai R,G,, 2) )p 11. 1z)p(())))p))))))))))) dz, p(R,zra-1)) is constant drop it
pl P(R,
1
p( R.ilR 01--12 2)
p( (R p(z) dz,
is constant add it
p(R, 21i-1111122 v(2) log pp R2 x1 (2) pl R, 11:1-1111
R, R, p (xxpxxxxxx((())))))e
u(2) log p( R, 11:111111 1/2*)
(L (L(pp(r))l((rl))
benefit, as f
detriment denote as
in Equation 60 can be transformed as:
[ (21(- 21s,
dz
(61)
pl (R (R,
p(ail 200
(62)
p(R,I11-1)
2-13*
0 2-112)
plaal R 1xx
x15-122=)1 and p(R, x1a-1) are constants)
(63)
R3
R, 2 ()))(=))))) dz (Equation 9)
(64)
2 (2-)
Now the Theorem has been proven.
18

---

Published as a conference paper at ICLR 2025
E PROOF FOR THEOREM 2
In this section, we try to prove that the gap between values of benefit and detriment is approximately positively correlated with the similarity (+) between p(2; R, f1 and PR ( To achieve this, we can start from Equation 160 prove that the gap between values of benefit and detriment is negatively correlated with the difference (D) between p(r; (1:--) and pR ( which is actually the reciprocal of similarity (g) Specifically, we want to prove that the gap between values of benefit and detriment (KL(px(r)lp((()
(L(px(r))lp(r ())))(()=))))))))))))))))))))() is negatively correlated with both lower and upper bound in Equation 60 For ease of description, we denote benefit (l.(p((r)llp(ree as Q and denote detriment Kl.(pn (r)((pn()(((())))) (p)(() as Y.
Proof. Recapping Equation 5
and
/
O, (15-1))
p( RR Ti 2.t
F1a-1) dz + p(x; Fi5 p(2" RR x1
12-13
denote A
denote 4
p(ay 21
dz
p(a. R. R1 pl R, x1 115-112)22pp
de
p( 11:1-1)
p(a R. 114-1, z) R, 11:1-11 12)p(() dz, pl p(R,311111) constant so we drop it
2
p(a R. 014-1, 2) plR, 2R p(2) dz, p(R, x1---1/2*)
a constant so we add it
p( 1R )
pRR, 1: 12) pla;l xi4- 1 2xxp(( dz, r(z)= = log pRR, x x 12*)
u(2) = log p(R, 11;-11111
KL(PR())lp((12)
KL(pa(r)
pl R2 $1:1 |z*)
benotit,denote f
detriment, denote as T
+ in Equation 60 can be transformed as:
[ 2-1 2p 21 dz
(65)
p(R, 111))
(66)
p( x11-1) 2-(0")
plaa R, 21,-1, sexp(())))))p( (1a-1e") and p(R, 214-1) are constants)
12-13
(67)
23 12-13
pl 1, ())))(((--())))))) (Fquation 9.
(68)
Therefore, the lower bound of Equation
60 is:
|| VT all /
p(z, R, 2xx1,, (0oxp(((00)))(()dde))))))))))- dz/l1 V2T
(69)
exp(-((1-1)) /2T
(70)
and the upper bound of Equation i0]
+ l1 V2Y exp(-((8 ((--(((())))) T)) V2T
(71)
19

---

Published as a conference paper at ICLR 2025
Due to both Q and Y being variables, analyzing the result of subtraction between Ω and Y under their simultancous changes is complex. Therefore, we use the "Separation variables" to simplify analysis. Specifically, we first assume that one is constant, and then analyze the changes caused by the variation of another:
Assume Q2 is constant, as the value of increases, Y decreases and the upper bound exp(-(22 T)) + v2T also deceases. the lower bound exp(-(9 1)) v2T, since the first term is exponential function and the second term is square root function, decrease of Y leads to the decrease the entire lower bound. Therefore, both lower and upper bounds in Equation 60 decrease Ω YY increases.
Assume T is constant, as the value of Ω increases, Q increases and the upper bound exp(-(88 (Q2 r)) + /2Y deceases. In the lower bound exp(-(22 - T)) v2Y, since the first term is an exponential function and the second term is square root function, an increase of Z2 leads to the decrease in the entire lower bound. Therefore, both lower and upper bounds in Equation 60 decrease Ω increases.
On behalf of the analysis above, negatively correlated with the gap between values of benefit and detriment. Therefore, the difference D between p(a,,| R, (1: and (3;] (10//:/111111) 111--1) is approximately negatively correlated with the gap between values of benefit and detriment.
derve that both lower and upper bounds in Equation 60 are
Then we try derive the maximum
for "approximation" in this approximate correlation and 71, the lower bound of D is:
make it more standard. According to Eqn and
| - /2T x exp( ( - T)) - V2T,
And the upper bound of 'D is:
| + /2Y x exp(- T)) + V2r,
Therefore, when the upper and lower bounds are determined, the maximum range of D is 2/2T and the average value of DD approaches ||+|||. Since Equ 68 shows ||6||] is an exponential function of T - Ω so the maximum error rate Cmaz is function of the detriment Y brought by RAG, which can be described as:
exp
According to Eqn 9 the detriment Y is the KI. divergence between two distribution, so the value range of T is (0, foo). So we can get:
VT
lim T ++ exp(P)
Therefore, the greater the detriment brought by RAG the smaller the maximum error rate of this correlation and the more robust of our theory. This shows that our method is effective avoid the detriment brought by RAG at token-level.
Since 1 can be treated as the similarity between p(x;|/t x14-1) and pr x=][=)=1)) and is approx- imately positively correlated with the gap between values of benefit and detriment.:
1
(L(p(rr)ll(rla)) KL(pp(r))l((r)")) x D'
(72)
benefit
detriment
So we have proved that the gap between values of benefit and detriment is approximately positively correlated with and the maximum error rate of this approximation decreases the detriment T increases and approaches 0.
20

---

Published as a conference paper at ICLR 2025
F
PROOF FOR THEOREM 3
This section aims to prove:
(L(pR(r)l < (L(pa())ll(rrr")) detriment outweighs benefit. if (L(r)))) = (l.(pn(r))lo(rra"))) detriment is equal to benefit. (L(pr()))l((rk)) > <L(pp<r)llp(r(<> benefit outweighs detriment.
if D M (73) if T > M
in which 1 is the similarity between p(z, #1 and p(as(Z1a-1) (LLMs' pre-trained knowl- .M
edge), D is the similarity between p(a; [,, m14-1) and pn(mi(21s01 (distribution of retrieved texts)
Proof. When benefit is equal to detriment:
(L(pR())l)(r11)) (L(pa(r)ll(r12")) = O,
(74)
which means that:
(o( p(r(2) O,
(75)
p(rlz*)
since (r(r) cannot be 0, then:
log p(r[z) 0,
(76)
p(r|22) 22))
p(r(2)
1,
(77)
p(r(2*)
p(r[2) = p(rlz*),
(78)
Recapping Equation 2 that z" is sampled from retrieved texts and z is sampled from LLMs' pre- trained knowledge, Equation 78 indicates that the knowledge of retrieved texts has been involved in LLLs' pre-trained knowledge, so:
((s//ra-1) = 1)nos)
(79)
then:
Ip(2; (R, 01:5-1) p(1::/11:: )|
pn (213-1)
(80)
which means that D M is an important dividing point. When 1
M we can get that benefit
equal to detriment and (L(pR(r))lp(ri that the gap between values of benefit and detriment (KL(pR(r) [[((lz])
- (rpppr)) )) 0. Equation 72 indicates
(L(p [p)
approximately positively correlated with $. Therefore, when D .M we can get that benefit
outweighs detriment ([L(((())ll((()))) - (L(pp(())ll(rr=") > 0). When D
< .MM we can get
that detriment outweighs benefit (KL(pp(r)llp(rt)) KL(PR(T)ll(r(z") Ó). Now the proof of Theorem 3 has been finished.
G PROOF FOR RAG IS ACTUALLY UNSUPERVISED IN-CONTEXT LEARNING
This section aims to prove that RAG actually unsupervised ICL from two perspectives. One is that previous studies find that ICL performs gradient descent as meta-optimizer (Von Oswald et al. 2023 Akyürek et al.] 2022; Dai et al.) (022). We prove that in this perspective, the distribution of texts in context drives the learning even without explicit input-output supervision. Therefore, the distribution of unsupervised retrieved texts RAG, which is actually the distribution of context for query, can also drives the learning. Then we can prove that RAG is actually unsupervised in-context learning. The specific proof is:
Proof. From the perspective that ICL performs gradient descent as meta-optimizers, ICL can be formalized the following:
Gradient descent in optimization of linear layers have dual form of linear attention (Irie et al. 2022
Aizerman et al. 1964), define a liner layer as:
f(x)
Wox.
(81)
21

---

Published as a conference paper at ICLR 2025
in which Wi is the initial weight matrix. Given sequence of historical input vectors x; R and corresponding error signals e E Rdout, [1 obtained by gradient descent, the update of the weight matrix can be represented as:
N
W = Wo + AW = Wo Σ ei 8 xi.
(82)
Recap that the linear attention can be formulated as:
LinearAttn(V. K, = Σ vi(kg 9).
(83)
Then the dual form of updated linear layer with new input XN++1 is:
f((x) = (Wo + AW)xN+1
(84)
-Wo + Σ c x1
(85)
= W + Σ (e, @ xx)xx+1
(86)
N
= WWNN++1 + Σ @ (xF xN+1)
(87)
Wo + LinearAttt(E E, XIN, XN+1)
(88)
In in-context learning, the attention of a head is: fici(q) = At, K29)
(89)
WK|B': B) Vi
= Wv [B': B]softmax
(90)
in which q = Wob, b is the input t-th token in query. Wo. WK. W. are projection matrices, B' is demonstrations of the context and X is the prefix for b of the query. To simplify qualitative analysis, we follow (Dai et al. (022) to estimate the standard attention as relaxed linear attention, achieved by
eliminating the softmax function and the scaling factor:
(g) = [VIB': B] [[W](BB (B': B))²q q
(91)
BVW B) + Wy X'Wx'''
(92)
Wy B) q + LinearAttnWW, B', WK B', q)
(93)
According to Equation
88 the dual form of the Transformer attention is:
fI ~ Wy T q + LinearAttnWWy B', B' q))
(94)
WyB(W BY q + Σ (WWx())) =
(95)
WWvB(WBB) q+ Σ ((Wvb') @ (Wxb()) q.
(96)
Based on above derivation, we have this finding: comparing Equation 86 with Equation 96 we
find that Wv B(Wx k B)" is equal to the initial weight matrix Wo. which zero-shot prediction give query prefix B without demonstrations in the context. Besides, Wy-b/ is equal to which is the meta-gradient used to update the weighted matrix. Wxb, b, is equal to the historical input vectors:
Wyb' = e;
(97)
Wx bb; Tj.
(98)
dc In the standard gradient descent with loss L, e -17 dy; and 17 is the learning rate and y: W.r, is the output of the linear layer using the weight matrix W, at step t (Irie et al. 2022). So we can get:
dc
ac
e, = 1 oy,
- oW,
(99)
ac
(100)
22

---

Published as a conference paper at ICLR 2025
Therefore:
ac 11 H.WWWbb. Wvb,
(101)
So the loss L ca be represented as:
L L=
Wyb, d(W:Wxb))
(102)
Equation 97 and 102 show that the supervision signal, both loss and gradient, are directly related to the semantic representation of the tokens (b;) in the demonstration of context. This suggests that the distribution the text in context a direct learning signal for in-context learning, without the need for explicit input-ouiput pairs in the demonstration. From the perspective that ICI. performs gradient descent as meta-optimizer, the proof has been finished.
The other pespective is from our theoretical results in Equation 5 and 9
that:
p(a, (1a-1) = 21=-1, (1a-1) dz
(103)
{
31. mis + p( R, 21 R1.-111)
2-120
(104)
plai pur.::.aeee.uu.p.p..a.m---s1pp):
(105)
(1o r(z)= - log p( 11
(106)
(2, x1:i-1
v(z)= log p( 115--122)
II ((1)p(rs)z)
R3 log
(107)
p( R, 21 --*)
O(1)p(r,/2*)
0(1) m log p(ri(z)
p(r:12)
0(1) + n * ERPP
log
p(r,|2*)
p(r:
(108)
x pn(())og p(r|z) pn( log pn(r)
(n(r)log pn(r)
(109)
p(r|2*)
p(rl2*)
p(r(2)
-((KI((p(()) (p((l=) - KL.(Pn (I)((p(())l(((r=)))))
(110)
We explain Equation 110 from the perspective of the loss function in gradient descent to explain the
learning from retrieved texts of LLMs in RAG. (R()) is the distribution of retrieved texts and it can serve ground truth distribution in loss functions. p(r(z) and p(r[2*) are distribution estimated by
LLMs. Two KL-divergence between ground truth distribution (R( and estimated distribution p(r(a) and p(r]=") respectively
loss functions. is the retrieved passage that invariant in generation process, so what contributes to the change of loss function is sampling more and more accurate retrieved concept from the retrieved texts. So (L(pR(r) [Ip(r[=*) (r)pprr))((((a>')))))))))))))))) is the actual loss that can be
meta-optimized in RAG, in which (R(r) is the ground truth distribution and p(r(2*) ) is the estimated distribution. As this loss decreases, the value of v(z) when not equal to also decreases, which means that in Equation [04] the ratio of the knowledge from LLMs' pre-trained distribution decreases and meanwhile the ratio of knowledge from the retrieved texts increases. Lower loss means that the output of RAG is closer to the distribution of retrieved texts, which is actually that LLMS learning the distribution from retrieved texts in input context. Since pr(rr is the ground truth in this learning but dose not have any explicit input-output supervision like demonstrations in in in-context learning, RAG the unsupervised in-context learning and distribution of retrieved texts (pR(T) is the unsupervised learning signal.
Based on the above two perspectives, we successfully prove that: The distribution of retrieved passage in RAG (i.e., pR(r)) can serve the unsupervised learning signal for LLMs learning from context, even without explicit input-output supervision. RAG is actually unsupervised in-context Learning that fuses the distribution from retrieved texts with L.LMs" pre-trained distribution.
23

---

Published as a conference paper at ICLR 2025
H EXPLORE THE MECHANISM OF RAG
We also perform experiments in Section 3.1 using more LLMs (OPT-6.7B, Mistral-7B) on more datasets (TriviaQA, WebQ and Squad). The experimental results are shown in Table 5 and 6 and the conclusions of them are consistent with Section 3.1
Datasets
Value
0
8
12
16
20
24
28
32
TriviaQA
Attention Score
0.12
0.79
0.65
0.51
0.59
0.72
0.86
0.69
0.43
TriviaQA
Word Distribution Change
0.10
0.11
0.10
0.10
0.13
0.22
0.38
0.56
0.82
WebQ
Attention Score
0.10
0.69
0.65
0.54
0.55
0.70
0.82
0.64
0.40
WebQ
Word Distribution Change
0.15 0.13
0.13
0.14
0.16
0.27
0.45
0.55
0.79
Squad
Attention Score
0.16
0.82
0.64
0.55
0.63
0.69
0.90
0.67
0.48
Squad
Word Distribution Change
0.11 0.11 0.13 0.12
0.12
0.24
0.36
0.69
0.89
Table 5: Attention score for and difference of word distribution change vary with layers based on OPT-6.7B on TriviaQA, WebQ and Squad.
Datasets
Value
0
4
8
12
16
20
24
28
32
TriviaQA
Attention Score
0.05
0.64
0.50
0.45
0.60
0.69
0.81
0.65
0.37
TriviaQA
Word Distribution Change
0.08
0.10
0.11
0.15
0.19
0.30
0.42
0.65
0.82
WebQ
Attention Score
0.07
0.69
0.55
0.46
0.66
0.70
0.82
0.60
0.42
WebQ
Word Distribution Change
0.12
0.14
0.14
0.17
0.29
0.38
0.50
0.67
0.85
Squad
Attention Score
0.10
0.75
0.60
0.48
0.65
0.72
0,86
0.62
0.40
Squad
Word Distribution Change
0.13
0.13 0.15
0.19
0.35
0.40
0.57
0.70
0.84
Table 6: Attention score for and difference of word distribution change vary with layers based on Mistral-7B on TriviaQA, WebQ and Squad.
EXPERIMENTAL DETAILS
1.1 RETRIEVAL
As for the retrieval model and retrieval database, for Slot Filling and Language Modeling, we use ColBERTv2, a late-interaction model with excellent generalization ability as the retriever, and use Wikipedia consisting of 21,015,324 passages as retrieval database. For Code Generation, we SCODE- as code retriever and use deduplicated source codes in CodeSearchNET as retrieval database. For all the above tasks, give Top-5 retrieved passages each example. For ELI5, Dialog, the list contextual passages provided in the datasets as the retrieved list (distractor setting). All baselines and our method share the same retrieved documents.
1.2 TASKS, DATASETS AND METRICS
Open-domain Question Answering Open-domain QA is a typical knowledge-intensiv task that can directly evaluate the knowledge LLMs. We use TriviaQA Joshi et al. (2017). Rajpurka et al.) (2016) and WebQuestions (WebQ) as the datasets. We use cover Exact Match (EM) to determine whether the ground truth exactly appears in the output and the accuracy is used as the evaluation metric, following Schick et al. (2023); Xu et al. (2025)
Slot Filling Slot filling requires LLMs to output the object entities for the input subject entity and relation. We use a knowledge-intensivg dataset T-REx Elsahar al. (018). We use the same evaluation metric as Open-domain QA.
Long-Form Question Answering Compared with open-domain QA, LFQA is the QA task whose ground truth answer is a relatively long text. We use ELI5 Fan et al. (2019). a knowledge-intensive dataset for LFQA. We use ROUGE-L as the evaluation metric Petroni al. 2020)
24

---

Published as a conference paper at ICLR 2025
Dialogue Dialogue in our experiment focuses the factual knowledge. We use Wizard of Wikipedia Dinan et al. (2018) (WoW), a knowledge -oggeeed dialogue dataset whose conversa- tion is grounded with knowledge. We use Fl as the evaluation metric/PPtron et al. (2020).
Language Modeling We use WikiText- 103 Merity (2016), popular dataset for language modeling. We use ROUGE-L the evaluation metric.
Code Generation Code generation aims to generate the code for the given natural language. We use Java and Python in CodeXGLUE Lyer et al. (2018) for this task. We use CodeBLEU Ren et al. (0020) as the evaluation metric.
1.3 BASELINES
For benefit-detriment comparison experiment that needs methods to determine the value order between benefit and detriment for each token, it is actually a binary classification task (benefit outweigh detriment or not). The mainstream methods in this area are detecting and comparing the degree of hallucination between tokens generated by LLMs (w/o RAG) and RAG. Below we will describe in detail how we apply these baselines this task.
Logprobs. Logprobs can indicate the confidence for LLMs in generating the tokens Kuhn et al. 2023). We use the value order between top- probabilbty of the tokens output by pure LLM and RAG to determine the value order between benefit and detriment for these tokens. If the logprobs of tokens generated by RAG greater than the logprobs of tokens generated by pure L.LM, the benefit outweigh the detriment, otherwise the detriment outweigh the benefit.
Uncertainty. We use Length- Length-normalized normalized Entropy (Malinin Gales 2020 to measure the uncertainty of the tokens generated by pure LLM and RAG respectively. If the uncertainty of tokens generated by RAG lower than the uncertainty of tokens generated by pure LLM, the benefit outweigh the detriment, otherwise the detriment outweigh the benefit.
Consistency-LLxxica (Lin et al. 2022). Consistency- based methods make LLMs perform multiple generations for question and calculate consistency score among multiple answers. If the consistency score of tokens generated by RAG is greater than the consistency score of tokens generated by pure LLM, the benefit outweigh the detriment, otherwise the detriment outweigh the benefit. Lexical-based consistency means calculating consistency score by lexical-sing shiial--imilarity among multiple answers. Since the experiment is at token level, we use the number of tokens that are completely consistent in multiple generations as the consistency score.
(Chen et al., 2024). We follow (Chen et al. 2024) to EigenScore to calculate the semantic similarity among hidden states of tokens in multiple generations and use it as the consistency score.
Ponsistency-Semanne
For open-domain Q&A under practical autoregressive generation setting, baselines for this include the methods that introduce additional modules to filter irrelevant passages (NLI+RAG (Yoran et al. 20244) or as action triggers (CRAG (Yan et al.] 2024)), train more robust LLMs for RAG (RetRobust (Yoran al., 2024) and INFO-RAG (Xu et al. (024b)) and train LLMs to dynamically retrieve and critique retrieved texts (SeIf-RAG (Asai et al. 2023)).
This method use Natural Language Inference model to filter the possible irrelevant documents in retrieved results and provide the remaining documents to LLMs for generation. We follow YYoran et al. 2024) to use a BART-Large model (Lewis et al. 2019) with 407 million parameters trained on the MNLI dataset (Williams et al.) (007)). We consider qucry-document pair as entailed the probability for the entailment label is 2 0.5 and filter the documents with probability for the entailment label < 0.5.
NLI+RAG.
This method uses a retrieval evaluator to assess the correctness of retrieved texts trigger different actions based on the evaluation results. One of the actions is using additional google search API for web search, which is unfair for baselines and our method. So we remove this action and use its knowledge refinement strategy for document filtering (aan et al. 2024).
CRAG.
RetRobust. This method fine-tunes LLMs to properly leverage retrieved passages with a mix of relevant and irrelevant contexts (Yoran et al. 2024
25

---

Published as a conference paper at ICLR 2025
This method uses unsupervised method make LLMs learn the retrieved texts robustly. It enables LLMs to judge the correctness of the retrieved texts, extract the correct content and revise the wrong content (Xu al. 2024b).
INFO-RAG.
Self-RAG. This method trains LLMs to dynamically retrieve and critique retrieved texts. Self-RAG first decodes retrieval token to evaluate the utilityy of retrieval and control retrieval component. retrieval is required, LLM calls an external retrieval module to find top relevant documents, using input query and previous generation. If retrieval If retrieval is needed, LLM first generates critique token evaluating whether retrieved documents are relevant and support generation, and then generates continuation conditioned on the retrieved passages (Asai et al. 2023).
required, LLM continues generation.
1.4 IMPLEMENTATION DETAILS
All models are run on a V100 GPU with Pytorch Paszke et al. 2019) and accelerated by DeepSpeed As for retrieval for RAG, follow (Xu et al.
2023 2024b to use C (Santhanam et al. 202I), an excellent generalizable model as the retriever, and use Wikipedia consisting of 21,015,324 passages (Karpukhin et al. 2020)) retrieval database. All baselines and Tok-RAG share the same retrieval setup and prompt. We use OPT-6.7B, LLaMA-2-78 and Histral-7B-vol as LLMs in benefit-detriment comparison experiment and use greedy-decoding strategy for generation.
1.5 ABLATION STUDY.
Figure 4 shows the effectiveness of our dynamic layer selection strategy in Equation [4 and supports our finding that RAG performs matching in middle layers. Figure 4 Figure&shoss shows the AUC when 1" in Equation 14 is set as a fixed value from 0 to 32. Our dynamic layer selection strategy (dashed line) is always better than any fixed layers (solid line). Besides, AUC is higher in middle layers, which supports that RAG performs matching in middle layers and the knowledge in retrieved texts is extracted in the turning point. After the turning point, LLMs instead perform distribution fusion, the matching cannot reflect the distribution of retrieved texts, so AUC decreases.
10
30
Layor
Figure 4: AUC varies with layer.
J CASE STUDY
Figure 5 shows the case study for collaborative generation between pure LLM and RAG at token level in our Tok-RAG. At the step that pure LLM and RAG generates the different tokens, Tok-RAG use our theoretical results in Theorem to compare the benefit and detriment. If benefit greater than detriment, the token from RAG is selected, otherwise, the token from pure LLM is selected. The selected tokens are marked by green color and bold. Then discarded tokens are marked by gray. The orange arrow represents the direction of token selection and usage. The selected tokens used for the next step generation of both pure LLM and RAG. This case study visually demonstrates that our Tok-RAG effectively enables pure LLM and RAG for collaborative generation preserve benefit and avoid detriment.
26

---

Published as a conference paper at ICLR 2025
Question:
Who is the book of galatians written to?
Pure LLM: It was written by the Apostle Peter to the churches in Galatia, region of present-day Turkey, Henefir wit
Detiment win
RAG: It written by the Apostle Paul the churches in Corinth, region of present-day Turkey,
Output: It was written by the Apostle Paul to the churches in Galatia, a region of present-day Turkey.
Question: Who conducted the opening concert at carnegie hall? In which year?
Pure LLM: It was conducted by Leopold Damrosch and composer Pyotr Ilyich Tchaikovsky in 1891.
Banefit win
Detrient wis
RAG: It conducted by Walter Damrosch and composer Pyotr llycch Tchaikovsky in 1901.
Output: It was conducted by Walter Damrosch and composer Pyotr Ilyich Tchaikovsky in 1891.
Question: Who says that which call rose?
Pure LLM: It is said by the character Juliet in
Christopher Shakespeare's play Romeo and Juliet. Besefit win
Deiiinentiin
RAG: It is said by the character Lysander in William Shakespeare's play "Romeo and Juliet.
Output: It is said by the character Juliet in William Shakespeare's play "Romeo and Juliet."
Question: Who signed the largest on the declaration of independence?
Pure LLM: belongs to John
Hancock, who was the president of the Continental Congress
Desiment wis
Demieet win
RAG: It belongs to Charles 'Hancock, who was the member of the Continental Congress
Output: It belongs to John Hancock, who was the president the Continental Congress
Question: Who is under the mask of darth vader?
Pure I.M:: It is the character
James Skywalker, as revealed in the Star Wars series. Bendfi win]
RAG: It is the character Anakin Skywalker, as revealed in the Star Wars series.
Output: is the character Anakin Skywalker, as revealed in the Star Wars series.
Question: Where does US highway
start and end?
Pure LLM: It starts in Chicago
Kent, Maine, at the Canadian border and ends in
Key
West, Florida.
Boncit win
RAG: It starts in
Fort
Kent, Maine, at the Canadian border and ends in
New
West, Florida.
Oulput: starts in Fort Kent, Maine, at the Canadian border and ends in Key West, Florida.
Figure 5: Case study for collaborative generation between pure LLM and RAG at token level in our Tok-RAG. Pure LLM and RAG generate the texts in parallel at token level. At the step that pure
LLM and RAG generate the different tokens, Tok-RAG use our theoretical results in Theorem to compare the benefit and detriment. If benefit is greater than detriment, the token from RAG is selected, otherwise, the token from pure LLM is selected. The selected tokens are marked by green color and bold. The discarded tokens are marked by gray. The orange arrow represents the direction of token selection and usage. The selected tokens are used for the next step generation of both pure LLM and RAG.
27


Reviews:
[{"review_id": "gmc4gDwZiB", "reviewer": "Reviewer_aryh", "summary": "In RAG settings, retrieved passages may contradict the parametric memory of an LLM (obtained during the LLM training process). The paper explores how to reason and trade-off between the two sources of information in a more optimal manner compared to existing methods. The authors provide a theoretical framework for understanding this in terms of a latent concept variable $z$ and later connect it to a practical method that relies on calculating some of the required quantities by peeking inside the transformer architecture. Without any additional training or models, but with a whitebox approach, the method is able to outperform other RAG baselines that attempt this task.", "questions": "1. The one question that completely blocked me from understanding the paper was the definition of $p_R$. It first appears in Eqns 36, 37 and is defined as \"$p_R(·)$ is the distribution of the retrieved texts\". Can the authors more rigorously define this quantity? What is the probability over? i.e. is $p_R(r_i)$ the probability of $r_i$ being retrieved? if $p$ stands for probability of an event, then is it simply $p(r_i)$? Could it just be a uniform distribution then? If not, is there a \"retriever\" conditioned on a \"query\" that determines this distribution? In which case, $p_R(r_i) = p(r_i | q)$. And neither of these definitions make sense for Eqn 42, where $p_R(x_i|x_{1:i−1})$ appears. What does it even mean? \" distribution of the retrieved texts where the event is x_i conditioned on the previous tokens\"? This blocked me from making sense of the proofs beyond Theorem 1. \n\n2. Assumption 1 implies that there exists **some** hidden state $h$ lower bounds it that $p(x|h, z^∗) > c_1 > 0$. Then how is the variable used $c_1$ used in the denominator of Eqn 29 where the sum is over all $h$? If the **some** $h = h^*$, then the denominator can be said to be $c_1 * p(h^*| R, z^*)$ (note the lack of a summation), but this quantity can be arbitrarily small as a function of $z^*$ and $R$. Thus bringing into question the $O(1)$ claim that is made in Eqn 33.\n\n3. Assumption 2 is even more unclear in the definition. Until now there was no mention of the form of $R$ and this is the first time delimiters are introduced. What is the \"delimiter hidden state\" precisely? Are there any assumptions on input format of $R$. Ideally, in a theoretical justification, the format of the variables should make no difference. It is also very unclear how this assumption leads to $O(1)$ in line 765. \n\n4. It took me a while to understand how Eqn 33 could be potentially derived from Eqn 32. It would be prudent to be more rigorous here for future readers. \n\n5. Even if I assume theorem 1 is correct, in theorem 2, authors claim \"approximate positive correlation\", but I don't think this is a standard or well defined term. In lines 978 - 983, the authors go from upper and lower bounds to the claim on \"approximate correlation\" which doesn't seem theoretically sound. \n\n6. Just do double check, is Eqn 29 the same as Eqn 30?\n\nOverall, the paper would be a lot easier to understand and review if there weren't so many leaps of logic in the proofs.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper has a fair bit of novel theoretical underpinnings for thinking about retrieved knowledge in terms of a latent concept variable. And if sound and practical, there is significance to the findings especially when it comes to merging two sources of information. The writing is overall clear (except for the items in weaknesses) and quality is fair as well.", "weaknesses": "The paper tries to do two things at once, provide both theoretical and practical justification. However, the theoretical justification is at best very convoluted and at worst incorrect. The practical justification is a good start, but the authors test on QA datasets with factoid answers potentially hiding the fact that their method may not generalize outside of this setting. It would be more convincing if the authors did one thing (either theoretical or practical) in a sound and convincing manner rather than attempting both but being unconvincing. In its current form, even if correct and significant, this paper isn't at the acceptance threshold for ICLR. I am willing to change my score if I'm convinced about at least one aspect (theoretical or practical) and that is highlighted in the paper. \n\n1. The theoretical justification needs to be more rigorous and clear (see questions). There are many places where new quantities are defined in an unclear manner which make the proofs hard to follow. There are a couple of assumptions that are likely incorrectly applied or unclearly defined. And terms like \"approximate positive correlation\" appear in a theorem which to the best of my understanding isn't a standard term. While I suspect the final conclusion that authors draw (Eqn 12) could be practically correct, I am not convinced that the theoretical justification for it is sound. \n\n2. While the authors do not need an external model or a separate training process, they do make a whitebox assumption, i.e. the underlying language model is a transformer and they have access to the layers and logits. This needs to be made the abstract/intro where they downplay other methods but do not mention this limitation of their method. \n\n3. Central to the practical implementation of their method are Eqn 19 and the authors measure performance on factiod QA tasks where the attention distributions are likely to be low entropy (peaky) because one passage likely has the answer and has a strong lexical overlap with the question. However, it isn't clear how much their practical implementation generalizes beyond factiod QA where the attention distribution could be complex and not clearly delineated. For instance, in lines 322-323 the authors claim that \"LLMs use the selected knowledge at the maximum point for distribution fusion to predict $x_i$, the attention shifts from $R$ to prefix $x_{1:i-1}$\" which seems like pure speculation about the mechanism as we don't have any justification for it and we also don't know that this speculated mechanism generalizes beyond QA tasks. \n\n4. Why do the the authors claim \"distribution completion\" = \"benefit\" and \"distribution contradiction\" = \"Detriment\"? It could be that the LLM was trained on data that was considered true at the time, but has changed since. In which case the distribution contradiction is beneficial and desirable. The claim about benefit vs detriment is orthogonal to the point of the paper and should be avoid.  \n\n5. There is no way a reader can understand this paper without reading the appendix. At least the important assumptions need to be elevated to the main paper, even if the math isn't.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OcLo5LuXDa", "reviewer": "Reviewer_s84W", "summary": "This paper provides a theory to explain and trade off the 'benefit' and 'detriment' in next token prediction of RAG. Specifically, the author uses the distribution of the LLM's knowledge and the distribution of retrieved texts to calculate the distribution difference, which can trade off the 'benefit' and 'detriment'. Then, the author employs a series of theoretical analyses and some experiments to explain how to compare the values of 'benefit' and 'detriment'. Finally, the author proves the effectiveness of the method and supports the theoretical findings with experiments.", "questions": "1. For the calculation of the $Att$ score, it appears that the LLM is used for inference, and then the $Att$ score is computed for each layer. Essentially, it is still based on the $Wq$ and $W_k$ trained by the LLM itself? This suggests that the distribution of retrieved texts $p_R(x_i|x_{1:i-1})$ has a certain correlation with the knowledge distribution of the LLM $p(x_i|x_{1:i-1})$, and the same applies to word similarity. Is it possible to provide a more rigorous proof of the relationship between  $p_R(x_i|x_{1:i-1})$ and the LLM? For instance, analyzing the relationship between the knowledge distribution $p(x_i,x_{1:i-1})$ of LLM A itself and the retrieval texts distribution $p_R(x_i|x_{1:i-1})$ of LLM B when LLM B performing RAG.\n2. In Equation 16, $p_R(x_i|x_{1:i-1})$ is approximated using element-wise multiplication of $Att$ and $WordSim$. Are there other methods that could be used to combine $Att$ and $WordSim$ that could improve the estimation of $p_R(x_i|x_{1:i-1})$. For example, using element-wise addition. You could provide experimental or theoretical explanations.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper connects theory with practice, and the formula derivation is logically clear. Most of the papers are well written and explanatory, and are supported by experiments.\n2. The standpoint of this paper is novel, using the perspective of latent variable models to explain RAG, and analyzing the distribution differences between LLM and external knowledge.", "weaknesses": "1. In Equation 2, it seems a bit forced to split it into two terms, even though we know that distribution fusion is not a simple addition of distributions.\n2. From Equation 4 to 5, the symbol is used incorrectly; An equal cannot be used.\n3. In Equation 7, the $P_r$  is not explained. Is it $p_R(r)$?\n4. In Section 3.1, during the exploratory experiments on the distribution of retrieved texts $p_R(x_i|x_{1:i-1})$, it is mentioned that this distribution can be approximated using Equations 15 and 16. However, the calculation of the $Att$ score is not clearly explained. It still uses the QKV matrix from LLM? Appendix H only compares different LLMs and datasets, lacking a discussion on the calculation of the $Att$ score. Similarly, the hidden layer state $h$ in Equation 13 also originates from the pre-trained LLM itself.\n5. The paper does not evaluate Tok-RAG on LLMs with 33B and 65B scales. Since Tok-RAG generates texts in parallel, it means that two models need to be run at the same time, which may lead to a sharp increase in the demand for computing resources, especially when using larger-scale LLM.\n\nOthers:\n1. In Equation 2, it is better to write the right side of the equation as $ = p(..) + \\int .. dz$ to avoid misunderstandings.\n2. There is a mismatch in the notation of log, $ z*$ in Equation 7 compared to Equations 6 and 8. They should be standardized and modified to use $log$ and $z^*$.\n3. About table format, the titles of the tables should be placed above the tables. Additionally, the sequence of the tables in the paper goes directly from Table 1 to Table 3. The order of the tables should be revised accordingly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bHUIzT1AAN", "reviewer": "Reviewer_Rtos", "summary": "This paper explores an interesting and important problem of detecting the benefit and detriment of RAG from a theoretical perspective and propose a method called token-RAG. The authors first develop a theory to understand the benefit (retrieval can find valuable external information) and detriment (retrieved documents can be misleading or noisy) at token level. They find that the benefit and detriment can be compared by measuring the relative value of D and M, which is easily calculated in practice. However, calculating p_R for D and M is hard. The authors then introduce a heuristic method to use the attention score and word embedding similarity inside some specific LLM layers to estimate p_R. They finally conduct experiments on several datasets to demonstrate the effectiveness of the proposed token-RAG.", "questions": "See the weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper provides a theory to understand the benefit and detriment of the retrieved documents in RAG.\n- The paper introduces a method called token-RAG to leverage benefit and prevent detriment.\n- The authors conduct experiments on several datasets.", "weaknesses": "- The meaning of some notations is not well explained. What is z*?\n\n- Some formulations may be incorrect. Is there any problem with Eq.(2)? Given that p(z|R,x1:{i-1}) is a continuous function, how can you get the term corresponding to p(z|R,x1:{i-1}) out from the integral?\n\n- Why can you go from (4) to (5)? I believe it should not be “=”. Why p(R,x1:{i-1}|z*) is a constant?\n\n- The estimation of p_R(x_i|x_1:{i-1}) seems to be based on heuristics. I am not persuaded that this method can generalize to different datasets and domains.\n\n- In (17) and (18), why use the average word embedding rather than use the real predicted probability of the words? The latter seems to be naturally describing p in M and D.\n\n- The exact token-RAG method is not very clear. From (19), does it mean that if benefit wins, token-RAG will generate the token predicted by RAG, otherwise by LLM directly without RAG?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)