Retrieval Head Mechanistically Explains Long-Context Factuality

Authors: Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, Yao Fu
Primary area: foundation or frontier models, including LLMs
Keywords: Large language models, long context, interpretability, attention

Abstract:
Despite the recent progress in long-context language models, it remains elusive how transformer-based models exhibit the capability to retrieve relevant information from arbitrary locations within the long context. This paper aims to address this question. Our systematic investigation across a wide spectrum of models reveals that a special type of attention heads are largely responsible for retrieving information, which we dub retrieval heads. We identify intriguing properties of retrieval heads:(1) universal: all the explored models with long-context capability have a set of retrieval heads; (2) sparse: only a small portion (less than 5\%) of the attention heads are retrieval. (3) intrinsic: retrieval heads already exist in models pretrained with short context. When extending the context length by continual pretraining, it is still the same set of heads that perform information retrieval. (4) dynamically activated: take Llama-2 7B for example, 12 retrieval heads always attend to the required information no matter how the context is changed. The rest of the retrieval heads are activated in different contexts. (5) causal: completely pruning retrieval heads leads to failure in retrieving relevant information and results in hallucination, while pruning random non-retrieval heads does not affect the model's retrieval ability. We further show that retrieval heads strongly influence chain-of-thought (CoT) reasoning, where the model needs to frequently refer back the question and previously-generated context. Conversely, tasks where the model directly generates the answer using its intrinsic knowledge are less impacted by masking out retrieval heads. These observations collectively explain which internal part of the model seeks information from the input tokens. We believe our insights will foster future research on reducing hallucination, improving reasoning, and compressing the KV cache.

Decision: accept (oral)
Reviewer scores: 8, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1730605048182

thread_type:
official_review

<thread_object>
creator: reviewer_xfgx

<content>
summary:
The paper investigates how transformer-based models extract relevant information from long context. It identifies a specific type of attention head, named retrieval head, which plays a significant role in the information retrieval process. The authors demonstrate how they detect retrieval heads and describe their characteristics through various experimental settings. Additionally, they conducted experiments that involved pruning the retrieval heads to show that these heads are essential for recalling specific information amidst vast amounts of data. Their findings indicate that retrieval heads are crucial for extractive question-answering and chain-of-thought reasoning.

strengths:
- The paper shows originality by exploring the retrieval capabilities of transformer language models, a topic that has not been extensively studied.
- The paper is well-organized and easy to follow. The paper first defines the special type of attention head that plays a significant role in recalling relevant information during generation. Then, it successfully proves the existence of the retrieval heads within Needle-in-a-Haystack (NIAH) task and further illustrates several properties of the retrival heads which are quite interesting.
- Figures are clear and intuitive to comprehend  
- The paper is impactful as it proposes prospective research directions involving retrieval heads.

weaknesses:
- The work is limited to the Needle-in-a-Haystack (NIAH) task. Although NIAH is a good task to prove the existence of the retrieval heads, we do not know if the similar findings and significance would transfer to other tasks where the LM needs to paraphrase or utilize the previous context (not just copy-and-paste), which are more complex and closely related to real-world applications.

questions:
- It would be better to show that similar findings transfer to paraphrasing tasks. 
- More appendix figures for evaluation results of Retrieval Head Detection Algorithm on various settings (line 199)

- Minor fixes
line 247-250 repetition
Figure 9 caption typo: needels
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700700152

thread_type:
official_review

<thread_object>
creator: reviewer_f46b

<content>
summary:
This paper identifies and analyzes the properties of retrieval heads, a specialized type of attention head primarily responsible for retrieving information. Key properties include: 

1. **Universality** – retrieval heads are present across all explored models.
2. **Sparsity** – only a small subset of attention heads serve this retrieval function.
3. **Intrinsic nature** – these heads exist in pretrained models, even those trained on short contexts.
4. **Dynamic activation** – their activation varies depending on specific tokens and contexts.
5. **Causality** – pruning these heads leads to significant performance drops.

By examining the influence of retrieval heads across various tasks, these findings shed light on which internal model components actively seek information from input tokens.

strengths:
1. Provides in-depth analysis and extensive experiments on various properties of retrieval heads.

2. Well-written, with clear graphics, easy-to-follow explanations, and a well-organized structure.

3. Numerous examples and case studies effectively illustrate the properties, making the concepts easy to understand.

weaknesses:
Nothing major needs to be addressed. Please address some discussion questions in question sections.

questions:
1. Are retrieval heads consistent across different architectures? Does the role of retrieval heads vary with different transformer architectures (e.g., decoder-only vs. encoder-decoder models), or are these properties universally applicable?

2. How does the model determine when retrieval heads should be dynamically activated? Is there a mechanism or threshold within the model that dictates when these retrieval heads become active, especially in different contexts?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730755211211

thread_type:
official_review

<thread_object>
creator: reviewer_zppp

<content>
summary:
The paper provides a systematic examination of a specific type of attention head, termed "retrieval heads," which primarily handle information retrieval from input data. It introduces an approach based on the Needle-In-a-Haystack (NIAH) setup to empirically identify retrieval attention heads across various transformer-based architectures. The findings demonstrate that: 
1) retrieval heads are present across a diverse set of models,
2) only a small subset of attention heads function as retrieval heads,
3) these retrieval heads exist even in models pretrained with limited context, suggesting they are an intrinsic artifact of pretraining,
4) they are dynamically activated rather than continuously active, and
5) there is a causal link between retrieval heads and the model’s capability to retrieve relevant information.

strengths:
Overall, the paper is well-structured, the assumptions are clear, and, with a few exceptions listed in the "weaknesses", the methodology is clear. The results and experiments robustly support the authors’ claims.

weaknesses:
The main concern with the paper is the lack of details on the benchmark used to generate the needle-and-haystack pairs. The paper does not clarify how these pairs are created, the diversity of pairs (e.g., in topic, token variety), or the validation methods used. Including these details would provide reviewers with valuable insights into the experimental design, helping them better assess the generalizability and universality of the findings.


Additional minor points:
1) The acronym "KV" is not defined anywhere in the paper. Based on the context, it likely stands for "Key-Value," but this should be explicitly stated.
2) The caption for Figure 3 could benefit from significant revision. It’s challenging to interpret without detailed reference to the discussion, so adding clarifying information in the caption itself would help.
3) Figure 5 is difficult to interpret; it’s unclear what is being visualized. Overlaying the heatmaps for comparison could enhance clarity, or, if this visualization is redundant given the results in Figure 6, consider omitting it.
4) The discussion in Section 4, in particular, would benefit from qualitative examples (perhaps as an appendix) to illustrate and substantiate claims related to downstream tasks.

questions:
Please check the list of minor points in "weaknesses" section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730791453781

thread_type:
official_review

<thread_object>
creator: reviewer_br4p

<content>
summary:
This paper investigates the mechanism with which transformer-based language models "retrieve" information in the long context. It experimented with four model families, six model scales, and three types of post-training variants, and reveals that a special type of attention heads are largely responsible for retrieving information (either copy-paste or paraphrase) from long contexts. Such attention heads are named “retrieval heads”. The authors find that these retrieval heads
- (1) exist in all the explored models,
- (2) are only 5% of the attention heads,
- (3) exist in models large-scale-pretrained with long or only short contexts and remain the same when the models are continually pretrained on longer contexts,
- (4) are dynamically activated given different contexts, and
- (5) will cause degradation in retrieval abilities or chain-of-thought abilities if pruned.

strengths:
- The paper is well-written, with a clear overarching research question (”How do transformer-based language models acquire long-context capabilities?”), the substantial findings of the existence of “retrieval heads” and their properties, and experiments to support each finding.
- The experiments are extensively conducted on LLaMA, Yi, Qwen, and Mistral model families, at various scales from 6B to 8x7B, on base and chat models, and model dense and Mixture-of-Experts models.
- To identify which attention head is contributing to the retrieval from contexts, the authors proposed a novel retrieval score to measure the frequency of a head’s copy-paste behavior during autoregressive decoding. This retrieval score is both analyzed in different models and used to guide the experiments that prune or mask retrieval heads to understand the causal importance of retrieval heads.
- The authors also considerately report empirical results that these identified retrieval heads are activated during paraphrasing, QA, CoT reasoning tasks, and not just in copy-paste tasks.

weaknesses:
- It’d facilitate reading to clarify that "$k$ is a sentence that is irrelevant to $x$" in L146, instead of, for example, a short phrase or a single word. Can add a reference to Figure 2 so that readers see an example.
- The paper misses dataset details (L195, L355, L427). Are NIAH samples created manually or by prompting large language models? What datasets are used to begin with in Sec 2-4? What additional evaluation tests did you create in Sec 4.1-4.3?
- The paper misses experimental details, such as prompts used, links to existing assets used, etc.
- The questions below need to be addressed.

questions:
- L156: By $\mathcal{R}$, do you mean real numbers? If so, perhaps use $\mathbb{R}$ instead and clarify that $a$ refers to unnormalized attention scores.
- Figure 3: Seems that in fact less that 1% of attention heads are activated more than 50% times. The 5% in the caption could probably be changed to 1%.
- L194: Does it happen that the model generates a word $w$ that is not the next token in the needle $k$ that should be generated? If this happens, do you skip the example? Or consider that as a case when all attention heads do not perform a copy-paste operation, even if an attention head actually pays the most attention to the token that should be generated next?
- L203: What do you mean by “stabilize” and “converge”? Please either provide definitions or plots to illustrate.
- Figure 7: Could be nice to include the dataset name in the caption.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731995416231

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
[Update on Experiments] We also found Jamba has retrieval heads

comment:
We are excited to observe that Jamba, a hybrid model combining **state-space layers (Mamba)**, mixture-of-experts (MoE), and a limited number of attention layers (4 layers × 32 heads = 128 total attention heads), also exhibits retrieval heads. 
 Interestingly, these retrieval heads within the attention layers appear to play a key role in Jamba’s retrieval capabilities. 
Below, we provide a comparison of **masking top-k retrieval heads**  versus **random heads** in  [Jamba-v0.1](https://huggingface.co/ai21labs/Jamba-v0.1)  (12B active parameters, with a total of 52B parameters across all experts) oon needle-in-a-haystack experiments using the same settings as Figure 7 in our paper:
| Masking Head number | 0   | 2    | 5    | 10   | 15   | 20   | 30   | 50   | 100  |
|-------------------|-----|------|------|------|------|------|------|------|------|
| Masking Random Head  | 100 | 99.1 | 98.0 | 94.7 | 90.3 | 85.1 | 70.3 | 44.1 | 9.0  |
| Masking Top Retrieval head| 100 | 98.1 | 38.5 | 61.9 | 33.0 | 15.3 | 10.3 | 12.9 | 3.5  |

These results provide strong evidence supporting our hypothesis that “full attention is crucial for effective long-context information retrieval” (L477) in the section “Relationship to Local and Linear Attention and State-Space Models”. Masking top retrieval heads significantly impacts retrieval performance for Jamba, confirming their essential role in maintaining model capabilities.
We will incorporate these updated results into Figure 7, along with comparisons to other models, to further strengthen our findings.

Please refer to Section 7.1 in the appendix to see the new figure with Jamba.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731997177736

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Update on Experimental Details

comment:
We thank all reviewers for their valuable feedback regarding the experimental details. In response to these comments, we have updated the paper to include the required details. Also we elaborate these information below:

### **1. NIAH experiments details:**
Before rebuttal:
- We manually construct / write four sets of (question, answer) pairs that are semantically irrelevant to the long document ("haystack"). The long documents for all four sets are randomly sampled from the publicly available [SlimPajama](https://huggingface.co/datasets/yaofu/slimpajama-per-source-length-upsample/viewer/default/train).
- We use three sets of (question, answer) pairs for retrieval head detection (i.e., calculating the retrieval scores discussed in Section 2), and reserve one set of (question, answer) pair for testing (generating Figure 1).
- A maximum sequence length of 50K tokens is used for retrieval head detection, while the full 128K token length is used during testing. This ensures that the retrieval head generalizes to longer lengths than detecting them.

The needles for retrieval head detection are: 

1. Needle: A new report from the WMO shows that records were once again broken, and in some cases smashed, for greenhouse gas levels, surface temperatures, ocean heat and acidification. Question: What does a new report from WMO shows?
2. Needle: The best thing to do in Beijing is to take a walk in Chaoyang Park and have a cup of Espresso in the evening. Question: What is the best thing to do in Beijing?
3. Needle: Mr Green is disliked by everyone because he is a mean person and also he can't ride a horse or dive a car. Question: Why does everyone dislikes Mr Green?

During rebuttal:

Since reviewer ZPPp is concerned about the diversity of the needle, we use the above three cases as in-context examples to a language model, and ask the language model to generate 100 more examples. By doing this we get a synthetic data of 1.2K unique tokens spanning 10 topics (Technology, Transportation, Education, Festivals, Health .etc). Below are three examples of the synthetic data. The full needle set is uploaded as a supplementary material: 

1. Needle: The newly constructed SkyBridge connects three cities—Everdale, Pinehurst, and Riverpoint—allowing citizens to travel with ease, admire scenic views, and save significant commute time. Question: What does the newly constructed SkyBridge connect and offer?
2. Needle: The upcoming GalaxyFest will feature over 100 sci-fi authors, 50 exclusive book signings, and a virtual reality experience of Mars colonization. Question: What features will the upcoming GalaxyFest have?
3. Needle: The Aurora Conservatory is renowned for its collection of rare Arctic flora, cutting-edge climate research, and eco-friendly glass dome architecture. Question: Why is the Aurora Conservatory renowned?

Repeating the retrieval head detection algorithm on Mistral 7B with our newly generated 100 needles, we get the same set of retrieval heads as before. This is to say, we confirm that out conclusion holds when scaling the number of needles from 3 to 100. We have added the details in appendix Figure 14. 

We apologize for any confusion caused by the initial omission of these details and have updated the paper to include this information. Our full reproducible code and data is open-sourced, and we will reveal the link after the anonymity period.

### **2. Other Experiments and Prompts**
As our primary focus is the fundamental properties of retrieval heads, we do not emphasize prompt engineering, opting instead for the simplest possible prompts.
The datasets used, excluding **ExtractiveQA** (Details are mentioned in section 4.2 and we will open source it) and **NIAH**, are publicly available:
| **Dataset**       | **Prompt & Source**                                                                                                                  |
|--------------------|--------------------------------------------------------------------------------------------------------------------------------------|
| **Musique**        | Prompt and data from [LongBench](https://huggingface.co/datasets/yanbingzheng/LongBench). The CoT version adds: "please first think step by step." |
| **GSM8K and MMLU**    | Prompt and data from [Chain-of-thought-hub](https://github.com/FranxYao/chain-of-thought-hub).                                      |

We hope these clarifications address the reviewers' concerns
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper investigates how transformer-based models extract relevant information from long context. It identifies a specific type of attention head, named retrieval head, which plays a significant role in the information retrieval process. The authors demonstrate how they detect retrieval heads and describe their characteristics through various experimental settings. Additionally, they conducted experiments that involved pruning the retrieval heads to show that these heads are essential for recalling specific information amidst vast amounts of data. Their findings indicate that retrieval heads are crucial for extractive question-answering and chain-of-thought reasoning.