Venue: NAACL 2024

# Unveiling the Magic: Investigating Attention Distillation in Retrieval-Augmented Generation

Zizhong Li Haopeng Zhang Jiawei Zhang

IFM Lab, University of California, Davis {zzoli, haphahang, jiwzhang}@ucdavis.edu

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

+
Footnote †: Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 2: Short Papers), pages 745–754

June 16-21, 2024 02024 Association for Computational Linguistics

## 1 Introduction

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, in the inference phase, their frozen 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). Retrievail-augmented 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 Shi et al. (2023); Shuster et al. (2021).

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 Karpukhin et al. (2020); Jiang et al. (2023). Among these, attention score-based knowledge distillation stands out due to its notable effectiveness in question-answering QA tasks Izacard and Grave (2020), outperforming other established methods Karpukhin et al. (2020); Lewis et al. (2020); Izacard and Grave (2020). 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

Figure 1: Training _Contriever_ on _NaturalQuestions_ for the QA task with attention distillation shows an improved Hit Rate @ 5 with a fine-tuned reader but a significant decline with an off-the-shelf reader.



satisfactory results (Hu et al., 2023; Wang et al., 2023).

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 token-level 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:

* We conduct an extensive analysis of attention scores in large language models, mainly focusing on the prevalent decoder-only structure, to understand their impacts 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.

## 2 Method

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 question-answering 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 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 \(\alpha_{t}\) but also the norm of the value vector **v\({}_{t}\)** should be taken into account (Izacard et al., 2022b). Af

Figure 2: The framework of the Retrieval-augmented Language Model of our experiment.

 

terwards, 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}_{t=1}\alpha_{t}||\mathbf{v}_{t}||_{2}) \tag{1}\]

where \(T\) represents the total token count 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}_{k=1}exp(s(n_{k},Q)/ \theta)} \tag{2}\]

where \(s\) denotes the dot-product between the representation vectors of the input question \(Q\) and document candidate \(n_{i}\), and \(\theta\) is the temperature hyper-parameter.

During training, the attention score's distribution is distilled into the retriever by minimizing KL-divergence between \(P_{ATTN}(n_{i})\) and \(P_{RETR}(n_{i})\), which aligns the retriever's behavior more closely with the insights derived from attention scores. Figure 2 visually illustrates the retrieval process and the utilization of attention scores during training.

## 3 Experiments

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. (2022) in using _Contriver_ as the retriever model. During the retrieval process, we set the number of retrieved documents \(D_{n}\) to a fixed size \(k=5\) to balance training costs with the amount of information retrieved, thereby avoiding inefficiencies of either extreme.

### Experiment Setup

**Dataset** We assess the model's performance using the _NaturalQuestions_Kwiatkowski et al. (2019) and the _TriviaQA_Joshi et al. (2017) benchmarks, which are the two most popular dataset in the QA task. For the knowledge base, we utilize data from Wikipedia as of December 20, 2018.

**Experimental Settings** Specifically, we use the following settings for our experiments.

**1) Off-the-shelf Distillation Training**: We synchronously train the model using the initial _Falcon-1b_Penedo et al. (2023) as the reader and _Contrierer_Izacard et al. (2022) as the retriever.

**2) Fine-tuned Distillation Training**: This experiment involves two steps:

**Step1.** We start with the initial _Falcon-1b_ checkpoint as a reader and _Contriver_ as a retriever, only fine-tuning the reader model while keeping the retriever model's parameters fixed.

**Step2.** We continue training the retriever model using the fine-tuned reader checkpoint from Step1, updating the knowledge base index periodically.

**Evaluation Metrics**: We assess the model performance in terms of retrieval quality and question-answering correctness, given the involvement of both retriever and reader models. We use the _top-5_ retrieval Hit Rate (HR@5), which is the proportion of retrieved documents \(D_{n}\) containing at least one answer \(A\), to measure the retriever's effectiveness. For the reader's QA performance, we employ the standard Exact Match (EM) metric and F1-Score.

### Results and Discussion

In this section, we empirically analyze the effectiveness of attention distillation training by answering the following research questions:

_RQ1: When does the attention distillation work?_

As shown in Table 1, the _Fine-tuned Distillation Training_ after Step2 shows the best performance in both EM, F1 and HR@5. In contrast, _Off-the-shelf Distillation Training_ performs the worst, with its retriever even underperforming the initial Contriever model (i.e., the retriever model of _Fine-tuned Distillation Training_ Step1). Notice that the critical difference lies in the quality of the reader models: _Off-the-shelf Distillation Training_ uses the initial Falcon-1b model, whereas _Fine-tuned Distillation Training_ employs a well-tuned Falcon-1b. These experimental results strongly suggest that the quality of attention scores is pivotal: **attention scores from the high-quality readers enhance training, whereas low-quality ones lead to poor interaction between the retriever and the reader.**

_RQ2: Are there any commonalities in attention scores from the high-quality readers?_

We sample 1000 data instances from each exper

\begin{table}
\begin{tabular}{c c c c c} \hline \multirow{2}{*}{**Method**} & \multicolumn{4}{c}{**Evaluation Metrics**} \\ \cline{2-6}  & **Dataset** & **EM’** & **FN1’** & **HR@5† \\ \hline Off-the-shelf Distillation & NQ & 27.24 & 33.62 & 0.030 \\ TriviaQA & 30.55 & 35.24 & 0.022 \\ \hline Fine-tuned Distillation (Step1) & NQ & 31.76 & 38.72 & 0.391 \\ TriviaQA & 44.62 & 50.79 & 0.516 \\ \hline Fine-tuned Distillation (Step2) & NQ & 35.22 & 43.44 & 0.645 \\ TriviaQA & 54.59 & 61.04 & 0.643 \\ \hline \end{tabular}
\end{table}
Table 1: Model’s Performance of Different Experimental Settings 

iment to obtain reliable analysis results. We focus on the attention score characteristics **at token level** to identify which tokens receive more attention from high-quality signals. Our analysis firstly finds that in the high-quality readers, the tokens most related to _answer_ and _nouns in question_ receive the most attention. Based on our initial observations, we secondly focus on studying the distribution of attention scores for _answer-related_ and _question-related_1 tokens. We use token embedding's _cosine similarity_ to measure its proximity to targets (i.e., answer or nouns in question), selecting the top 5% and top 10% of closest tokens and analyzing their average _attention scores_ and _Spearman correlation with similarity to target tokens_, as shown in Table 23. We also include the _Off-the-shelf Checkpoint_ as a baseline to observe attention score evolution in different settings. This analysis identifies the key commonalities in high-quality attention scores.

Footnote 1: We only focus on the nouns in the question in selecting _question-related_ tokens.

Footnote 2: The highest values in the table are highlighted in bold on the NQ Dataset and underlined on the TriviaQA Dataset.

**Commonality1. Higher attention to answer tokens in higher-quality models.** In all training settings, tokens closer to answer tokens (i.e., from a similarity higher than \(90^{th}\) percentile to a similarity higher than \(95^{th}\) percentile) receive increasingly higher attention scores. It can be observed that for both two measure metrics, the _Off-the-shelf Distillation Training_ results are lower compared to the _Off-the-shelf Checkpoint_, while _Fine-tuned Distillation Training_ shows improvement in both Step1 and Step2. The results suggest that in _Off-the-shelf Distillation_, the reader's attention does not effectively "highlight" key information, leading to suboptimal training. In contrast, _Fine-tuned Distillation_ after Step1 and Step2 both indicate that high-quality readers focus more on relevant answer tokens, thereby enhancing both the retriever's performance and the relevance of attention allocated to these tokens, which is also revealed in Figure 3.

**Communality 2. Tokens similar to question nouns receive more attention in high-quality models.** Table 2 also indicates that tokens closer to the nouns in question tokens receive higher attention scores. The _Fine-tuned Distillation_ experiments exhibit much higher values in both metrics compared to _Off-the-shelf Checkpoint_ and _Off-the-shelf Attention Distillation_, aligning with their superior performance. However, unlike Commonality 1, the Spearman correlation between attention to question-related tokens and model performance isn't consistent: while _Fine-tuned Attention Distillation_ Step2 surpasses Step1, its metric values do not consistently align with this improvement, suggesting a more complex relationship.

_RQ3: How do we evaluate the quality of attention distillation on decoder-only readers based on the analysis results?_

**Indicator1.** Focusing on the attention scores of the nearest tokens to answer \(A\), denoted as \(M_{A}=\{ma_{1},...,ma_{k}\}\). Higher average \(P_{ATTN}(ma_{i})\) values indicate better attention distillation quality. Additionally, a higher average Spearman correlation between the \(P_{ATTN}(ma_{i})\) and their semantic similarity to \(A\) also signifies better quality.

**Indicator2.** Examining the attention scores of tokens closest to nouns in question \(Q\), denoted as \(M_{Q}=\{mq_{1},...,mq_{k}\}\). An increase in average \(P_{ATTN}(mq_{i})\) suggests better quality. Moreover, if the average Spearman correlation between the attention scores of \(M_{Q}\) and their similarity to \(Q\) is above the threshold for a weak monotonic relationship (i.e., value > 0.3), the attention distillation quality is considered good.

_RQ4: Can we extend the proposed indicators to encoder-to-decoder structure readers?_

\begin{table}
\begin{tabular}{c c c c c c c c c c} \hline \hline  & & \multicolumn{6}{c}{**Answer-related**} & \multicolumn{3}{c}{**Question-related**} \\ \cline{3-8} Experiment & **Dataset** & \multicolumn{2}{c}{**90\({}^{th}\) percentile**} & \multicolumn{2}{c}{95\({}^{th}\) percentile} & \multicolumn{2}{c}{90\({}^{th}\) percentile} & \multicolumn{2}{c}{95\({}^{th}\) percentile} & \multicolumn{2}{c}{95\({}^{th}\) percentile} \\ \cline{3-8}  & & **Attin.** & **Corr.** & **Attin.** & **Corr.** & **Attin.** & **Corr.** & **Attin.** & **Corr.** \\ \hline Off-the-shelf Checkpoint & NJQ & 0.033 & 0.227 & 0.039 & 0.196 & 0.023 & 0.103 & 0.024 & 0.092 \\  & TriviaQA & 0.027 & 0.218 & 0.032 & 0.206 & 0.021 & 0.103 & 0.023 & 0.067 \\ \hline Off-the-shelf Attention Distillation & NO & 0.017 & 0.145 & 0.017 & 0.076 & 0.027 & 0.139 & 0.039 & 0.153 \\  & TriviaQA & 0.031 & 0.160 & 0.035 & 0.172 & 0.047 & 0.144 & 0.063 & 0.260 \\ \hline Fine-tuned Attention Distillation (Step1) & ND & 0.039 & 0.308 & 0.052 & 0.282 & **0.035** & **0.343** & **0.045** & **0.333** \\  & TriviaQA & 0.058 & 0.259 & 0.074 & 0.258 & 0.058 & 0.349 & 0.078 & 0.372 \\ Fine-tuned Attention Distillation (Step2) & MQ & **0.049** & **0.316** & **0.066** & **0.350** & 0.032 & 0.310 & 0.039 & 0.225 \\  & TriviaQA & 0.069 & 0.290 & 0.089 & 0.320 & 0.060 & 0.367 & 0.078 & 0.326 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Average values of attention scores and Spearman correlation in _answer-related_ and _question-related_ tokens 

An analysis with the fine-tuned encoder-to-decoder structure _Atlas-large_ model is presented in Figure 4. The results show that the performance of _Atlas-large_ surpasses _Fine-tuned Distillation Training (Step2)_. However, only the average \(P_{ATTN}(ma_{i})\) trend from Indicator1 applies to this encoder-to-decoder structure model, while _Allas-large_ exhibits a polarized distribution for the Spearman correlation values, as shown in Figure 3 and Appendix A.

**RQ5: Can we extend the proposed indicators to perplexity distillation training?**

Finally, we want to determine if our indicators can apply to perplexity distillation, another popular knowledge distillation method used in training the retriever model. We fine-tune _Atlas-large_ model with the perplexity distillation method and find that the perplexity distribution does not align with either Commonality 1 or Commonality 2, saying that our indicators are not suitable for perplexity distillation (details in Appendix A and B).

## 4 Related Work

The concept of using attention scores for knowledge distillation was introduced by Izacard and Grave (2020), 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 Shuster et al. (2021) and dependency on pre-training data Kandpal et al. (2023), or enhancing retriever efficiency through techniques like specific data sampling Hofstatter et al. (2021). 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).

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.

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

Figure 4: Model performance (top) and their attention distillation analysis (bottom) of _Atlas-large_ model (yellow) for the answer-related tokens, comparing with _Finetuned Distillation Training (Step2)_ (blue).

Figure 3: The attention score distribution histogram (left) and Spearman correlation distribution histogram of \(95^{th}\) percentile _answer-related_ tokens under NQ dataset.

 



## Acknowledgement

This work is partially supported by NSF through grants IIS-1763365 and IIS-2106972. We thank the anonymous reviewers for their helpful feedback.

## Limitation

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 the future work, we will extend the study to larger-scale language models, focusing on validating the accuracy of our analysis on them to enhance the generalizability and applicability of our proposed methods.

## References

* B. BehnamGhader, S. Miret, and S. Reddy (2022)Can retriever-augmented language models reason? the blame game between the retriever and the language model. arXiv preprint arXiv:2212.09146. Cited by: SS1.
* K. Guu, K. Lee, Z. Tung, P. Pasupat, and M. Chang (2020)Retrieval augmented language model pre-training. In International conference on machine learning, pp. 3929-3938. External Links: Link Cited by: SS1.
* S. Hofstatter, S. Lin, J. Yang, J. Lin, and A. Hanbury (2021)Efficiently teaching an effective dense retriever with balanced topic aware sampling. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 113-122. Cited by: SS1.
* L. Hu, Z. Liu, Z. Zhao, L. Hou, L. Nie, and J. Li (2023)A survey of knowledge enhanced pre-trained language models. IEEE Transactions on Knowledge and Data Engineering. Cited by: SS1.
* G. Izacard and E. Grave (2020)Destilling knowledge from reader to retriever for question answering. arXiv preprint arXiv:2012.04584. Cited by: SS1.
* G. Izacard and E. Grave (2020)Leveraging passage retrieval with generative models for open domain question answering. arXiv preprint arXiv:2007.01282. Cited by: SS1.
* G. Izacard and E. Grave (2020)Level-aging passage retrieval with generative models for open domain question answering. arXiv preprint arXiv:2007.01282. Cited by: SS1.
* G. Izacard and E. Grave (2020)User-aging passage retrieval with generative models for open domain question answering. arXiv preprint arXiv:2007.01282. Cited by: SS1.
* G. Izacard and E. Grave (2020)Use

Eneko Agirre, Ilana Heintz, and Dan Roth. 2023. Recent advances in natural language processing via large pre-trained language models: A survey. ACM Computing Surveys56(2), pp. 1-40. Cited by: SS1.
* M. Nasr, N. Carlini, J. Hayase, M. Jagielski, A. Cooper, D. Ipolito, C. A. Choquette-Choo, E. Wallace, F. Tramer, and K. Lee (2023)Scalable extraction of training data from (production) language models. External Links: Link Cited by: SS1.
* L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. (2022)Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems35, pp. 27730-27744. Cited by: SS1.
* G. Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, A. Cappelli, H. Alobeidli, B. Pennier, E. Almazrouei, and J. Launay (2023)The refinedweb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116. Cited by: SS1.
* O. Ram, Y. Levine, I. Dalmedigos, D. Muhlgay, A. Shashua, K. Leyton-Brown, and Y. Shoham (2023)In-context retrieval-augmented language models. arXiv preprint arXiv:2302.00083. Cited by: SS1.
* Weijia Shi, J. Michael, S. Gururangan, and L. Zettlemoyer (2022)Rnn-prompt: nearest neighbor zero-shot inference. Note: URL https://arxiv. org/abs/2205.13792 Cited by: SS1.
* Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih (2023)Replug: Retrieval-augmented black-box language models. arXiv preprint arXiv:2301.12652. Cited by: SS1.
* K. Shuster, S. Poff, M. Chen, D. Kiela, and J. Weston (2021)Retrieval augmentation reduces hallucination in conversation. arXiv preprint arXiv:2104.07567. Cited by: SS1.
* Lei Wang, C. Ma, X. Feng, Z. Zhang, H. Yang, J. Zhang, Z. Chen, J. Tang, X. Chen, Y. Lin, et al. (2023)A survey on large language model based autonomous agents. arXiv preprint arXiv:2308.11432. Cited by: SS1.
* Haopeng Zhang, Xiao Liu, and Jiawei Zhang (2023)Extractive summarization via ChatGPT for faithful summary generation. In Findings of the Association for Computational Linguistics: EMNLP 2023, Singapore, pp. 3270-3278. Cited by: SS1.
* H. Zhang, X. Liu, and J. Zhang (2023)SumMLt: iterative text summarization via chatgrt. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 10644-10657. Cited by: SS1.
* H. Zhang, S. Yavuz, W. Kryscinski, K. Hashimoto, and Y. Zhou (2022)Improving the faithfulness of abstractive summarization via entity coverage control. In Findings of the Association for Computational Linguistics: NAACL 2022, Seattle, United States, pp. 528-535. External Links: Document, ISSN 0021-900X Cited by: SS1.
* M. Zhang, O. Press, W. Merrill, A. Liu, and N. A Smith (2023)How language model hallucinations can snowball. arXiv preprint arXiv:2305.13534. Cited by: SS1.

 

Quantitative Analysis of

Answer-Related Tokens

We present a detailed analysis of _answer-related_ tokens' attention score distribution (or perplexity distribution of _Perplexity Distillation Training_) shown in Table 3. In addition to the histograms visualized for the attention score distribution of \(95^{th}\) percentile _answer-related_ tokens under the NQ dataset (i.e., Figure 3), we also plot the corresponding attention score distribution figure under the TriviaQA dataset, as shown in Figure 5.

## Appendix B Quantitative Analysis of

Question-Related Tokens

We present a detailed analysis of _question-related_ tokens' attention score distribution (or perplexity distribution of _Perplexity Distillation Training_) shown in Table 4, Figure 6, and Figure 7.

## Appendix C The Internal Relationship within Indicators

We also find a strong positive correlation (i.e., Pearson correlation > 0.5 in most cases) between our two proposed indicators in decoder-only language model structure. In contrast, this correlation does not appear in encoder-to-decoder language model structure, which also indicates that Indicator2 is inapplicable to this language model structure.

## Appendix D Dataset Statistics

For the _NaturalQuestions_ dataset, we split it according to the number of 79,168/8,757/3,610 to form the train/validation/test dataset; for the _TriviaQA_ dataset, we split it according to the number of 78,785/8,837/11,313 to form the train/validation/test dataset.

## Appendix E Implementation Details

We conducted all computations on a Nvidia A100 GPU. For the _Off-the-shelf Distillation Training_ and the _Fine-tuned Distillation Training_, we use _Falcon-1b_ as the initial reader model and _Contriever_ as the initial retriever model, which has about 1 billion and 110 million training parameters respectively. For the _Atlas-large Distillation Training_ and _Perplexity Distillation Training_, we use _T5-large_ as the initial reader model and _Contriever_ as the initial retriever model, which has about 770 million and 110 million training parameters respectively.

**Off-the-shelf Distillation Training** We set the batch size to 1, the maximum length of the input prompt to 128, and limit the generation max length to 32. We set the learning rate to 1e-5 and used the Adam optimizer. For _NaturalQuestions_ dataset, we set the total training steps to 160,000 with approximately 2000 warmup steps, training for about 40 hours. For _TriviaQA_ dataset, we set the total training steps to 320,000 with approximately 4000 warmup steps, training for about 60 hours.

**Fine-tuned Distillation Training** For Step 1, we set the batch size to 1, the maximum length of the input prompt to 128, and limit the generation max length to 32. We set the learning rate to 1e-5 and used the Adam optimizer. For _TriviaQA_ dataset, we set the total training steps to 128, and limit the generation max length to 32. We set the learning rate to 1e-5 and used the Adam optimizer. For with approximately 2000 warmup steps, training for about 30 hours. For dataset, we set the total training steps to 320,000 with approximately 4000 warmup steps, training for about 45 hours.

For Step 2, we set the batch size to 1, the maximum length of the input prompt to 128, and limit the generation max length to 32. We set the learning rate to 5e-7 and used the Adam optimizer. For _NaturalQuestions_ dataset, we set the total training steps to 6,000 with approximately 300 warmup steps, training for about 2 hours. For dataset, we set the total training steps to 32,000 with approximately 600 warmup steps, training for about 3 hours.

**Atlas-large Distillation Training** We set the batch size to 1, the maximum length of the input prompt to 128, and limit the generation max length to 32. We set the learning rate to 1e-5 and used the Adam optimizer. For was the total training steps to 6,000 with approximately 300 warmup steps, training for about 2 hours. For dataset, we set the total training steps to 32,000 with approximately 600 warmup steps, training for about 40 hours.

**Perplexity Distillation Training** We set the batch size to 1, the maximum length of the input prompt to 128, and limit the generation max length to 32. We set the learning rate to 4e-5 and used the Adam optimizer. For was the total training steps to 20,000 with approximately 400 warmup steps, training for about 40 hours. For dataset, we set the total training steps to 10,000 with approximately 500 warmup steps, training for about 15 hours.

 



\begin{table}
\begin{tabular}{c c c c c c} \hline \hline
**Experiment** & **Dataset** & **Avg. Attn. (p90)** & **Spearman Corr. (p90)** & **Avg. Attn. (p95)** & **Spearman Corr. (p95)** \\ \hline Off-the-shelf Model Checkpoint & NQ & \(0.023\pm 0.011\) & \(0.103\pm 0.253\) & \(0.024\pm 0.014\) & \(0.092\pm 0.309\) \\  & TriviaQA & \(0.021\pm 0.010\) & \(0.103\pm 0.245\) & \(0.023\pm 0.013\) & \(0.067\pm 0.304\) \\ \hline Off-the-shelf Attention Distillation & NQ & \(0.027\pm 0.010\) & \(0.139\pm 0.237\) & \(0.039\pm 0.017\) & \(0.153\pm 0.341\) \\  & TriviaQA & \(0.047\pm 0.016\) & \(0.144\pm 0.220\) & \(0.063\pm 0.025\) & \(0.260\pm 0.280\) \\ \hline Fine-tuned Distillation Training (Step1) & NQ & \(0.035\pm 0.015\) & \(0.343\pm 0.238\) & \(0.045\pm 0.023\) & \(0.333\pm 0.303\) \\  & TriviaQA & \(0.058\pm 0.024\) & \(0.349\pm 0.222\) & \(0.078\pm 0.037\) & \(0.372\pm 0.285\) \\ Fine-tuned Distillation Training (Step2) & NQ & \(0.032\pm 0.014\) & \(0.310\pm 0.256\) & \(0.039\pm 0.021\) & \(0.225\pm 0.340\) \\  & TriviaQA & \(0.060\pm 0.025\) & \(0.367\pm 0.227\) & \(0.078\pm 0.037\) & \(0.326\pm 0.311\) \\ \hline Atlas-large Distillation Training & NQ & \(0.037\pm 0.027\) & \(0.082\pm 0.251\) & \(0.038\pm 0.032\) & \(0.086\pm 0.345\) \\  & TriviaQA & \(0.047\pm 0.245\) & \(0.076\pm 0.249\) & \(0.050\pm 0.038\) & \(0.081\pm 0.348\) \\ \hline Perplexity Distillation Training & TriviaQA & \(0.063\pm 0.038\) & \(-0.012\pm 0.207\) & \(0.060\pm 0.042\) & \(-0.036\pm 0.297\) \\ \hline \hline \end{tabular}
\end{table}
Table 4: Mean and std. of average attention scores (or perplexity distribution in _Perplexity Distillation Training_) and Spearman correlations of the question-related tokens

Figure 5: The attention score distribution histogram (left) and Spearman correlation distribution histogram of \(95^{th}\) percentile _answer-related_ tokens under TriviaQA dataset.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline
**Experiment** & **Dataset** & **Avg. Attn. (p90)** & **Spearman Corr. (p90)** & **Avg. Attn. (p95)** & **Spearman Corr. (p95)** \\ \hline Off-the-shelf Model Checkpoint & NQ & \(0.033\pm 0.016\) & \(0.227\pm 0.259\) & \(0.039\pm 0.023\) & \(0.196\pm 0.349\) \\  & TriviaQA & \(0.027\pm 0.013\) & \(0.218\pm 0.252\) & \(0.032\pm 0.019\) & \(0.206\pm 0.331\) \\ \hline Off-the-shelf Attention Distillation & NQ & \(0.017\pm 0.008\) & \(0.145\pm 0.193\) & \(0.017\pm 0.010\) & \(0.076\pm 0.254\) \\  & TriviaQA & \(0.031\pm 0.012\) & \(0.160\pm 0.174\) & \(0.035\pm 0.017\) & \(0.172\pm 0.236\) \\ \hline Fine-tuned Distillation Training (Step1) & NQ & \(0.039\pm 0.023\) & \(0.308\pm 0.276\) & \(0.052\pm 0.036\) & \(0.282\pm 0.336\) \\  & TriviaQA & \(0.058\pm 0.031\) & \(0.259\pm 0.261\) & \(0.074\pm 0.050\) & \(0.258\pm 0.331\) \\ Fine-tuned Distillation Training (Step2) & NQ & \(0.049\pm 0.023\) & \(0.316\pm 0.280\) & \(0.066\pm 0.036\) & \(0.350\pm 0.336\) \\  & TriviaQA & \(0.069\pm 0.036\) & \(0.290\pm 0.267\) & \(0.089\pm 0.061\) & \(0.320\pm 0.323\) \\ \hline Atlas-large Distillation Training & NG & \(0.062\pm 0.036\) & \(0.171\pm 0.462\) & \(0.083\pm 0.058\) & \(0.307\pm 0.471\) \\  & TriviaQA & \(0.072\pm 0.045\) & \(0.141\pm 0.379\) & \(0.091\pm 0.067\) & \(0.217\pm 0.438\) \\ \hline Perplexity Distillation Training & TriviaQA & \(0.072\pm 0.039\) & \(0.029\pm 0.142\) & \(0.071\pm 0.042\) & \(0.013\pm 0.202\) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Mean and std. of attention scores (or perplexity distribution in _Perplexity Distillation Training_) and the Spearman correlations of the answer-related tokens.



Figure 6: The attention score distribution histogram (left) and Spearman correlation distribution histogram of \(95^{th}\) percentile _question-related_ tokens under NQ dataset.

Figure 7: The attention score distribution histogram (left) and Spearman correlation distribution histogram of \(95^{th}\) percentile _question-related_ tokens under the TriviaQA dataset.