Unveiling the Magic: Investigating Attention Distillation in Retrieval-augmented Generation

Large language models (LLMs) have showcased remarkable capabilities across various natural language processing tasks  Min et al. (2023); OpenAI (2023); Ouyang et al. (2022); Zhang et al. (2023a, b). However, their fixed parameters limit their ability to update knowledge in real time, resulting in the hallucination problem during generation Zhang et al. (2022, 2023c). Additionally, these models also lack protection for sensitive training data Nasr et al. (2023); Lin et al. (2021). One promising method to overcome these limitations is using retrieval-augmented language models Ram et al. (2023); Shi et al. (2022); Izacard et al. (2022b); Guu et al. (2020); Karpukhin et al. (2020); Khandelwal et al. (2019). Retrieval-augmented language models typically comprise two main components: (1) the retriever, which selects relevant information, and (2) the reader, incorporates this information into the generation process. Combining these two components, retrieval-augmented language models not only improve accuracy and reliability by dynamically using external knowledge but also reduce training costs with fewer trainable parameters Shi et al. (2023); Shuster et al. (2021).

Various methods have been proposed to improve the coordination between the retriever and the reader Karpukhin et al. (2020); Jiang et al. (2023). Among these, attention score-based knowledge distillation has shown its effectiveness Izacard and Grave (2020a), outperforming other established methods Karpukhin et al. (2020); Lewis et al. (2020); Izacard and Grave (2020b) in QA tasks. In this process, the attention scores from the reader are captured and conveyed to the retriever as the supervisory signal, enabling the retrieval model to more effectively identify information candidates that can significantly improve the language model’s responses. This efficient strategy reduces the need for manual annotation of the knowledge corpus, saving resources while achieving satisfactory results Hu et al. (2023); Wang et al. (2023).

However, its efficiency heavily relies on the reader model’s quality. As Figure 1 shows, low-quality reader models yield ineffective supervision signals, detrimentally impacting the retriever’s performance. A fundamental hypothesis underpinning this mechanism is that more attention to certain tokens suggests greater relevance in answering questions Izacard and Grave (2020a). Yet this correlation is not clearly defined: no existing work has quantitatively analyzed the role of attention scores within the retrieval-augmented generation framework. Therefore, our research seeks to understand which text segments gather more attention and how to assess attention quality.

In this paper, we first experimentally confirmed that attention scores are not always effective as training supervisors in certain experiment settings under QA tasks. Following this, we performed a thorough token-level quantitative analysis to identify commonalities in the distribution of attention scores that align with high-quality supervisory signals. Furthermore, we derived two key indicators to measure distillation quality based on the commonalities. Our main contributions are as follows:

• •

  We conduct an extensive analysis of attention scores in language models, mainly focusing on the prevalent decoder-only structure, to understand their impact on retriever model training and the overall performance of retrieval-augmented language models, thereby identifying key factors that significantly influence the model’s performance.

• •

  We introduce novel metrics to evaluate the reader model’s proficiency in attention distillation, aiming to improve training performance by leaning on effective training sessions.

In our experiment, we adapt the ATLAS architecture Izacard et al. (2022b) but use a decoder-only language model structure for our empirical analysis, focusing on QA tasks to study attention score distillation mechanisms. Specifically, for a given question $Q$, we supply models with a knowledge base $D=\{d_{1},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_{n}=\{n_{1},n_{2},...,n_{k}\}\subseteq D$ using the retriever, and then generate the answer $A$ using the reader.

The attention distillation uses attention scores to gauge the importance of each input document in answer generation. To adapt to changes in the reader’s structure, we utilize the self-attention scores related to the output tokens as an indicator for indicating document relevance, rather than relying on cross-attention scores between input documents and output. In addition, notice that the contribution of a token $t$ is not only evaluated from the self-attention score $\alpha_{t}$ but also the norm of the value vector $\textbf{v}_{t}$ should be taken into account Izacard et al. (2022b). Then the Softmax operator is applied to obtain the attention score distribution over the question-relevant documents $D_{n}$:

$$p_{ATTN}(n_{i}|Q,A)=Softmax(\sum_{t=1}^{T}\alpha_{t}||\textbf{v}_{t}||_{2})$$

(1)

where $T$ represents the total number of tokens in $n_{i}$. On the other hand, the retriever’s probability distribution $p_{RETR}$ over $D_{n}$ can be defined as:

$$p_{RETR}(n_{i}|Q)=\frac{exp(s(n_{i},Q)/\theta)}{\sum_{k=1}^{K}exp(s(n_{k},Q)/\theta)}$$

(2)

where $s$ denotes the dot-product of query and document vectors, and $\theta$ is the temperature hyper-parameter. During training, the attention distribution is distilled into the retriever by minimizing KL-divergence between $p_{ATTN}(n_{i})$ and $p_{RETR}(n_{i})$. Figure 2 visually illustrates the retrieval process and the utilization of attention scores during training.

We chose Falcon-1b Penedo et al. (2023) as our primary decoder-only reader model for its performance and flexibility, and we follow ATLAS Izacard et al. (2022b) in using Contriver as the retriever model. During the retrieval process, we fix the retrieved documents $D_{n}$’s size to $k=5$ to balance training costs with the amount of information retrieved, avoiding inefficiencies of either extreme.

The concept of using attention scores for knowledge distillation was introduced by Izacard and Grave (2020a), and the following research has mainly focused on independently optimizing the reader (i.e., the large language model) and the retriever. Previous studies have explored improving large language model performance within the retriever-then-read framework by addressing issues like hallucination Shuster et al. (2021) and dependency on pre-training data Kandpal et al. (2023), and enhancing retriever efficiency through techniques like specific data sampling Hofstätter et al. (2021) or using large language model generators as retrievers Yu et al. (2022). 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 BehnamGhader et al. (2022).

Unlike the previous work, our paper conducts 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 propose novel metrics to identify whether the training process is beneficial, which can apply to all general neural-retrieval-in-the-loop architecture.

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.

This paper analyzes the attention score-based knowledge distillation quality in training retrieval-augmented language models under various experimental settings in QA tasks. Furthermore, based on our findings, we have developed two indicators to assess the quality of attention score supervision. However, our exploration is conducted based on lightweight language models (i.e., language models with about one billion parameters) due to their flexibility and have yet to extend to larger-scale language models. In future work, we will focus on validating the accuracy of our methods on more extensive language models to enhance the generalizability and applicability of our results.