Unveiling the Magic: Investigating Attention Distillation in Retrieval-Augmented Generation (2024)

Abstract: Retrieval-augmented generation framework addresses the limitations of large language models by enabling real-time knowledge updates for more accurate answers. An efficient way in the training phase of retrieval-augmented models is attention distillation, which uses attention scores as supervision signals instead of manually annotated query-document pairs. Despite its growing popularity, the detailed mechanisms behind the success of attention distillation remain unexplored, particularly the specific patterns it leverages to benefit training. In this paper, we address this gap by conducting a comprehensive investigation of attention distillation workflow and identifying key factors influencing the learning performance of retrieval-augmented language models. We further propose several insightful indicators for optimizing models’ training methods and avoiding ineffective training.

## introduction large language models (llms) have showcased remarkable capabilities across various natural language processing tasks @xcite @xcite . however, in the inference phase, their frozen parameters limit their ability to update knowledge in real-time, resulting in the hallucination problem during generation @xcite . additionally, these models also lack protection for sensitive training data @xcite . one promising method to overcome these limitations is using retrieval-augmented language models @xcite @xcite @xcite . retrievalaugmented language models typically comprise two essential components: (1) the retriever, which selects relevant information, and (2) the reader, which incorporates this information into the generation process. the integration of these two components allows retrieval-augmented language models to enhance accuracy and reliability by dynamically utilizing external knowledge, while also reducing training costs due to fewer trainable parameters @xcite . the performance of retrieval-augmented language models may significantly depend on the effective synergy between the retriever and the reader. to this end, various methods have been proposed to improve the coordination between these two components @xcite . among these, attention score-based knowledge distillation stands out due to its notable effectiveness in question-answering (qa) tasks @xcite , outperforming other established methods @xcite @xcite . in this process, the attention scores from the reader are captured and conveyed to the retriever as the supervisory signal (i.e., the retriever uses the attention scores as the basis for assessing the relevance of retrieved information), enabling the retrieval model to identify information candidates more effectively that can significantly improve the language model's responses. this efficient strategy reduces the need for manual annotation of the knowledge corpus, leading to resource savings while still achieving satisfactory results @xcite . however, the success of attention distillation heavily relies on the quality of the reader model. as shown in figure 1 , reader models of low quality yield ineffective supervision signals, detrimentally impacting the retriever's performance. given the critical nature of the issue, it becomes imperative to delve into the mechanism of attention distillation and identify characteristics of low-quality readers to avoid ineffective training. a fundamental hypothesis underpinning attention distillation is that more attention to certain tokens suggests that these tokens are likely to be of greater relevance in answering questions (izacard and grave, 2020a). however, this correlation between attention scores and tokens has not to be clearly established yet, as existing works lack a quantitative analysis in attention scores' impacts within the training process of retrieval-augmented language models. therefore, our research seeks to understand which text segments receive more attention and how to assess the attention distillation quality. in this paper, we first experimentally confirm that attention scores are not always effective as training supervisors across different experimental settings under question-answering tasks. motivated by this observation, we conduct an in-depth tokenlevel quantitative analysis, seeking to uncover patterns within attention scores that correspond to high-quality supervision. we analyze the attention scores from reader models of various qualities and identify a clear, stable correlation between these scores and supervisory quality, especially in tokens related to answers and questions. building on these insights, we derived two key indicators to measure the distillation quality based on the commonalities. our main contributions are as follows: ## method in our experiment, we adapt the atlas architecture @xcite but use a decoder-only language model structure for our empirical analysis, focusing on question-answering tasks to study attention score distillation mechanisms. specifically, for a given question q, we supply models with a knowledge base @xmath0, d 2 , ..., d m }, where each d i is a unique document. the objective of the models is to find the question-relevant documents d @xmath1, n 2 , ..., n k } ⊆ d using the retriever, and then incorporate the query and d n as the input for the reader to generate the answer a for the given question. the attention distillation approach uses attention scores to gauge the importance of each input document d i during the answer generation process. to accommodate changes in the reader model's structure, we utilize the self-attention scores related to the output tokens as a measure of document relevance, rather than relying on cross-attention scores between input documents and output that atlas uses. in addition, the attention level of a token t is not only evaluated from the self-attention score α t but also the norm of the value vector v t should be taken into account @xcite . af-terwards, the softmax operator is applied to obtain the attention score distribution over the questionrelevant documents d n : on the other hand, the retriever's probability distribution p ret r over d n can be defined as: ## experiments we chose falcon-1b @xcite as our primary decoder-only reader model for its performance and flexibility, and we follow atlas @xcite in using contriver as the retriever model. during the retrieval process, we set the number of retrieved documents d n to a fixed size @xmath2, thereby avoiding inefficiencies of either extreme. ## related work the concept of using attention scores for knowledge distillation was introduced by @xcite , and the following research has mainly focused on independently optimizing the reader and the retriever. previous studies have explored improving large language model performance within the retriever-then-read framework by addressing issues like hallucination @xcite and dependency on pre-training data @xcite , or enhancing retriever efficiency through techniques like specific data sampling @xcite . only one study has examined the reader-retriever interaction within a neural-retrieval-in-the-loop architecture, noting that imperfect retrievers can degrade reader performance, though it lacked quantitative analysis @xcite . our study offers a comprehensive quantitative analysis of how the reader and the retriever interact during the neural-retrieval-in-the-loop training architecture under the attention distillation mechanism. we introduce novel metrics to evaluate the efficacy of the training process across all general reader-retriever framework. ## conclusion in this paper, we comprehensively evaluate attention distillation for training retrieval-augmented language models, emphasizing the importance of attention to answer and question-related tokens. we further introduce novel metrics for assessing language models' attention distillation ability to optimize the training process.