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

## Abstract

Abstract Retrieval-augmented generation (RAG) utilizes retrieved texts to enhance large language models (LLMs). Studies show that while RAG provides valuable external information (benefit), it may also mislead LLMs (detriment) with noisy or incorrect retrieved texts. Although many existing methods attempt to preserve benefit and avoid detriment, they lack a theoretical explanation for RAG. The benefit and detriment in the next token prediction of RAG remain a ’black box’ that cannot be quantified or compared in an explainable manner, so existing methods are data-driven, need additional utility evaluators or post-hoc. This paper takes the first step towards providing a theory to explain and trade off the benefit and detriment in RAG. First, we model RAG as the fusion between distribution of 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, 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.

## 1 Introduction

Retrieval-augmented generation (RAG) has shown promising performance in enhancing Large Language Models (LLMs) by integrating retrieved texts  (Xu et al., ; Shi et al., ; Asai et al., ; Ram et al., ). Studies indicate that while RAG provides LLMs with valuable additional knowledge (benefit), it also poses a risk of misleading them (detriment) due to noisy or incorrect retrieved texts (Ram et al., ; Xu et al., ; ; Jin et al., ; Xie et al., ; Jin et al., ). Existing methods attempt to preserve benefit and avoid detriment by adding utility evaluators for retrieval, prompt engineering, or fine-tuning LLMs (Asai et al., ; Ding et al., ; Xu et al., ; Yoran et al., ; Ren et al., ; Feng et al., ; Mallen et al., ; Jiang et al., ). However, existing methods are data-driven, need evaluator for utility of retrieved texts or post-hoc. A theory-based method, focusing on core principles of RAG is urgently needed, which is crucial for consistent and reliable improvements without relying on additional training or utility evaluators and improving our understanding for RAG.

Figure: Figure 1: 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 our 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 or utility evaluators.
Refer to caption: x1.png

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 a novel method to preserve benefit and avoid detriment based on our theoretical findings. Specifically, this paper pioneers in modeling next token prediction in RAG as the fusion between the distribution of LLM’s knowledge and the distribution of retrieved texts as shown in Figure 1. Our theoretical derivation based on this formalizes the core of this fusion as the subtraction between two terms measured by the distribution difference: one is distribution completion and the other is distribution contradiction. Further analysis indicates that the distribution completion measures how much out-of-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 LLM’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 as subtraction.

In this way, we successfully decouple benefit and detriment from next token prediction of RAG and formalize them as the subtraction between two 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 output and the representation of retrieval texts. Finally, we establish a theory to predict the actual effect of RAG at 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 of the token outweighs the detriment.

Based on our theoretical results, we propose a practical novel method called Tok-RAG that can achieve collaborative generation between pure LLM and RAG at token level to preserve benefit and avoid detriment without any training or additional modules for utility evaluation of retrieved texts. As shown in Figure 1, 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 theoretical 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, which indicates that our theoretical results are essential and fundamental for RAG. The main contributions of this paper are:

$\bullet$ 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.

$\bullet$ We model RAG as distribution fusion and formalize the core of this fusion as the subtraction between distribution completion and distribution contradiction, our analysis shows the former is benefit and the latter is detriment for next token prediction of RAG. This subtraction describes the trade-off between the value of external knowledge and its potential risk of misleading LLM.

$\bullet$ 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.

$\bullet$ Based on the theoretical results, we propose a practical novel method to enable pure LLM and RAG to collaboratively generate at token level. Experimental results on real-world tasks across different LLMs show the effectiveness of our method and support our theoretical results.

## 2 Theory to Understand Benefit and Detriment at Token Level

RAG exhibits the duality, 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 to use latent variable model to model the next token prediction in RAG as the fusion between the distribution of LLM’s knowledge and the distribution of retrieved texts. Second, our theoretical derivation formalizes the core of this fusion as 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 (i.e., the comparison between benefit and detriment) at token level. Figure 2 is the derivation path of our theory.

### 2.1 Model Next Token Prediction in RAG as Distribution Fusion

To begin the analysis, inspired by previous studies that prove LLMs implicitly perform latent variable inference (Xie et al., ; Zhang et al., ; Wang et al., ), we first propose to analyse RAG by latent variable model to model the next token prediction in RAG as the fusion between the distribution of LLM’s knowledge and the distribution of retrieved texts. Specifically, given the prefix $x_{1:i-1}=\{x_{1},x_{2},...x_{i-1}\}$ generated from time step $1$ to $i-1$, from the perspective of the latent variable model, the probability distribution of the token $x_{i}$ at the $i$-th step is described as this:

$$ $p(x_{i}|x_{1:i-1})=\int_{\mathcal{Z}}p(x_{i}|x_{1:i-1},z)p(z|x_{1:i-1})\,dz,$ (1) $$

in which $\mathcal{Z}$ is the space of high dimensional concept variable, $p(z|x_{1:i-1})$ is the probability that the model samples latent concept $z$ from $\mathcal{Z}$ given the prefix $x_{1:i-1}$, and $p(x_{i}|x_{1:i-1},z)$ means the probability for token $x_{i}$ conditioned on the prefix $x_{1:i-1}$ and the sampled latent concept $z$. $p(x_{i}|x_{1:i-1})$ can be obtained by integrating over all latent concepts from the space $\mathcal{Z}$. Latent variable model has been applied in many methods such as LDA (Blei et al., ). 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., ; Zhang et al., ). Inspired by this, we analyse the next token prediction of RAG given prefix $x_{1:i-1}$ as sampling the shared Retrieved Concept $z^{*}$ from the input retrieved texts list $R=\{r_{1},r_{2},...,r_{n}\}$ ($r_{i}$ is a retrieved passage), and then predicting $p(x_{i}|R,x_{1:i-1})$, which can be formalized as:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i-1})\,dz$ (2) $\displaystyle=\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i -1})\,dz+p(x_{i}|R,x_{1:i-1},z^{*})p(z^{*}|R,x_{1:i-1}).$ $$

Equation 2 describes RAG as distribution fusion. The first term is the prediction that marginalizes out all latent concepts except $z^{*}$, which is the knowledge distribution in LLMs. The second term is the prediction that is only conditioned on $z^{*}$, which is the knowledge distribution from retrieved texts.

### 2.2 Formalize and Explain benefit and detriment by Distribution Difference

Figure: Figure 2: Derivation path of our theory. Reference: Equation 9, Theorem 1, 2, 3 and Corollary 1, 2, 3.
Refer to caption: x2.png

In this section, we make further derivation on the distribution fusion in Equation 2 to formalize the core of this fusion as the subtraction between two terms measured by the distribution difference: one is distribution completion, and the other is distribution contradiction. Our analysis reveals that for the next token prediction, distribution completion is actually the benefit and distribution contradiction is the detriment. Specifically, inspired by (Xie et al., ), the Equation 2 can be transformed as (detailed proof can be found in Appendix A):

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i-1})\,dz$ (3) $\displaystyle\propto\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(R,x_{1:i-1}|z)p( z)\,dz$ (4) $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\textrm{exp}(v(z))p(z)\, dz,\quad v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ (5) $$

Define $r_{i}$ is a passage in the retrieved texts list $R$, we can get (see detailed proof in Appendix B):

$$ $\displaystyle v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ $\displaystyle\approx\textrm{log}\frac{\prod_{i=1}^{n}O(1)p(r_{i}|z)}{\prod_{i= 1}^{n}O(1)p(r_{i}|z^{*})}$ (6) $\displaystyle\rightarrow n*\frac{1}{n}\sum_{i=1}^{n}log\frac{p(r_{i}|z)}{p(r_{ i}|z*)}=n*\mathbb{E}_{r\sim P_{r}}\left[log\frac{p(r|z)}{p(r|z^{*})}\right]$ (7) $\displaystyle\propto p_{R}(r)\textrm{log}\frac{p(r|z)}{p(r|z^{*})}=p_{R}(r) \textrm{log}\frac{p_{R}(r)}{p(r|z^{*})}-p_{R}(r)\textrm{log}\frac{p_{R}(r)}{p( r|z)}$ (8) $\displaystyle=-(\underbrace{\text{KL}(p_{R}(r)\|p(r|z))}_{\textbf{Distribution Completion: Benefit}}-\underbrace{\text{KL}(p_{R}(r)\|p(r|z^{*}))}_{\textbf{ Distribution Contradiction: Detriment}}),$ (9) $$

$p_{R}(\cdot)$ is the distribution of the retrieved texts, $p(\cdot)$ is the distribution of the LLM’s knowledge. $v(z)$ is an important term in distribution fusion because it reflects the proportion between the latent concept from the space of LLMs and from the retrieved texts. Details are in Appendix C.

Discuss the Benefit and Detriment Based on the Theoretical Results. In Equation 9, the first term represents the distribution difference (KL divergence) between retrieved texts ($p_{R}(r)$) and LLM’s knowledge ($p(r|z)$) given the concept $z$, which is sampled from $\mathcal{Z}$ (latent variables in LLM). This can be defined as distribution completion that measures how much out-of-distribution knowledge that retrieved texts provide to LLM in the prediction of the token $x_{i}$. This term is actually the benefit for the prediction of token $x_{i}$ in RAG. This is because that if the retrieved text $r$ is perfect, $p_{R}(r)$ would be infinitely close to the ground-truth distribution. The larger the difference between $p_{R}(r)$ and $p(r|z)$, the more the knowledge distribution about $x_{i}$ in the LLM’s knowledge deviates from the ground truth. This means the retrieved texts can provide more valuable out-of-distribution knowledge to the LLM to predict $x_{i}$, 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 distribution contradiction can be used to measure this risk. It represents the distribution difference (KL divergence) between the retrieved texts ($p_{R}(r)$) and LLM’s knowledge given the concept $z^{*}$, which is sampled from retrieved texts ($p(r|z^{*})$). $p(r|z^{*})$ is the prediction made by LLM conditioned on the concept $z^{*}$ sampled from the retrieved texts. If the external knowledge in the retrieved texts contradicts LLM’s knowledge, $p(r|z^{*})$ will deviate from $p_{R}(r)$ (the actual distribution of the retrieved texts). Therefore, the difference between $p(r|z^{*})$ and $p_{R}(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 $x_{i}$ in RAG.

###### Corollary 1 .

Two terms about distribution difference in Equation 9 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.3 Actual Effect of RAG can be Predicted at Token Level

Based on 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 1 and Theorem 2 from Equation 2:

###### Theorem 1 .

Define $\mathcal{D}=\|p(x_{i}|R,x_{1:i-1})-p_{R}(x_{i}|x_{1:i-1})\|_{1}$ to measure the difference between the distribution of RAG ($p(x_{i}|R,x_{1:i-1})$) and the distribution of retrieved texts ($p_{R}(x_{i}|x_{1:i-1})$) in token prediction of $x_{i}$ conditioned on prefix $x_{1:i-1}$. Both benefit and detriment are important terms of the upper and lower bounds of $\mathcal{D}$, which can be described as:

$$ $\displaystyle\|\Phi\|_{1}-\sqrt{2\textnormal{KL}(p_{R}(r)\|p(r|z^{*}))}$ $\displaystyle\leq\mathcal{D}\leq\|\Phi\|_{1}+\sqrt{2\textnormal{KL}(p_{R}(r)\| p(r|z^{*}))},$ (10) $\displaystyle\Phi\approx\alpha\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1 },z)\textnormal{exp}$ $\displaystyle\left[-(\underbrace{\textnormal{KL}(p_{R}(r)\|p(r|z))}_{ \textnormal{{benefit}}}-\underbrace{\textnormal{KL}(p_{R}(r)\|p(r|z^{*}))}_{ \textnormal{{detriment}}}\right]p(z)\,dz,$ $$

in which $\alpha$ is a constant. Our detailed proof of Theorem 1 can be found in Appendix D.

###### Theorem 2 .

$\mathcal{D}$ is the difference, so $\frac{1}{\mathcal{D}}$ can be treated as similarity between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$. The result of benefit minus detriment is approximately positively correlated with $\frac{1}{D}$:

$$ $\underbrace{\textnormal{KL}(p_{R}(r)\|p(r|z))}_{\textnormal{{benefit}}}- \underbrace{\textnormal{KL}(p_{R}(r)\|p(r|z^{*}))}_{\textnormal{{detriment}}} \propto\frac{1}{\mathcal{D}}.$ (11) $$

Our detailed proof of Theorem 2 is in Appendix E.

###### Corollary 2 .

The difference between values of benefit and detriment in Equation 11 indicates the extent to which the benefit (value of external knowledge) outweighs the detriment (potential risk of misleading LLM) in the prediction of token $x_{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 $x_{i}$ (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 $\frac{1}{D}$. Therefore, the value of $\frac{1}{D}$ at which the benefit minus detriment equals zero is an important threshold. A $\frac{1}{D}$ value greater than this threshold indicates that benefit exceeds detriment, while a value less than this threshold indicates that detriment outweighs benefit. We derive Theorem 3 to identify this dividing point and map the value order between benefit and detriment of token $x_{i}$ to the relationship between representation similarity, which can be calculated in practical applications:

###### Theorem 3 .

Define $\mathcal{M}=\|p(x_{i}|R,x_{1:i-1})-p(x_{i}|x_{1:i-1})\|_{1}$ to measure the difference between distribution of RAG ($p(x_{i}|R,x_{1:i-1})$) and pure LLM ($p(x_{i}|x_{1:i-1})$), so $\frac{1}{\mathcal{M}}$ can be treated as the similarity between them. $\frac{1}{\mathcal{D}}=\frac{1}{\mathcal{M}}$ is the dividing point in which benefit is equal to detriment, and the value order between $\frac{1}{\mathcal{D}}$ and $\frac{1}{\mathcal{M}}$ can indicate the value order between benefit and detriment as:

$$ $\displaystyle J=\begin{cases}\textnormal{KL}(p_{R}(r)\|p(r|z))<\textnormal{KL} (p_{R}(r)\|p(r|z^{*}))\text{, detriment outweighs benefit.}&\text{if }\frac{1} {\mathcal{D}}<\frac{1}{\mathcal{M}}\\ \textnormal{KL}(p_{R}(r)\|p(r|z))=\textnormal{KL}(p_{R}(r)\|p(r|z^{*}))\text{, detriment is equal to benefit.}&\text{if }\frac{1}{\mathcal{D}}=\frac{1}{ \mathcal{M}}\\ \textnormal{KL}(p_{R}(r)\|p(r|z))>\textnormal{KL}(p_{R}(r)\|p(r|z^{*}))\text{, benefit outweighs detriment.}&\text{if }\frac{1}{\mathcal{D}}>\frac{1}{ \mathcal{M}}\end{cases}$ (12) $$

Our detailed proof of Theorem 3 can be found in Appendix F. Equation 12 is a novel principle that can compare the values of benefit and detriment in RAG at token level. It does not rely on additional modules to access the utility of retrieved texts or training but simply compares $\frac{1}{\mathcal{D}}$ and $\frac{1}{\mathcal{M}}$.

###### Corollary 3 .

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(x_{i}|R,x_{1:i-1})$, $p(x_{i}|x_{1:i-1})$, and $p_{R}(x_{i}|x_{1:i-1})$, without additional modules to access the utility of retrieved texts.

The next section introduces how to apply above theoretical findings to improve RAG in practice.

###### Theorem 1 .

###### Theorem 2 .

###### Corollary 2 .

###### Theorem 3 .

###### Corollary 3 .

## 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 1. 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 $\frac{1}{\mathcal{D}}$ and $\frac{1}{\mathcal{M}}$ in Equation 12. The terms related to the comparison between $\frac{1}{\mathcal{D}}$ and $\frac{1}{\mathcal{M}}$ consist of three parts: (1) $p(x_{i}|R,x_{1:i-1})$ can be directly obtained from the prediction of RAG; (2) $p(x_{i}|x_{1:i-1})$ can be directly obtained from the prediction of pure LLM; (3) however, the distribution of retrieved texts conditioned on the prefix $x_{1:i-1}$, $p_{R}(x_{i}|x_{1:i-1})$, is hard to 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.1 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 $p_{R}(x_{i}|x_{1:i-1})$ 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 $p_{R}(x_{i}|x_{1:i-1})$. To address these challenges, in the following parts, we first explore the operating mechanism of RAG and then propose a novel method that dynamically determines the layers where distribution fusion occurs and use the signal from retrieved texts in these layers as $p_{R}(x_{i}|x_{1:i-1})$.

Figure: Figure 3: Attention score for $x_{i}$ (blue line) and difference of word distribution change (yellow line) vary with layers. stage 1: Lexical and Syntactic. stage 2: Text Matching. stage 3: Distribution Fusion.
Refer to caption: x3.png

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 $x_{i}$ from the retrieved texts $R$. The second is distribution fusion, which means fusing 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 $x_{1:i-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 at 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 3. For the second challenge, we use the matching information in the layer where distribution fusion starts to approximate the distribution $p_{R}(x_{i}|x_{1:i-1})$. Our experiments about these findings are introduced in Exp 1. and Exp 2.. Experiments are conducted based on LLaMA-2-7B on Natural Question dataset. We also perform these experiments using more LLMs (OPT-6.7B, Mistral-7B) on more datasets (TriviaQA, WebQ and Squad). The conclusions of them are consistent with LLaMA-2-7B 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 $x_{i}$ by the attention score between token $x_{i}$ and the tokens in the retrieved texts $R$. We analyze how the sum of attention scores from token $x_{i}$ 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 0-$5$), which is mainly because LLMs capture the low-level lexical and syntactic information on the entire input (Tenney et al., ). (2) The value first decreases, then increases to a maximum point in the middle layers ($5$-$23$), which is mainly because LLMs select the relevant semantics that can be used to generate $x_{i}$ from $R$ 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 $x_{i}$, the attention shifts from $R$ to prefix $x_{1:i-1}$.

Exp 2. For distribution fusion, since distribution fusion is often marked by a change in word distribution (Bengio et al., ), we identify the occurrence of distribution fusion in RAG by comparing the change in word distribution between pure LLM and RAG.  (Chuang et al., ; Schuster et al., ) 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 $\phi$ as $\phi(h_{i}^{l})$, in which $h_{i}^{l}$ is the hidden states for token $x_{i}$ in the $l$-th layer. We then measure the word distribution change in the $l$-th layer by Jensen-Shannon Divergence (JSD) between $\phi(h_{i}^{l-1})$ and $\phi(h_{i}^{l})$ as: $C=\textrm{JSD}(\phi(h_{i}^{l-1})\|\phi(h_{i}^{l}))$. The difference of word distribution change between pure LLM and RAG in the $l$-th layer is described as:

$$ $\displaystyle D^{l}=|\textrm{JSD}(\phi(\tilde{h}_{i}^{l-1})\|\phi(\tilde{h}_{i }^{l}))-\textrm{JSD}(\phi(h_{i}^{l-1})\|\phi(h_{i}^{l}))|,$ (13) $$

in which $\tilde{h}_{i}^{l-1}$ and $\tilde{h}_{i}^{l}$ are from RAG, $h_{i}^{l-1}$ and $h_{i}^{l}$ are from pure LLM. The yellow line in Figure 3 shows $D^{l}$ is very small in the shallow and middle layers (0 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(x_{i}|R,x_{1:i-1})$, the layer where distribution fusion starts can be located by detecting the turning point in Figure 3. Specifically, we use $f(l)$ to denote the attention score for $x_{i}$ and $g(l)$ to denote the difference of word distribution change in Equation 13 varies with layer $l$. The layer where distribution fusion starts is:

$$ $\displaystyle l^{*}=\lfloor\frac{1}{2}(\arg\max_{l}f(l)+\min\{l:g(l)>a\})\rfloor.$ (14) $$

The first term is the $l$ that maximizes $f(l)$, which is the third turning point in the blue line of Figure 3. The second term means the minimum $l$ value for which $g(l)$ is greater than $a$ (hyperparameter, can be set to 5e-7 according to our statistics), which is the turning point in the yellow line of Figure 3. We take the average of the two values and round down as the layer $l^{*}$ where distribution fusion starts.

Matching as Distribution. The matching information between $R=[rt_{1},rt_{2},...,rt_{m}]$ ($rt$ is the token in $R$) and token $x_{i}$ at turning point (the $l^{*}$-th layer) can be used to approximate the distribution $p_{R}(x_{i}|x_{1:i-1})$ of the retrieved texts $R$ conditioned on $x_{1:i-1}$. The matching information consists of two parts, one is the attention score, which can measure the matching between retrieved tokens and current token $x_{i}$ 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_{i}$ at the word distribution level:

$$ $\displaystyle Att$ $\displaystyle=\textrm{softmax}\left(\frac{(\tilde{h}_{i}^{l^{*}}W_{q})(\tilde{ h}_{1:m}^{l^{*}}W_{k})^{T}}{\sqrt{d_{k}}}\right),WordSim=\textrm{softmax}\left ((x_{i}^{l^{\prime}-l^{*}}A)(rt_{1:m}^{l^{*}}A)^{T}\right),$ (15) $$

$W_{q}$ and $W_{k}$ are matrices in attention (Vaswani et al., ), $\tilde{h}_{i}^{l^{*}}$ is the hidden state of token $x_{i}$ and $\tilde{h}_{1:m}^{l^{*}}$ are hidden states of $R$. $A$ is word embedding matrix, $x_{i}^{l^{\prime}-l^{*}}$ is the token with the largest logits increase in word distribution from layer $l^{*}$ to the final layer $l^{\prime}$, $rt_{1:m}^{l^{*}}$ are tokens in $R$. Then:

$$ $\displaystyle p_{R}(x_{i}|x_{1:i-1})=\textrm{softmax}\left(Att\odot WordSim \right),\odot\textnormal{ is element-wise multiplication.}$ (16) $$

### 3.2 Token-Level Comparison between Benefit and Detriment in Practice

Equation 12 shows that the relationship between $\textnormal{Sim}(p(x_{i}|R,x_{1:i-1}),p_{R}(x_{i}|x_{1:i-1}))$ and $\textnormal{Sim}(p(x_{i}|R,x_{1:i-1}),p(x_{i}|x_{1:i-1}))$ indicates the value order between benefit and detriment ($\textnormal{Sim}(\cdot,\cdot)$ is the similarity). We propose to use the token semantics as the representation for $p(x_{i}|R,x_{1:i-1})$, $p_{R}(x_{i}|x_{1:i-1})$ and $p(x_{i}|x_{1:i-1})$ and use cosine to compute the similarity. It not only follows the principle of Equation 12 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(x_{i}|R,x_{1:i-1})$ as $\mathbf{w}_{RAG}=\frac{1}{\sum p^{\prime}}\sum_{(p^{\prime},\mathbf{w})\in
\mathbb{V}}p^{\prime}\mathbf{w}$,
for each token in vocabulary $\mathbb{V}$, $p^{\prime}$ is its logits from $p(x_{i}|R,x_{1:i-1})$ and $\mathbf{w}$ is its word embedding. We can also use this to get the weighted average word embedding $\mathbf{w}_{LLM}$ for $p(x_{i}|x_{1:i-1})$ and $\mathbf{w}_{IR}$ for $p_{R}(x_{i}|x_{1:i-1})$. The similarity between them can be calculated via cosine similarity as:

$$ $\displaystyle\textnormal{Sim}(p(x_{i}|R,x_{1:i-1}),p_{R}(x_{i}|x_{1:i-1}))$ $\displaystyle=\textrm{cos}(\mathbf{w}_{RAG},\mathbf{w}_{IR})$ (17) $\displaystyle\textnormal{Sim}(p(x_{i}|R,x_{1:i-1}),p(x_{i}|x_{1:i-1}))$ $\displaystyle=\textrm{cos}(\mathbf{w}_{RAG},\mathbf{w}_{LLM})$ (18) $$

Combining our theoretical analysis of Theorem 1, 2 and 3, we can derive this principle to compare the values of benefit and detriment brought by RAG to the token $x_{i}$ in practical applications:

$$ $s=\begin{cases}{\color[rgb]{0.0,0.69,0.0}\textrm{benefit win}}&\text{if } \textrm{cos}(\mathbf{w}_{RAG},\mathbf{w}_{IR})\geq\textrm{cos}(\mathbf{w}_{RAG },\mathbf{w}_{LLM}),\\ {\color[rgb]{1,0,0}\textrm{detriment win}}&\text{if }\textrm{cos}(\mathbf{w}_{ RAG},\mathbf{w}_{IR})<\textrm{cos}(\mathbf{w}_{RAG},\mathbf{w}_{LLM}),\end{cases}$ (19) $$

## 4 Experiments

### 4.1 Experimental Details

**Table 1: Accuracy on 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-RAG does not need any training or additional modules while baselines need.**
| Methods | Train<br>LLM | Add<br>Module | TriviaQA | WebQ | Squad |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Ratio of Hard Negative Passages | Ratio of Hard Negative Passages | Ratio of Hard Negative Passages |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| 100% | 80% | 60% | 40% | 20% | 0% | 100% | 80% | 60% | 40% | 20% | 0% | 100% | 80% | 60% | 40% | 20% | 0% |  |  |  |
| Standard RAG | no ✔ | no ✔ | 43.8 | 67.0 | 71.3 | 76.2 | 78.2 | 81.9 | 23.9 | 35.8 | 40.6 | 43.4 | 48.4 | 53.1 | 8.6 | 31.0 | 43.2 | 53.0 | 58.8 | 67.2 |
| NLI+RAG | no ✔ | need ✗ | 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 | no ✔ | need ✗ | 48.2 | 68.3 | 72.5 | 76.7 | 81.5 | 82.2 | 25.6 | 37.4 | 41.9 | 46.2 | 51.5 | 54.9 | 7.4 | 28.7 | 39.6 | 50.7 | 53.2 | 61.1 |
| RetRobust | need ✗ | no ✔ | 49.2 | 67.3 | 72.9 | 77.5 | 79.4 | 82.3 | 30.0 | 38.9 | 42.5 | 48.2 | 49.8 | 54.3 | 10.5 | 30.8 | 43.3 | 52.5 | 58.4 | 66.0 |
| Self-RAG | need ✗ | no ✔ | 43.0 | 68.7 | 73.5 | 76.4 | 80.8 | 82.2 | 18.3 | 34.8 | 42.2 | 47.2 | 51.3 | 57.0 | 5.5 | 27.8 | 38.9 | 46.4 | 52.5 | 58.3 |
| INFO-RAG | need ✗ | no ✔ | 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) | no ✔ | no ✔ | 53.5 | 72.9 | 77.6 | 81.3 | 83.4 | 85.7 | 32.9 | 43.8 | 47.3 | 50.0 | 52.9 | 57.3 | 12.8 | 31.3 | 44.5 | 54.1 | 60.8 | 68.1 |

Practical RAG for Open-domain Q&A. One experiment is in a practical autoregressive generation setting for open-domain Q&A given retrieved texts with different qualities, it evaluates the robustness and performance of RAG. Baselines include the methods that use additional modules to filter irrelevant texts (NLI+RAG (Yoran et al., )) or as action triggers (CRAG (Yan et al., )), fine-tune LLMs for robust RAG (RetRobust (Yoran et al., ) and INFO-RAG (Xu et al., )) and fine-tune LLMs to dynamically retrieve and critique retrieved texts (Self-RAG (Asai et al., )). We follow baselines to use Cover-EM (Rosset et al., ) that indicates the accuracy as the metric.

Setup for Benefit-Detriment Comparison Experiment. The other is benefit-detriment comparison experiment. Given prefix $x_{1:i-1}$ and retrieved texts $R$, our motivation aims to
build a theory to predict the actual effect of RAG on the prediction of $x_{i}$, i.e., whether the positive impact of the retrieved texts $R$ on $x_{i}$ (benefit) outweighs the potential risk of misleading LLM (detriment). This is a binary classification task at token-level. To evaluate this task, 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 $x_{1:i-1}$ and the next token $x_{i}$.

(2) We input the prefix $x_{1:i-1}$ to LLM w/o RAG and LLM w/ RAG to get the next predicted token $a$ and $b$ respectively.

(3) If $b$ is $x_{i}$ but $a$ not, it means LLM w/ RAG performs better than LLM w/o RAG, so the benefit is greater than detriment, the ground-truth label is $1$.

(4) If $a$ is $x_{i}$ but $b$ not, it means LLM w/o RAG performs better than LLM w/ RAG, so the benefit is lower than detriment, the ground-truth label is 0 0.

We use the above method to traverse all sentences in the datasets to obtain prefix-ground-truth pairs 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 compare the values of benefit and detriment by comparing the hallucination degree at token level between RAG and pure LLM (details in Appendix I.1). Baselines include: (1) Logprobs-based (Kuhn et al., ), we use the value order between top-1 log-probability of the tokens output by pure LLM and RAG to determine the value order between benefit and detriment. (2) Uncertainty-based, we use Length-normalized Entropy (Malinin & Gales, ) to measure the uncertainty of the tokens and compare it between pure LLM and RAG. (3) Consistency-based, we run LLMs multiple times and calculate consistency scores among multiple answers using Lexical and Semantic Similarity (Lin et al., ; Chen et al., ). We compare these scores between pure LLM and RAG to indicate the comparison between benefit and detriment.

Datasets. For the open-domain Q&A experiment, since long-form tasks are not conducive to objectively and accurately evaluating the factual correctness of the answers, we use three short-form Q&A tasks including WebQuestions (WebQ) (Berant et al., ), TriviaQA (Joshi et al., ) and SQuAD v1.1 (Rajpurkar et al., ). For the token level binary classification task in benefit-detriment comparison experiment, we use three long-form generation tasks including long-form Q&A (ASQA (Stelmakh et al., )), people biographies generation (Bio (Min et al., )) and language modeling (Wikitext103 (Merity et al., )) and one short-form task includes Q&A (Natural Questions (Kwiatkowski et al., )).

Implementation details. As for retrieval in RAG, we follow (Xu et al., ) to use ColBERTv2 (Santhanam et al., ), an excellent generalizable model as the retriever, and use Wikipedia consisting of 21,015,324 passages (Karpukhin et al., ) 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 benefit-detriment comparison experiment and use greedy-decoding strategy for generation. More details of Tok-RAG and baselines are in Appendix I.

### 4.2 Experimental Results

Experiment on Open-domain Q&A. This experiment is under the practical autoregressive generation setting for open-domain Q&A.
Table 1 shows that in RAG given the retrieved texts with various qualities, our Tok-RAG shows better robustness and performance than baselines. 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 LLMs. This means our Tok-RAG achieves a 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. In this experiment, 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. Since Open-domain Q&A is a short-form task,

**Table 2: Comparison of methods based on GPU Memory, Time, and Performance.**
| Method | GPU (GB) ↓ | Time (s) ↓ | Per. (Acc) ↑ |
| --- | --- | --- | --- |
| Standard | 17.50 | 1.95 | 51.96 |
| NLI+RAG | 20.25 | 2.03 | 49.76 |
| CRAG | 20.60 | 2.15 | 51.53 |
| RetRobust | 17.50 | 2.00 | 52.98 |
| Self-RAG | 17.50 | 2.98 | 50.26 |
| INFO-RAG | 17.50 | 2.10 | 52.35 |
| Tok-RAG | 23.98 | 2.24 | 56.12 |

we calculate the accuracy by judging whether the ground truth appears exactly in the generated texts (Cover-EM (Rosset et al., )). LLM in this is LLaMA-2-7B.

**Table 3: Performance on comparing the valurs of benefit and detriment at token level. Significant test with p-value $\leq 0.05$ compared with all baselines are denoted as ‘$+$’. Setup in this experiment is introduced in Setup for Benefit-Detriment Comparison Experiment in Section 4.1.**
| LLMs | Methods | # Generation | Wikitext | ASQA | Bio | NQ |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Times | AUC | F1 | AUC | F1 | AUC | F1 | AUC | F1 |  |  |
| OPT-6.7B | Logprobs | 2 | 65.25 | 64.33 | 68.96 | 67.55 | 65.24 | 64.59 | 55.31 | 51.41 |
| Uncertainty | 2 | 64.12 | 63.50 | 66.14 | 63.96 | 65.78 | 64.60 | 56.03 | 52.15 |  |
| Consistency-Lexical | 10 | 64.01 | 62.17 | 69.42 | 67.04 | 65.41 | 65.28 | 55.06 | 51.13 |  |
| Consistency-Semantic | 10 | 65.93 | 64.22 | 70.11 | 69.50 | 65.76 | 64.37 | 56.24 | 52.88 |  |
| Tok-RAG (Ours) | 2 | 68.64^+ | 66.88^+ | 72.28^+ | 72.05^+ | 66.27^+ | 66.04^+ | 57.92^+ | 52.90^+ |  |
| Mistral-7B | Logprobs | 2 | 73.52 | 72.90 | 68.05 | 66.86 | 65.22 | 64.39 | 57.04 | 57.23 |
| Uncertainty | 2 | 73.72 | 72.71 | 67.47 | 65.63 | 65.59 | 65.83 | 57.19 | 57.10 |  |
| Consistency-Lexical | 10 | 72.15 | 70.44 | 69.16 | 67.33 | 64.79 | 64.33 | 56.95 | 54.37 |  |
| Consistency-Semantic | 10 | 73.98 | 72.26 | 70.05 | 69.54 | 65.68 | 65.12 | 57.43 | 56.12 |  |
| Tok-RAG (Ours) | 2 | 75.85^+ | 74.11^+ | 71.51^+ | 71.47^+ | 66.37^+ | 66.04^+ | 58.52^+ | 57.56^+ |  |
| LLaMA-2-7B | Logprobs | 2 | 73.47 | 72.95 | 68.50 | 68.04 | 62.11 | 60.94 | 67.40 | 69.24 |
| Uncertainty | 2 | 73.98 | 73.01 | 68.72 | 67.63 | 63.67 | 63.50 | 68.03 | 69.15 |  |
| Consistency-Lexical | 10 | 73.51 | 71.62 | 70.09 | 68.45 | 62.49 | 61.98 | 68.17 | 70.09 |  |
| Consistency-Semantic | 10 | 74.96 | 74.23 | 71.23 | 69.38 | 63.77 | 62.10 | 69.72 | 71.14 |  |
| Tok-RAG (Ours) | 2 | 81.89^+ | 80.42^+ | 76.96^+ | 76.80^+ | 64.08^+ | 64.19^+ | 70.50^+ | 72.45^+ |  |

Experiment on Benefit-Detriment Comparison.
Table 3 shows that our Tok-RAG achieves better performance in comparing the values of benefit and detriment at token level in RAG than baselines across different tasks and LLMs. Baselines compare the benefit and detriment by detecting the degree of hallucination, while our Tok-RAG can directly compare the benefit and detriment based on our theoretical analysis, which is more fundamental so it performs better. The detailed setup for this experiment can be found in Setup for Benefit-Detriment Comparison Experiment in Section 4.1.

Figure: Figure 4: AUC varies with layer.
Refer to caption: x4.png

Analysis on Computational Costs. Table 2 shows that out Tok-RAG achieves significant performance (accuracy) improvement with little increase of GPU memory and running time. In addition, the baselines all require training LLM or introducing additional modules, while our method is based on our solid theoretical 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 or GPU memory overhead. Experiments are performed on three Q&A datasets (TrviaQA, WebQ, Squad) with V100 GPU, the LLM is LLaMA-2-7B.

Ablation study. Figure 4 shows the effectiveness of our dynamic layer selection strategy in Equation 14 and supports our finding that RAG performs matching in middle layers. Figure 4 shows the AUC when $l^{*}$ in Equation 14 is set as a fixed value from 0 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. LLM used in this experiment is LLaMA-2-7B.

Case study. Case study of Figure 5 in Appendix J intuitively shows the collaborative generation between pure LLM and RAG in our Tok-RAG. Tok-RAG is effective to preserve benefit and avoid detriment at token level by dynamically selecting suitable tokens among pure LLM and RAG.

## 5 Related work

Robust RAG. To make LLMs robust in RAG to avoid the detriment caused from noisy in retrieved texts, some methods use additional modules to filter out irrelevant documents (Yoran et al., ; Yan et al., ). Some methods train LLMs to make them robust to noisy in retrieved texts (Xu et al., ; Yoran et al., ). Some methods let LLMs dynamically determine whether the query needs RAG (Asai et al., ; Xu et al., ; Ren et al., ; Feng et al., ; Mallen et al., ; Jiang et al., ). 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 to solve it. Therefore, they rely on additional modules or fine-tuning LLMs. 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., ; Akyürek et al., ; Dai et al., ). Besides, there are other explanations of ICL such as Bayes inference (Xie et al., ), Bayes model averaging (Zhang et al., ), leaning topic structure (Li et al., ) and kernel regression (Han et al., ). They focus on explaining why ICL occurs. Our contribution lies in analyzing the benefit and detriment in RAG and proposing a practical method based on our theory.

## 6 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, we propose a practical novel 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. Limitations and Societal Impact: The main limitation of this paper is that due to immense resource cost, we do not evaluate our method on LLMs with 33B and 65B scales. Our paper deepens society’s understanding of LLMs’ usage of external retrieved knowledge through theoretical analysis. After careful consideration, we believe that our paper does not have any potential negative societal impact.

## Appendix A Proof for Equation 5

###### Proof.

The transformation is motivated by (Xie et al., ) and we apply it to the analysis of RAG:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i-1})\,dz$ (20) $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\frac{p(R,x_{1:i-1}|z)p( z)}{p(R,x_{1:i-1})}\,dz$ (21) $\displaystyle\propto\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(R,x_{1:i-1}|z)p( z)\,dz,\textrm{\quad$p(R,x_{1:i-1})$ is a constant so we drop it}$ (22) $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\frac{p(R,x_{1:i-1}|z)}{ p(R,x_{1:i-1}|z^{*})}p(z)\,dz,\textrm{\quad$\frac{1}{p(R,x_{1:i-1}|z^{*})}$ is a constant so we add it}$ (23) $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\textrm{exp}(v(z))p(z)\, dz,\quad v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ (24) $$

∎

###### Proof.

## Appendix B Proof for Equation 6

###### Proof.

For $p(R,x_{1:i-1}|z)$ in $v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$, we can make further derivation as:

$$ $\displaystyle p(R,x_{1:i-1}|z)=p(x_{1:i-1}|R,z)p(R|z)$ (25) $$

According to the definition of latent variable model in the analysis of in-context learning from (Xie et al., ) that views the latent variable inference as Hidden Markov Model (HMM) and the latent concept $z$ determines the transition probability matrix in HMM hidden states $h$, we can get the following derivations :

$$ $\displaystyle p(x_{1:i-1}|R,z)p(R|z)$ $\displaystyle=\sum_{h}p(x_{1:i-1}|h,z)p(h|R,z)p(R|z),$ (26) $\displaystyle v(z)$ $\displaystyle=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ (27) $\displaystyle=\textrm{log}\frac{\sum_{h}p(x_{1:i-1}|h,z)p(h|R,z)}{\sum_{h}p(x_ {1:i-1}|h,z^{*})p(h|R,z^{*})}+\textrm{log}\frac{p(R|z)}{p(R|z^{*})}.$ (28) $$

Based on previous work (Xie et al., ; Zhang et al., ), 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(x|h,z^{*})>c_{1}>0$.

###### Assumption 2 .

Delimiter is an important distinguishing signal between each passage $r$ in the retrieved texts $R$. For any delimiter hidden state $h^{d}$ and other hidden state $h$, there are upper and lower bounds on the transition probability from $h$ to $h^{d}$: $0\leq c_{2}\leq p(h^{d}|h,z)\leq c_{3}$.

Then we can get:

$$ $\displaystyle v(z)$ $\displaystyle\leq\textrm{log}\frac{\sum_{h}1\cdot p(h|R,z)}{\sum_{h}c_{1}\cdot p (h|R,z^{*})}+\textrm{log}\frac{p(R|z)}{p(R|z^{*})}$ (29) $\displaystyle=\textrm{log}\frac{\sum_{h}1\cdot p(h|R,z)}{\sum_{h}c_{1}\cdot p( h|R,z^{*})}+\textrm{log}\frac{p(R|z)}{p(R|z^{*})}$ (30) $\displaystyle=-\textrm{log}{c_{1}}+\textrm{log}\frac{p(R|z)}{p(R|z^{*})}$ (31) $\displaystyle=-\textrm{log}{c_{1}}+\textrm{log}\frac{\prod_{i=1}^{n}p(r_{i}|r_ {1:i-1},z)}{\prod_{i=1}^{n}p(r_{i}|r_{1:i-1},z^{*})}.$ (32) $$

So we can get:

$$ $\displaystyle p(R,x_{1:i-1}|z)=p(x_{1:i-1}|R,z)p(R|z)\approx\prod_{i=1}^{n}O(1 )p(r_{i}|r_{1:i-1},z)$ (33) $\displaystyle\prod_{i=1}^{n}O(1)p(r_{i}|r_{1:i-1},z)=\prod_{i=1}^{n}\sum_{h_{i -1}^{d}\in\mathcal{D}}p(r_{i}|h_{i-1}^{d},z)p(h_{i-1}^{d}|r_{1:i-1},z),$ (34) $$

$r_{i}$ is a passage in the retrieved texts list $R$, $h_{i-1}^{d}$ is the hidden state for the delimiter between $r_{i-1}$ and $r_{i}$ in $R$. According to the Assumption 2, $p(h_{i-1}^{d}|r_{1:i-1},z)=O(1)$, then Equation 34 is approximately equal to $\prod_{i=1}^{n}O(1)p(r_{i}|z)$, which means that $p(R,x_{1:i-1}|z)\approx\prod_{i=1}^{n}O(1)p(r_{i}|z)$, so we can get that:

$$ $\displaystyle v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ $\displaystyle\approx\textrm{log}\frac{\prod_{i=1}^{n}O(1)p(r_{i}|z)}{\prod_{i= 1}^{n}O(1)p(r_{i}|z^{*})}$ (35) $\displaystyle\rightarrow O(1)+n*\frac{1}{n}\sum_{i=1}^{n}log\frac{p(r_{i}|z)}{ p(r_{i}|z*)}=O(1)+n*\mathbb{E}_{r\sim P_{r}}\left[log\frac{p(r_{i}|z)}{p(r_{i} |z^{*})}\right]$ (36) $\displaystyle\propto p_{R}(r)\textrm{log}\frac{p(r|z)}{p(r|z^{*})}=p_{R}(r) \textrm{log}\frac{p_{R}(r)}{p(r|z^{*})}-p_{R}(r)\textrm{log}\frac{p_{R}(r)}{p( r|z)}$ (37) $\displaystyle=-(\underbrace{\text{KL}(p_{R}(r)\|p(r|z))}_{\textbf{benefit}, \textrm{denote as }\Omega}-\underbrace{\text{KL}(p_{R}(r)\|p(r|z^{*}))}_{ \textbf{detriment},\textrm{denote as }\Upsilon}),$ (38) $$

$p_{R}(\cdot)$ is the distribution of the retrieved texts, $p(\cdot)$ is the distribution of the LLMs’ pre-trianed knowledge.

∎

###### Proof.

###### Assumption 1 .

###### Assumption 2 .

## Appendix C Effect of v ⁢ ( z ) 𝑣 𝑧 v(z) italic_v ( italic_z ) in Distribution Fusion

Recapping the Equation 38, we find $v(z)$ actually regulates the proportion between LLMs’ pre-trained knowledge and retrieved knowledge in distribution fusion of RAG prediction:

- •
The more benefit outweigh detriment, $v(z)\to-\infty$ and exp($v(z)$) $\to 0$ for all $z\neq z^{*}$, this indicates that concepts $z$ sampled from LLMs’ space contribute little to $p(x_{i}|R,x_{1:i-1})$. When $z=z^{*}$, exp($r(z^{*})$) $=1$, which means that latent variable model concentrates more on $z^{*}$ 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, $v(z)\to+\infty$ and exp($v(z)$) $\to+\infty$ for all $z\neq z^{*}$ and when $z=z^{*}$, exp($r(z^{*})$) $=1$. This indicates that concepts $z$ sampled from LLMs’ space contribute more and more than $z^{*}$ sampled from retrieved texts as $v(z)$ increases.

## Appendix D Proof for Theorem 1

###### Proof.

Recapping the Equation 2 that describes the distribution fusion in RAG via latent variable model:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\underbrace{\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z) p(z|R,x_{1:i-1})\,dz}_{\textrm{denote as }\Phi}+\underbrace{p(x_{i}|R,x_{1:i-1 },z^{*})p(z^{*}|R,x_{1:i-1})}_{\textrm{denote as }\Lambda}.$ (39) $$

Since latent concept $z^{*}$ determines the hidden states $h$, $\Lambda$ can be transformed as:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1},z^{*})p(z^{*}|R,x_{1:i-1})$ $\displaystyle=\sum_{h}p(x_{i}|x_{1:i-1},h,z^{*})p(h|R,x_{1:i-1}.z^{*})p(z^{*}| R,x_{1:i-1}).$ (40) $$

Let $p(z^{*}|R,x_{1:i-1})$ = $\beta$:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\Phi+\beta\sum_{h}p(x_{i}|x_{1:i-1},h,z^{*})p(h|R,x_{1:i-1}.z^{*})$ (41) $\displaystyle p_{R}(x_{i}|x_{1:i-1})$ $\displaystyle=\sum_{h}p(x_{i}|x_{1:i-1},h,z^{*})p_{R}(h|x_{1:i-1})$ (42) $\displaystyle p_{R}(h|x_{1:i-1})$ $\displaystyle\propto p(x_{1:i-1}|h,z^{*})p_{R}(h)$ (43) $\displaystyle p(h|R,x_{1:i-1},z^{*})$ $\displaystyle\propto p(x_{1:i-1}|h,z^{*})p(h|R,z^{*})$ (44) $$

let probabilities $p(x_{i}|x_{1:i-1},h,z^{*})p(x_{1:i-1}|h,z^{*})$ in Equation 40 is represented as matrix $W\in\mathbb{R}^{|\mathcal{X}|\times|\mathcal{H}|}$ for all possible $x_{i}\in\mathcal{X}$ and $h\in\mathcal{H}$, $p(h|R,z^{*})$ in Equation 44 is matrix $B$, $p_{R}(h)$ in Equation 43 is $u\in\mathbb{R}^{|\mathcal{H}|}$. We use 1-norm to calculate the difference between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$, which can be formalized as:

$$ $\displaystyle\|p(x_{i}|R,x_{1:i-1})-p_{R}(x_{i}|x_{1:i-1})\|_{1}=\|\Phi+\beta WB -Wu\|_{1}.$ (45) $$

Then, according to the triangle inequality of 1-norm, the difference between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$ is bouned by:

$$ $\displaystyle\|\Phi\|_{1}-\|\beta WB-Wu\|_{1}\leq\|\Phi+\beta WB-Wu\|_{1}\leq \|\Phi\|_{1}+\|\beta WB-Wu\|_{1}.$ (46) $$

We consider to further analyze $\|\beta WB-Wu\|_{1}$:

$$ $\displaystyle\|\beta WB-Wu\|_{1}$ $\displaystyle=\sum_{i=1}^{|\mathcal{X}|}|W_{i}^{T}(\beta B-u)|_{i}$ (47) $\displaystyle=\sum_{i=1}^{|\mathcal{X}|}|\sum_{j=1}^{|\mathcal{H}|}W_{ij}( \beta B-u)_{j}|$ (48) $\displaystyle\leq\sum_{i=1}^{|\mathcal{X}|}\sum_{j=1}^{|\mathcal{H}|}W_{ij}|( \beta B-u)_{j}|$ (49) $\displaystyle=\sum_{j=1}^{|\mathcal{H}|}(\sum_{i=1}^{|\mathcal{X}|}W_{ij})|( \beta B-u)_{j}|$ (50) $\displaystyle=\sum_{j=1}^{|\mathcal{H}|}|(\beta B-u)_{j}|$ (51) $\displaystyle=\|\beta B-u\|_{1}$ (52) $$

Then:

$$ $\displaystyle\|\beta B-u\|_{1}$ $\displaystyle=2TV(p_{R}(\cdot),\beta p(\cdot|R,z^{*}))\text{\quad TV is Total Variation Distance.}$ (53) $\displaystyle\leq 2\beta TV(p_{R}(\cdot),p(\cdot|R,z^{*}))$ (54) $\displaystyle\leq\sqrt{2\text{KL}(p_{R}(\cdot)\|p(\cdot|R,z^{*}))}\text{\quad Pinsker 's Inequality.}$ (55) $\displaystyle\leq\sqrt{2\text{KL}(p_{R}(\cdot)\|p(\cdot|z^{*}))}$ (56) $\displaystyle\approx\sqrt{2\text{KL}(p_{R}(r)\|p(r|z^{*}))},$ (57) $$

in which $r$ is the passage in $R$, $\text{KL}(p_{R}(r)\|p(r|z^{*}))$ is actually the detrimentin Equation 9.
Recapping Equation 46, we can get:

$$ $\displaystyle\|\Phi+\beta WB-Wu\|_{1}\leq\|\Phi\|_{1}+\sqrt{2\text{KL}(p_{R}(r )\|p(r|z^{*}))}$ (58) $$

Since $0\leq\|\beta WB-Wu\|_{1}\leq 2\sqrt{2\text{KL}(p_{R}(\cdot)\|p(\cdot|z^{*}))}$ and $\|\Phi+\beta WB-Wu\|_{1}\geq\|\Phi\|_{1}-\|\beta WB-Wu\|_{1}$, then the lower bound for $\|\Phi+\beta WB-Wu\|_{1}$ is included in:

$$ $\displaystyle\left[\|\Phi\|_{1}-\sqrt{2\text{KL}(p_{R}(\cdot)\|p(\cdot|z^{*})) },\|\Phi\|_{1}\right],$ (59) $$

we take the minimum value as the lower bound. Define $\mathcal{D}=\|p(x_{i}|R,x_{1:i-1})-p_{R}(x_{i}|x_{1:i-1})\|_{1}$ is the difference between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$, according to Equation 45 and 46, the lower and upper bound for $\mathcal{D}$ is:

$$ $\displaystyle\|\Phi\|_{1}-\sqrt{2\underbrace{\text{KL}(p_{R}(r)\|p(r|z^{*}))}_ {\textbf{detriment}}}\leq\mathcal{D}\leq\|\Phi\|_{1}+\sqrt{2\underbrace{\text{ KL}(p_{R}(r)\|p(r|z^{*}))}_{\textbf{detriment}}},$ (60) $$

For ease of description, we denote benefit $\text{KL}(p_{R}(r)\|p(r|z)$ as $\Omega$ and denote detriment $\text{KL}(p_{R}(r)\|p(r|z^{*})$ as $\Upsilon$. Recapping Equation 5 and 9:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\underbrace{\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z) p(z|R,x_{1:i-1})\,dz}_{\textrm{denote as }\Phi}+\underbrace{p(x_{i}|R,x_{1:i-1 },z^{*})p(z^{*}|R,x_{1:i-1})}_{\textrm{denote as }\Lambda}.$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i-1})\,dz$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\frac{p(R,x_{1:i-1}|z)p( z)}{p(R,x_{1:i-1})}\,dz$ $\displaystyle\propto\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(R,x_{1:i-1}|z)p( z)\,dz,\textrm{\quad$p(R,x_{1:i-1})$ is a constant so we drop it}$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\frac{p(R,x_{1:i-1}|z)}{ p(R,x_{1:i-1}|z^{*})}p(z)\,dz,\textrm{\quad$\frac{1}{p(R,x_{1:i-1}|z^{*})}$ is a constant so we add it}$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\textrm{exp}(v(z))p(z)\, dz,\quad v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ $$

$$ $\displaystyle v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ $\displaystyle\approx-\left[\underbrace{\text{KL}(p_{R}(r)\|p(r|z))}_{\textbf{ benefit},\textrm{denote as }\Omega}-\underbrace{\text{KL}(p_{R}(r)\|p(r|z^{*}) )}_{\textbf{detriment},\textrm{denote as }\Upsilon}\right]$ $$

$\Phi$ in Equation 60 can be transformed as:

$$ $\displaystyle\Phi$ $\displaystyle=\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i -1})\,dz$ (61) $\displaystyle=\frac{p(R,x_{1:i-1}|z^{*})}{p(R,x_{1:i-1})}\int_{\mathcal{Z}-\{{ z}^{*}\}}p(x_{i}|R,x_{1:i-1},z)\textrm{exp}(v(z))p(z)\,dz$ (62) $\displaystyle=\alpha\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z) \textrm{exp}(v(z))p(z)\,dz,\quad(\text{$p(R,x_{1:i-1}|z^{*})$ and ${p(R,x_{1:i -1})}$ are constants})$ (63) $\displaystyle\approx\alpha\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z) \textrm{exp}(-(\Omega-\Upsilon))p(z)\,dz\quad\textrm{ (Equation~{}\ref{eq4})}.$ (64) $$

Now the Theorem 1 has been proven.

∎

###### Proof.

## Appendix 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 ($\frac{1}{\mathcal{D}}$) between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$. To achieve this, we can start from Equation 60 to prove that
the gap between values of benefit and detriment is negatively correlated with
the difference ($\mathcal{D}$) between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$, which is actually the reciprocal of similarity ($\frac{1}{\mathcal{D}}$). Specifically, we want to prove that the gap between values of benefit and detriment ($\text{KL}(p_{R}(r)\|p(r|z)-\text{KL}(p_{R}(r)\|p(r|z^{*})$) is negatively correlated with both lower and upper bound in Equation 60. For ease of description, we denote benefit $\text{KL}(p_{R}(r)\|p(r|z)$ as $\Omega$ and denote detriment $\text{KL}(p_{R}(r)\|p(r|z^{*})$ as $\Upsilon$.

###### Proof.

Recapping Equation 5 and 9:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\underbrace{\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z) p(z|R,x_{1:i-1})\,dz}_{\textrm{denote as }\Phi}+\underbrace{p(x_{i}|R,x_{1:i-1 },z^{*})p(z^{*}|R,x_{1:i-1})}_{\textrm{denote as }\Lambda}.$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i-1})\,dz$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\frac{p(R,x_{1:i-1}|z)p( z)}{p(R,x_{1:i-1})}\,dz$ $\displaystyle\propto\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(R,x_{1:i-1}|z)p( z)\,dz,\textrm{\quad$p(R,x_{1:i-1})$ is a constant so we drop it}$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\frac{p(R,x_{1:i-1}|z)}{ p(R,x_{1:i-1}|z^{*})}p(z)\,dz,\textrm{\quad$\frac{1}{p(R,x_{1:i-1}|z^{*})}$ is a constant so we add it}$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\textrm{exp}(v(z))p(z)\, dz,\quad v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ $$

$$ $\displaystyle v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ $\displaystyle\approx-\left[\underbrace{\text{KL}(p_{R}(r)\|p(r|z))}_{\textbf{ benefit},\textrm{denote as }\Omega}-\underbrace{\text{KL}(p_{R}(r)\|p(r|z^{*}) )}_{\textbf{detriment},\textrm{denote as }\Upsilon}\right]$ $$

$\Phi$ in Equation 60 can be transformed as:

$$ $\displaystyle\Phi$ $\displaystyle=\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i -1})\,dz$ (65) $\displaystyle=\frac{p(R,x_{1:i-1}|z^{*})}{p(R,x_{1:i-1})}\int_{\mathcal{Z}-\{{ z}^{*}\}}p(x_{i}|R,x_{1:i-1},z)\textrm{exp}(v(z))p(z)\,dz$ (66) $\displaystyle=\alpha\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z) \textrm{exp}(v(z))p(z)\,dz,\quad(\text{$p(R,x_{1:i-1}|z^{*})$ and ${p(R,x_{1:i -1})}$ are constants})$ (67) $\displaystyle\approx\alpha\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z) \textrm{exp}(-(\Omega-\Upsilon))p(z)\,dz\quad\textrm{ (Equation~{}\ref{eq4})}.$ (68) $$

Therefore, the lower bound of Equation 60 is:

$$ $\displaystyle\|\Phi\|_{1}-\sqrt{2\Upsilon}$ $\displaystyle\approx\alpha\|\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1}, z)\textrm{exp}(-(\Omega-\Upsilon))p(z)\,dz\|_{1}-\sqrt{2\Upsilon}$ (69) $\displaystyle\propto\text{exp}(-(\Omega-\Upsilon))-\sqrt{2\Upsilon}$ (70) $$

and the upper bound of Equation 60 is:

$$ $\displaystyle\|\Phi\|_{1}+\sqrt{2\Upsilon}$ $\displaystyle\propto\text{exp}(-(\Omega-\Upsilon))+\sqrt{2\Upsilon}$ (71) $$

Due to both $\Omega$ and $\Upsilon$ being variables, analyzing the result of subtraction between $\Omega$ and $\Upsilon$ under their simultaneous changes is complex. Therefore, we use the “Separation of variables“ to simplify our analysis. Specifically, we first assume that one is constant, and then analyze the changes caused by the variation of another:

- •
Assume $\Omega$ is constant, as the value of $\Omega-\Upsilon$ increases, $\Upsilon$ decreases and the upper bound $\text{exp}(-(\Omega-\Upsilon))+\sqrt{2\Upsilon}$ also deceases. In the lower bound $\text{exp}(-(\Omega-\Upsilon))-\sqrt{2\Upsilon}$, since the first term is an exponential function and the second term is a square root function, a decrease of $\Upsilon$ leads to the decrease in the entire lower bound. Therefore, both lower and upper bounds in Equation 60 decrease as $\Omega-\Upsilon$ increases.
- •
Assume $\Upsilon$ is constant, as the value of $\Omega-\Upsilon$ increases, $\Omega$ increases and the upper bound $\text{exp}(-(\Omega-\Upsilon))+\sqrt{2\Upsilon}$ deceases. In the lower bound $\text{exp}(-(\Omega-\Upsilon))-\sqrt{2\Upsilon}$, since the first term is an exponential function and the second term is a square root function, an increase of $\Omega$ leads to the decrease in the entire lower bound. Therefore, both lower and upper bounds in Equation 60 decrease as $\Omega-\Upsilon$ increases.

On behalf of the analysis above, we can derve that both lower and upper bounds in Equation 60 are approximately negatively correlated with the gap between values of benefit and detriment. Therefore, the difference $\mathcal{D}$ between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$ is approximately negatively correlated with the gap between values of benefit and detriment. In other words, $\frac{1}{D}$ can be treated as the similarity between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$ and it is approximately positively correlated with the gap between values of benefit and detriment.:

$$ $\underbrace{\text{KL}(p_{R}(r)\|p(r|z))}_{\textbf{benefit}}-\underbrace{\text{ KL}(p_{R}(r)\|p(r|z^{*}))}_{\textbf{detriment}}\propto\frac{1}{\mathcal{D}}.$ (72) $$

So we have proved that the gap between values of benefit and detriment is approximately positively correlated with $\frac{1}{D}$.

∎

###### Proof.

## Appendix F Proof for Theorem 3

This section aims to prove:

$$ $\displaystyle J=\begin{cases}\textnormal{KL}(p_{R}(r)\|p(r|z))<\textnormal{KL} (p_{R}(r)\|p(r|z^{*}))\text{, detriment outweighs benefit.}&\text{if }\frac{1} {\mathcal{D}}<\frac{1}{\mathcal{M}}\\ \textnormal{KL}(p_{R}(r)\|p(r|z))=\textnormal{KL}(p_{R}(r)\|p(r|z^{*}))\text{, detriment is equal to benefit.}&\text{if }\frac{1}{\mathcal{D}}=\frac{1}{ \mathcal{M}}\\ \textnormal{KL}(p_{R}(r)\|p(r|z))>\textnormal{KL}(p_{R}(r)\|p(r|z^{*}))\text{, benefit outweighs detriment.}&\text{if }\frac{1}{\mathcal{D}}>\frac{1}{ \mathcal{M}}\end{cases}$ (73) $$

in which $\frac{1}{\mathcal{M}}$ is the similarity between $p(x_{i}|R,x_{1:i-1})$ and $p(x_{i}|x_{1:i-1})$ (LLMs’ pre-trained knowledge), $\frac{1}{\mathcal{D}}$ is the similarity between $p(x_{i}|R,x_{1:i-1})$ and $p_{R}(x_{i}|x_{1:i-1})$ (distribution of retrieved texts)

###### Proof.

When benefit is equal to detriment:

$$ $\displaystyle\text{KL}(p_{R}(r)\|p(r|z))-\text{KL}(p_{R}(r)\|p(r|z^{*}))=0,$ (74) $$

which means that:

$$ $\displaystyle p_{R}(r)\textrm{log}\frac{p(r|z)}{p(r|z^{*})}=0,$ (75) $$

since $p_{R}(r)$ cannot be 0 0, then:

$$ $\displaystyle\textrm{log}\frac{p(r|z)}{p(r|z*)}$ $\displaystyle=0,$ (76) $\displaystyle\frac{p(r|z)}{p(r|z*)}$ $\displaystyle=1,$ (77) $\displaystyle p(r|z)$ $\displaystyle=p(r|z^{*}),$ (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:

$$ $\displaystyle p(x_{i}|x_{1:i-1})=p_{R}(x_{i}|x_{1:i-1}),$ (79) $$

then:

$$ $\displaystyle\|p(x_{i}|R,x_{1:i-1})-p(x_{i}\|x_{1:i-1})\|_{1}=\|p(x_{i}|R,x_{1 :i-1})-p_{R}(x_{i}\|x_{1:i-1})\|_{1},$ (80) $$

which means that $\frac{1}{\mathcal{D}}=\frac{1}{\mathcal{M}}$ is an important dividing point. When $\frac{1}{\mathcal{D}}=\frac{1}{\mathcal{M}}$, we can get that benefit is equal to detriment and $\text{KL}(p_{R}(r)\|p(r|z))-\text{KL}(p_{R}(r)\|p(r|z^{*}))=0$. Equation 72 indicates that the gap between values of benefit and detriment ($\text{KL}(p_{R}(r)\|p(r|z))-\text{KL}(p_{R}(r)\|p(r|z^{*}))$) is approximately positively correlated with $\frac{1}{\mathcal{D}}$. Therefore, when $\frac{1}{\mathcal{D}}>\frac{1}{\mathcal{M}}$ we can get that benefit outweighs detriment ($\text{KL}(p_{R}(r)\|p(r|z))-\text{KL}(p_{R}(r)\|p(r|z^{*}))>0$). When $\frac{1}{\mathcal{D}}<\frac{1}{\mathcal{M}}$ we can get that detriment outweighs benefit ($\text{KL}(p_{R}(r)\|p(r|z))-\text{KL}(p_{R}(r)\|p(r|z^{*}))<0$). Now the proof of Theorem 3 has been finished.
∎

###### Proof.

## Appendix G Proof for RAG is actually unsupervised In-context Learning

This section aims to prove that RAG is actually unsupervised ICL from two perspectives. One is that previous studies find that ICL performs gradient descent as meta-optimizer (Von Oswald et al., ; Akyürek et al., ; Dai et al., ). 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 in 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 as the following:

Gradient descent in optimization of linear layers have a dual form of linear attention (Irie et al., ; Aizerman et al., ), define a liner layer as:

$$ $\displaystyle f(x)=W_{0}x,$ (81) $$

in which $W_{0}$ is the initial weight matrix. Given a sequence of historical input vectors $\mathbf{x_{i}}\in\mathbb{R}^{d_{in}}$ and corresponding error signals $\mathbf{e_{i}}\in\mathbb{R}^{d_{out}}$, $i\in[1,N]$ obtained by gradient descent, the update of the weight matrix can be represented as:

$$ $\displaystyle W^{\prime}=W_{0}+\Delta W=W_{0}+\sum_{i}^{N}\mathbf{e_{i}} \otimes\mathbf{x_{i}}.$ (82) $$

Recap that the linear attention can be formulated as:

$$ $\displaystyle\text{LinearAttn}(V,K,\mathbf{q})=\sum_{i}\mathbf{v}_{i}(\mathbf{ k}^{T}_{i}\mathbf{q}).$ (83) $$

Then the dual form of updated linear layer with new input $\mathbf{x}_{N+1}$ is:

$$ $\displaystyle f^{\prime}(x)$ $\displaystyle=(W_{0}+\Delta W)\mathbf{x}_{N+1}$ (84) $\displaystyle=(W_{0}+\sum_{i}^{N}\mathbf{e}_{i}\otimes\mathbf{x}_{i})\mathbf{x }_{N+1}$ (85) $\displaystyle=W_{0}\mathbf{x}_{N+1}+\sum_{i}^{N}(\mathbf{e}_{i}\otimes\mathbf{ x}_{i})\mathbf{x}_{N+1}$ (86) $\displaystyle=W_{0}\mathbf{x}_{N+1}+\sum_{i}^{N}\mathbf{e}_{i}\otimes(\mathbf{ x}_{i}^{T}\mathbf{x}_{N+1})$ (87) $\displaystyle=W_{0}\mathbf{x}_{N+1}+\text{LinearAttn}(E,\mathbf{x}_{1:N}, \mathbf{x}_{N+1})$ (88) $$

In in-context learning, the attention of a head is:

$$ $\displaystyle f_{ICL}(\mathbf{q})$ $\displaystyle=\text{Attn}(V,K,\mathbf{q})$ (89) $\displaystyle=W_{V}[B^{\prime}:B]\text{softmax}\left(\frac{W_{K}[B^{\prime}:B] ^{T}\mathbf{q}}{\sqrt{d}}\right),$ (90) $$

in which $\mathbf{q}=W_{Q}\mathbf{b}$, $\mathbf{b}$ is the input $t$-th token in query, $W_{Q}$, $W_{K}$, $W_{v}$ are projection matrices, $B^{\prime}$ is demonstrations of the context and $X$ is the prefix for $\mathbf{b}$ of the query. To simplify qualitative analysis, we follow (Dai et al., ) to estimate the standard attention as relaxed linear attention, achieved by eliminating the softmax function and the scaling factor:

$$ $\displaystyle f_{ICL}(\mathbf{q})$ $\displaystyle\approx W_{V}[B^{\prime}:B](W_{K}[B^{\prime}:B])^{T}\mathbf{q}$ (91) $\displaystyle=W_{V}B(W_{K}B)^{T}\mathbf{q}+W_{V}X^{\prime}(W_{K}B^{\prime})^{T }\mathbf{q}$ (92) $\displaystyle=W_{V}B(W_{K}B)^{T}\mathbf{q}+\text{LinearAttn}(W_{V}B^{\prime},W _{K}B^{\prime},\mathbf{q})$ (93) $$

According to Equation 88, the dual form of the Transformer attention is:

$$ $\displaystyle f_{ICL}(\mathbf{q})$ $\displaystyle\approx W_{V}B(W_{K}X)^{T}\mathbf{q}+\text{LinearAttn}(W_{V}B^{ \prime},W_{K}B^{\prime},\mathbf{q})$ (94) $\displaystyle=W_{V}B(W_{K}B)^{T}\mathbf{q}+\sum_{i}W_{V}\mathbf{b}^{\prime}_{i }\left((W_{K}\mathbf{b}^{\prime}_{i})^{T}\mathbf{q}\right)$ (95) $\displaystyle=W_{V}B(W_{K}B)^{T}\mathbf{q}+\sum_{i}\left((W_{V}\mathbf{b}^{ \prime}_{i})\otimes(W_{K}\mathbf{b}^{\prime}_{i})\right)\mathbf{q}.$ (96) $$

Based on above derivation, we have this finding: comparing Equation 86 with Equation 96, we find that $W_{V}B(W_{K}B)^{T}$ is equal to the initial weight matrix $W_{0}$, which is zero-shot prediction give query prefix $B$ without demonstrations in the context. Besides, $W_{V}\mathbf{b}_{i}^{\prime}$ is equal to $\mathbf{e}_{i}$. which is the meta-gradient used to update the weighted matrix. $W_{K}\mathbf{b}_{i}$ is equal to the historical input vectors:

$$ $\displaystyle W_{V}\mathbf{b}_{i}^{\prime}$ $\displaystyle=\mathbf{e}_{i}$ (97) $\displaystyle W_{K}\mathbf{b}_{i}^{\prime}$ $\displaystyle=x_{i}.$ (98) $$

In the standard gradient descent with loss $\mathcal{L}$, $e_{i}=-\eta\frac{\partial\mathcal{L}}{\partial y_{i}}$ and $\eta$ is the learning rate and $y_{i}=W_{i}x_{i}$ is the output of the linear layer using the weight matrix $W_{i}$ at step $t$ (Irie et al., ). So we can get:

$$ $\displaystyle\mathbf{e}_{i}=-\eta\frac{\partial\mathcal{L}}{\partial y_{i}}$ $\displaystyle=-\eta\frac{\partial\mathcal{L}}{\partial W_{i}x_{i}}$ (99) $\displaystyle=-\eta\frac{\partial\mathcal{L}}{\partial W_{i}W_{K}\mathbf{b}_{i }^{\prime}}.$ (100) $$

Therefore:

$$ $\displaystyle-\eta\frac{\partial\mathcal{L}}{\partial W_{i}W_{K}\mathbf{b}^{ \prime}_{i}}=W_{V}\mathbf{b}_{i}^{\prime}.$ (101) $$

So the loss $\mathcal{L}$ ca be represented as:

$$ $\displaystyle\mathcal{L}=-\frac{1}{\eta}\int W_{V}\mathbf{b}_{i}^{\prime}\,d(W _{i}W_{K}\mathbf{b}_{i}^{\prime})$ (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_{i}$) in the demonstration of context. This suggests that the distribution of the text in context is a direct learning signal for in-context learning, without the need for explicit input-output pairs in the demonstration. From the perspective that ICL 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:

$$ $\displaystyle p(x_{i}|R,x_{1:i-1})$ $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i-1})\,dz$ (103) $\displaystyle=\int_{\mathcal{Z}-\{{z}^{*}\}}p(x_{i}|R,x_{1:i-1},z)p(z|R,x_{1:i -1})\,dz+p(x_{i}|R,x_{1:i-1},z^{*})p(z^{*}|R,x_{1:i-1}).$ (104) $\displaystyle\propto\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)p(R,x_{1:i-1}|z)p( z)\,dz$ (105) $\displaystyle=\int_{\mathcal{Z}}p(x_{i}|R,x_{1:i-1},z)\textrm{exp}(v(z))p(z)\, dz,\quad v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ (106) $$

$$ $\displaystyle v(z)=\textrm{log}\frac{p(R,x_{1:i-1}|z)}{p(R,x_{1:i-1}|z^{*})}$ $\displaystyle\approx\textrm{log}\frac{\prod_{i=1}^{n}O(1)p(r_{i}|z)}{\prod_{i= 1}^{n}O(1)p(r_{i}|z^{*})}$ (107) $\displaystyle\rightarrow O(1)+n*\frac{1}{n}\sum_{i=1}^{n}log\frac{p(r_{i}|z)}{ p(r_{i}|z*)}=O(1)+n*\mathbb{E}_{r\sim P_{r}}\left[log\frac{p(r_{i}|z)}{p(r_{i} |z^{*})}\right]$ (108) $\displaystyle\propto p_{R}(r)\textrm{log}\frac{p(r|z)}{p(r|z^{*})}=p_{R}(r) \textrm{log}\frac{p_{R}(r)}{p(r|z^{*})}-p_{R}(r)\textrm{log}\frac{p_{R}(r)}{p( r|z)}$ (109) $\displaystyle=-(\text{KL}(p_{R}(r)\|p(r|z))-\text{KL}(p_{R}(r)\|p(r|z^{*}))),$ (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. $p_{R}(r)$ is the distribution of retrieved texts and it can serve as ground truth distribution in loss functions. $p(r|z)$ and $p(r|z^{*})$ are distribution estimated by LLMs. Two KL-divergence between ground truth distribution $p_{R}(r)$ and estimated distribution $p(r|z)$ and $p(r|z^{*})$ respectively are loss functions. $r$ 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 $z^{*}$ from the retrieved texts. So $\text{KL}(p_{R}(r)\|p(r|z^{*})$ is the actual loss that can be meta-optimized in RAG, in which $p_{R}(r)$ is the ground truth distribution and $p(r|z^{*})$ is the estimated distribution. As this loss decreases, the value of $v(z)$ when $z$ is not equal to $z^{*}$ also decreases, which means that in Equation 104, 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 $p_{R}(r)$ is the ground truth in this learning but dose not have any explicit input-output supervision like demonstrations in in-context learning, RAG is the unsupervised in-context learning and distribution of retrieved texts ($p_{R}(r)$) is the unsupervised learning signal.

Based on the above two perspectives, we successfully prove that: The distribution of retrieved passage $r$ in RAG (i.e., $p_{R}(r)$) can serve as 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 LLMs’ pre-trained distribution.
∎

###### Proof.

## Appendix 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 4 and 5, and the conclusions of them are consistent with Section 3.1.

**Table 4: 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.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 Mistral-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 |

## Appendix I Experimental Details

### I.1 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 to this task.

Logprobs. Logprobs can indicate the confidence for LLMs in generating the tokens(Kuhn et al., ). We use the value order between top-1 log-probability 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 is greater than the logprobs of tokens generated by pure LLM, the benefit outweigh the detriment, otherwise the detriment outweigh the benefit.

Uncertainty. We use Length-normalized Entropy (Malinin & Gales, ) to measure the uncertainty of the tokens generated by pure LLM and RAG respectively. If the uncertainty of tokens generated by RAG is lower than the uncertainty of tokens generated by pure LLM, the benefit outweigh the detriment, otherwise the detriment outweigh the benefit.

Consistency-Lexical (Lin et al., ). Consistency-based methods make LLMs perform multiple generations for a 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-similarity 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.

Consistency-Semantic (Chen et al., ). We follow (Chen et al., ) to use EigenScore to calculate the semantic similarity among hidden states of tokens in multiple generations and use it as the consistency score.

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., )) or as action triggers (CRAG (Yan et al., )), train more robust LLMs for RAG (RetRobust (Yoran et al., ) and INFO-RAG (Xu et al., )) and train LLMs to dynamically retrieve and critique retrieved texts (Self-RAG (Asai et al., )).

NLI+RAG. This method use a Natural Language Inference model to filter the possible irrelevant documents in retrieved results and provide the remaining documents to LLMs for generation. We follow (Yoran et al., ) to use a BART-Large model (Lewis et al., ) with 407 million parameters trained on the MNLI dataset (Williams et al., ). We consider a query-document pair as entailed if the probability for the entailment label is $\geq 0.5$ and filter the documents with probability for the entailment label $<0.5$.

CRAG. 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 (Yan et al., ).

RetRobust. This method fine-tunes LLMs to properly leverage retrieved passages with a mix of relevant and irrelevant contexts (Yoran et al., ).

INFO-RAG. This method uses unsupervised method to make LLMs learn to use 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 et al., ).

Self-RAG. This method trains LLMs to dynamically retrieve and critique retrieved texts. Self-RAG first decodes a retrieval token to evaluate the utility of retrieval and control a retrieval component. If retrieval is required, LLM calls an external retrieval module to find top relevant documents, using input query and previous generation. If retrieval is not required, LLM continues generation. 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., ).

### I.2 Implementation Details

All models are run on a V100 GPU with Pytorch (Paszke et al., ) and accelerated by DeepSpeed (^1^11https://github.com/microsoft/DeepSpeed). As for retrieval for RAG, we follow (Xu et al., ; ) to use ColBERTv2 (Santhanam et al., ), an excellent generalizable model as the retriever, and use Wikipedia consisting of 21,015,324 passages (Karpukhin et al., ) as retrieval database. All baselines and Tok-RAG share the same retrieval setup and prompt. We use OPT-6.7B, LLaMA-2-7B and Mistral-7B-v0.1 as LLMs in benefit-detriment comparison experiment and use greedy-decoding strategy for generation.

## Appendix 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 3 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. Then 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. This case study visually demonstrates that our Tok-RAG effectively enables pure LLM and RAG for collaborative generation to preserve benefit and avoid detriment.

Figure: 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 3 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.
Refer to caption: x5.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
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. 

2. 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.

3. 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. 

4. 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. 

5. 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. 

6. Just do double check, is Eqn 29 the same as Eqn 30?

Overall, the paper would be a lot easier to understand and review if there weren't so many leaps of logic in the proofs.
soundness: 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.
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.
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. 

1. 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. 

2. 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. 

3. 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. 

4. 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.  

5. 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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
2. 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.
soundness: 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.
2. 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.
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.
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.
2. From Equation 4 to 5, the symbol is used incorrectly; An equal cannot be used.
3. In Equation 7, the $P_r$  is not explained. Is it $p_R(r)$?
4. 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.
5. 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.

Others:
1. In Equation 2, it is better to write the right side of the equation as $ = p(..) + \int .. dz$ to avoid misunderstandings.
2. 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^*$.
3. 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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: See the weakness section.
soundness: 3
strengths: - The paper provides a theory to understand the benefit and detriment of the retrieved documents in RAG.
- The paper introduces a method called token-RAG to leverage benefit and prevent detriment.
- The authors conduct experiments on several datasets.
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.
weaknesses: - The meaning of some notations is not well explained. What is z*?

- 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?

- Why can you go from (4) to (5)? I believe it should not be “=”. Why p(R,x1:{i-1}|z*) is a constant?

- 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.

- 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.

- 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?

Meta-review: This paper provides a theory to describe the trade-off between the value of external knowledge and its potential risk of misleading LLMs in next token prediction of RAG, and proposes a practical method to make pure LLM and RAG collaborate at token level. This paper connects theory with practice and has novel ideas. The reviewers raised major concerns about the presentation of the paper, and it has many writing issues, which need to be carefully addressed in the next version. One reviewer was still not convinced with the theoretical justification. After the author-reviewer discussion, all the reviewers were slightly positive about this paper. In all, this is a borderline paper and  I have mixed feelings about it, a bit positive.

Decision: Accept (Poster)