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

Thank you for your insightful feedback! We follow your suggestions to revise our paper and the new PDF version has been uploaded. The revised content is marked in blue font in our paper. We believe that our response and revisions address your concerns regarding the clarity of our theory and generalizability of our practical method. Following is our response, we look forward to your reply!

Response to Weakness

W1: The meaning of some notations is not well explained. What is z*?

Our Response: We apologize for the insufficient emphasis on $z^*$ in our paper. $z^*$ is the latent concept that is sampled from retrieved texts. In Line 169-170 of our latest uploaded PDF, we have explained this more clear.

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

Our Response: Thanks for your question! After our careful verification, Eqn (2) is correct. We can get $p(z|R,x_{1:{i-1}})$ out from the integral because the space of high dimensional concept is discrete. The core question for this concern is the about the property of $\mathcal{Z}$, the space of high dimensional concept. In this paper, $\mathcal{Z}$ is a discrete set rather than continuous. It is because that the theoretical framework of latent variable inference in our paper is based on Hidden Markov Model (see Line 140-144 of our paper). In Hidden Markov Model, the space of latent concept is a discrete set because it is the standard assumption in HMM [2] and both the transition probability matrix and the observation probability distribution are defined based on a set of discrete states [2].

[2] A Tutorial on Hidden Markov Models and Selected Applications in Speech Recognition

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

Our Response: We are sorry for this typo, the $=$ is actually $\propto$. In our latest uploaded PDF, we have revised this.

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

Our Response: Thanks for you propose this insightful suggestion to our method more solid. We follow your suggestions to compare our method and baselines in more various tasks (dialogue, code generation, slot filling, language modeling and long-form question answering) in RAG setting and the experimental results shown in the below Table indicate that our method can also achieve state-of-the-art performance in many other tasks. In our latest uploaded PDF, we have added this table to the main paper as Table 2 and the specific experimental settings are introduced in Appendix I.

In this experiment, as for the retrieval model and retrieval database, for Slot Filling and Language Modeling, we use ColBERTv2 , a late-interaction model with excellent generalization ability as the retriever, and use Wikipedia consisting of 21,015,324 passages as retrieval database. For Code Generation, we SCODE-R as code retriever and use deduplicated source codes in CodeSearchNET as retrieval database. For all the above tasks, we give Top-5 retrieved passages to each example. For LFQA, Dialog, we use the list of contextual passages provided in the datasets as the retrieved list (distractor setting). All baselines and our method share the same retrieved documents.

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

Our Response: Thanks for your insightful question and please allow us to explain it. Eqn (17) and (18) involves three distributions: distribution of RAG $p(x_i|R,x_{1:i-1})$, distribution of LLM $p(x_i|x_{1:i-1})$ and distribution of retrieved texts $p_R(x_i|x_{1:i-1})$. Although $p(x_i|x_{1:i-1})$ and $p(x_i|R,x_{1:i-1})$ can use the real predicted probability of the words, but $p_R(x_i|x_{1:i-1})$ cannot. So we use the heuristic method in Eqn (16) to calculate $p_R(x_i|x_{1:i-1})$ by attention score and word embedding similarity. Therefore, to ensure that the three distributions are on the same scale, we use average word embedding for $p(x_i|R,x_{1:i-1})$, $p(x_i|x_{1:i-1})$ and $p_R(x_i|x_{1:i-1})$ in Eqn (17) and (18).

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

Our Response: Our Tok-RAG method enables the LLM and RAG to collaborate at token level for generation to preserve benefit and avoid detriment. As described in Figure 1 of our paper, given the prefix $x_{1:i-1}$, RAG generates the token $x_i$ and LLM without RAG generates $x_{i}'$. We compare the value of benefit and detriment according to Eqn (19). If benefit wins, $x_i$ is selected, otherwise, $x_{i}'$ is selected. The selected token will be concatenated with $x_{1:i-1}$ as the new prefix for next step generation. This method is derived from our rigorous theoretical analysis in Section 2. In Line 396-398 of our latest uploaded PDF, we explain Equ (19) to make it more clear.

Thanks for your response. This is a insightful question! In fact, using integral to represent the latent variable inference is a standard formalization that does not vary with the properties of the high-dimensional latent variable space, even if it is discrete [1, 2, 3]. Using sum or integral does not affect the conclusion obtained from our theory.

[1] Language as a Latent Variable: Discrete Generative Models for Sentence Compression

[2] An explanation of in-context learning as implicit bayesian inference

[3] Scalable Inference in Latent Variable Models

Thank you for your insightful feedback! We follow your suggestions to revise our paper and the new PDF version has been uploaded. The revised content is marked in blue font in our paper. Following is our response, we look forward to your reply!

Response to Weakness

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

Our Response: Thanks for your feedback. We just use Equation 2 to show the phenomenon of distribution fusion in RAG in a relatively intuitive and easy-to-understand way. The more rigorous and specific description of distributed fusion is given by $v(z)$ in Equation 5. $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. For a detailed explanation, please refer to Appendix C in our paper.

W2: From Equation 4 to 5, the symbol is used incorrectly; An equal cannot be used.

Our Response: We are sorry for this typo, the $=$ is actually $\propto$. In our latest uploaded PDF, we have revised this.

W3: In Equation 7, the $P_r$ is not explained. Is it $P_R$?

Our Response: Yes, $P_r$ should be revised as $P_R$, we are sorry for this typo.

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

Our Response: Please allow us to explain this. $p_R(x_i|x_{1:i-1})$ means the next token prediction given prefix $x_{1:i-1}$ based on the distribution of retrieved texts $R$. The calculation of the $Att$ score is achieved through the interaction between the prefix $x_{1:i-1}$ and the retrieved text $R$, which measures the contribution of $R$ to the prediction of $x_i$. The parameters, such as QKV matrix, is from LLM but the objects of distribution statistics are the vocabulary and context information in the retrieved text $R$. So our paper calls this method as "approximate" the distribution. Analysis in our response to your question 1 below shows that our method is robust enough to approximate the distribution $p_R(x_i|x_{1:i-1})$, it is less affected by the knowledge distribution of LLM itself.

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

Our Response: Please allow us to clarify it. Although our Tok-RAG needs collaborative generation between pure LLM and RAG, it dose not need two models because it can be achieved by setting batch size as 2 in inference. In a batch, one input is only a query (pure LLM) and the other input is retrieved texts concatenated with the query (RAG). Toke-RAG is achieved in this setting and the GPU memory is not double the number of model parameters (2*33B or 2*65B) but only double the number of input and output tensors, which is significantly smaller than model parameters.

W6: In Equation 2, it is better to write the right side of the equation to avoid misunderstandings.

Our Response: Thanks for your suggestion, we have made the corresponding revision in our latest uploaded PDF.

W7: 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^*$.

Our Response: Thanks for your suggestion, we have made the corresponding revision in our latest uploaded PDF.

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

Our Response: Thanks for your suggestion, we have made the corresponding revision in our latest uploaded PDF.

Thanks for your response. After careful consideration, I will maintain my current scores.

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

Our Response: We appreciate your suggestions on the readability of the paper. In Line 121-154 of our latest uploaded PDF file, we follow your suggestions to move the important definitions and assumptions to the main body of the paper, ensuring that the core concepts are clearly presented. The mathematical derivations remain in the appendix to maintain the flow of the main text. We believe this change will improve the clarity and comprehensibility of the paper.

Response to Questions

Q1: 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". (Q1.1) Can the authors more rigorously define this quantity? (Q1.2) What is the probability over? i.e. is $p\_{R}(x\_i)$ the probability of $r_i$ being retrieved? if $p$ stands for probability of an event, then is it simply $p(r\_i)$? (Q1.3) Could it just be a uniform distribution then? If not, (Q1.4) 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\_{i-1})$ appears. (Q1.5) 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.

Our Response: Thanks for your valuable suggestions, please allow us to answer your questions below. We follow your suggestions to make $p\_{R}(\cdot)$ more clear in Line 127-139 of our latest uploaded PDF. (We found that OpenReview's rendering of latex formulas is unstable (depending on the browser version), so the display of the following formulas may be abnormal. We recommend you to read Line 127-139 of our latest uploaded paper.)

Q1.1: More rigorously definition of "$p\_{R}(\cdot)$ is the word distribution of retrieved passages list $R$".

A1.1: . $p\_{R}(\cdot)$ is the language modeling distribution just like the distribution $p(\cdot)$ of the LLM. This distribution $p(\cdot)$ is learned through training LLM on large corpus with next-token prediction paradigm. $p(x\_i | x\_{1:i-1})$ represents the probability distribution predicted by the LLM for $x\_i$, given the prefix $x\_{1:i-1}$. The LLM is capable of generating reasonable and coherent text based on the distribution patterns of vocabulary and context learned from the large training data. Here, $p\_{R}(\cdot)$ refers to the distribution $p(\cdot)$ learned when the training data that is limited to the retrieved passages list $R$, and it represents the vocabulary and contextual distribution patterns of the retrieved passages list $R$.

Q1.2: What is the probability $p\_{R}(\cdot)$ over?

A1.2: $r\_i$ is the $i$-$th$ passage in the retrieved passages list $R$. According to our A1.1, $p\_{R}(\cdot)$ is the language modeling distribution, so $p\_{R}(r\_i)$ represents the joint probability of the natural language sequence $r\_{i} = [r\_i^1,r\_i^2,r\_i^3,...,r\_i^l]$ under the distribution $p_{R}({\cdot})$, in which $r^x\_i$ is the $x$-$th$ token in passage $r\_i$. Specifically, according to the chain rule, $p\_{R}(r\_i)$ can be decomposed into the product of a series of conditional probabilities:

Q1.3: Could it just be a uniform distribution then? Q1.4: is there a "retriever" conditioned on a "query" that determines this distribution?

A1.3, A1.4: According to our A1.1 and A1.2, $p_{R}(\cdot)$ is the language modeling distribution and $p\_{R}(r\_i)$ represents the joint probability of the natural language sequence $r\_i$. This distribution is determined by the retrieved texts list $R$, not the retriever.

Q1.5: What does $p\_{R}(x\_i|x\_{1:i-1})$ in Equation 42 mean?

A1.5: According to our A1.1 and A1.2, $p\_{R}(\cdot)$ is the language modeling distribution, so $p\_{R}(x\_i|x\_{1:i-1})$ means the next token prediction for $x\_i$ given prefix $x\_{1:i-1}$.

Q2: 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^*)$, 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.

Our Response: There is some misunderstanding here. Please allow us to clarify it, the following part has also been revised in Line 738-783 of our latest uploaded PDF to make this point clear. (We found that OpenReview's rendering of latex formulas is unstable (depending on the browser version), so the display of the following formulas may be abnormal. We recommend you to read Line 738-783 of our latest uploaded paper.)

In Eqn 28, we get:

$

p(x\_{1:i-1}|h,z^*) > c\_1.

\notag $$ And the numeratorp(x_{1:i-1}|h,z)$in the numerator of the first term in$v(z)$is always lower than$1$:

$

v(z) &= \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^*)} \notag\\\\
& \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^*)} \text{  , // greater numerator and lower denominator}

\notag $$ In this way, we usec_1$to successfully transform Eqn 28 to get Eqn 29. **This can answer your core question about the usage of$c_1$**. As for the$O(1)$ claim made in Eqn 33, please see our response to Q4 below. In Line 738-783 of our latest uploaded PDF, we have made the corresponding improvements to make our derivation process clearer.

Q3: Assumption 2 is 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.

Our Response: We apologize for any inconvenience caused by our writing, please allow us to clarify it. The following part has also been revised in Line 122-126, Line 150-154 and Line 795-810 of our latest uploaded PDF to make this point clear. (We found that OpenReview's rendering of latex formulas is unstable (depending on the browser version), so the display of the following formulas may be abnormal. We recommend you to read Line 122-126, Line 150-154 and Line 795-810 of our latest uploaded paper.)

Q3.1: What is the "delimiter hidden state" precisely and what is the input format of $R$?

A3.1: $R$ is the list that contains many retrieved passages, the format of $R$ is like:

Thank you for your insightful feedback! We follow your suggestions to revise our paper and the new PDF version has been uploaded. The revised content is marked in blue font in our paper. We believe that our response and revisions address your concerns regarding the clarity of our theory and generalizability of our practical method.

• As for the theory, we follow your suggestions to make the definitions and assumptions more clear and add them to the main paper. We answer every one of your questions and solve each of them in the new PDF version to make them easier for readers to understand.
• As for the practical method, the main concern is about the generalization of our method to more tasks. We follow your suggestions to compare our method with baselines on more tasks such as dialogue, code generation, slot filling, language modeling and long-form question answering. Experimental results further show that our method can also achieve state-of-the-art performance in many other tasks besides QA. These experiments are highlighted in Table 2 of the new PDF version. The method based on our theory has shown excellent performance in various tasks without additional training (while baselines need), which further confirms the correctness of our theory.

Following is our response. We look forward to discussing with you!

Response to Weakness

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

Our Response: We appreciate your efforts to make our theory more understandable and sound, and apologize for the inconvenience our writing has caused you. Regarding this weakness you mentioned, we have responded it in Our Response to Questions below. We answer every one of your questions and solve each of them in the new PDF version to make them easier for readers to understand. We sincerely hope that our responses can make our theoretical analysis clearer and change your thoughts to our paper.

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

Our Response: We appreciate the reviewer for pointing out this point. Please allow us to clarify this. In fact, not only our method, but all the methods in RAG baseline that need to fine-tune LLMs, such as Retrobust, Self-RAG, and INFO-RAG, face the same problem, i.e. needing to access the parameters of LLM and updating them. Compared to these baselines, our method just needs to access LLMs' parameters at inference time without updating them and achieves better performance (Table 1 in our paper). We follow your advice to discuss this point in Line 92-94 of intro.

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

Our Response: Thanks for you propose this insightful suggestion to our method more solid. We follow your suggestions to compare our method and baselines in more various tasks (dialogue, code generation, slot filling, language modeling and long-form question answering) in RAG setting and the experimental results shown in the below Table indicate that our method can also achieve state-of-the-art performance in many other tasks besides QA. In our latest uploaded PDF, we have added this table to the main paper as Table 2 and the specific experimental settings are introduced in Appendix I.

In this experiment, as for the retrieval model and retrieval database, for Slot Filling and Language Modeling, we use ColBERTv2 , a late-interaction model with excellent generalization ability as the retriever, and use Wikipedia consisting of 21,015,324 passages as retrieval database. For Code Generation, we SCODE-R as code retriever and use deduplicated source codes in CodeSearchNET as retrieval database. For all the above tasks, we give Top-5 retrieved passages to each example. For ELI5, Dialog, we use the list of contextual passages provided in the datasets as the retrieved list (distractor setting). All baselines and our method share the same retrieved documents.

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

Our Response: We invite you to review Line 211-219 in our paper, which is our discussion about the connection between distribution completion (contradiction) and benefit (detriment).

• Distribution completion in our Eqn 9 measures how much out-of-distribution knowledge that retrieved texts provide to LLM in the prediction of the token, which can be defined as benefit.
• Distribution contradiction in our Eqn 9 measures the LLM’s resistance to any external knowledge in the retrieved texts that conflicts with LLM's knowledge, which can be defined as detriment.

Our Eqn 9 shows that the next token prediction in RAG is governed by the combined effect of benefit and detriment, whose relationship is described as subtraction (distribution completion - distribution contradiction).

As for your concern that "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", this situation has already been measured by our distribution completion rather than distribution contradiction. Our distribution contradiction focuses on LLM's resistance to external knowledge rather than just difference.

Q3.2: How assumption 2 leads to $O(1)$ in line 765 ($p(h\_{i-1}^{d}|r\_{1:i-1},z)=O(1)$)?

A3.2: In Line 795-810 of our latest uploaded PDF, we have made this point clear.

Let's decompose $p(h\_{i-1}^{d}|r\_{1:i-1},z)$ with chain rule:

$

p(h_{i-1}^{d}|r_{1:i-1},z) = \sum_h p(h_{i-1}^{d}|h,z)p(h|r_{1:i-1}z). \notag $$ Assumption 2 assumes that for any delimiter hidden stateh^d$and any other hidden state$h$, the transition probability from$h$to$h^d$has upper and lower bound as:$0 \leq c_2\leq p(h^d|h,z) \leq c_3$. Then we can get:

$

p(R,x_{1:i-1}|z)=p(x_{1:i-1}|R,z)p(R|z) \notag \text{ , // chain rule} \notag $

$

v(z) &= \textrm{log}\frac{p(R,x\_{1:i-1}|z)}{p(R,x\_{1:i-1}|z^*)} \notag \text{  , // our definition in Eqn 5 of our paper.} \\\\
&= \textrm{log}\frac{p(x\_{1:i-1}|R,z)p(R|z)}{p(x\_{1:i-1}|R,z^*)p(R|z^*)} \notag \\\\
&= \textrm{log}\frac{p(x\_{1:i-1}|R,z)}{p(x\_{1:i-1}|R,z^*)} + \textrm{log}\frac{p(R|z)}{p(R|z^*)} \notag  \notag \\\\
& \leq  -\textrm{log}{c\_1} + \textrm{log}\frac{p(R|z)}{p(R|z^*)} . \notag

\notag $

$

\textrm{log}\frac{p(x_{1:i-1}|R,z)}{p(x_{1:i-1}|R,z^)} &\leq -\textrm{log}{c_1}\notag \\ \frac{p(x_{1:i-1}|R,z)}{p(x_{1:i-1}|R,z^)} &\leq \frac{1}{c_1} \\ p(x_{1:i-1}|R,z) &\leq \frac{p(x_{1:i-1}|R,z^*)}{c_1} \notag $$ Sop(x_{1:i-1}|R,z)$is bound by$O(1)$ and then we can get:

$

p(R,x_{1:i-1}|z) \approx \prod_{i=1}^{n} O(1)p(r_i|r_{1:i-1},z).
\notag $

$

e_{max} = \frac{\sqrt{\Upsilon}}{\text{exp}(\Upsilon)}.
\notag $$ According to Eqn 9 in our paper, the detriment\Upsilon$is the KL divergence between two distribution, so the value range of$\Upsilon$is$(0,+\infty)$. So we can get:

Dear Reviewer aryh

Thanks for your response and your insightful suggestions. We are sorry to bother you but since the discussion phase will be ending in a few days, we would like to know can our responses address your concerns and change your assessment of our paper? If you still have any questions or concerns, we are glad to discuss with you!

We sincerely appreciate your time and we are looking forward to your reply.

Thank you!

We are willing to answer your question. Assumption 1 says there is some hidden state $h$ lower bounds $p(x|h,z*)$ > $c_1$, it means $h$ is the lower bound hidden state that makes $p(x|h,z*)$ greater than $c_1$, so all hidden states make $p(x|h,z*)$ greater than $c_1$, as you said. We apologize for this ambiguity and promise to make it clearer as you suggested.