Retrieval Head Mechanistically Explains Long-Context Factuality
Large language models;long context;interpretability;attention
Yes
No ethics review needed.
- 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.
- 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.
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.
- 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.
Yes
No ethics review needed.
Please check the list of minor points in "weaknesses" section.
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.
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.
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.
Yes
No ethics review needed.
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?
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.
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.
Nothing major needs to be addressed. Please address some discussion questions in question sections.
Yes
N/A
No ethics review needed.
- 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
- 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.
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.
- 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.
We study retrieval head, a special type of attention head that mechanistically explains long-context factuality
@inproceedings{
anonymous2024retrieval,
title={Retrieval Head Mechanistically Explains Long-Context Factuality},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=EytBpUGB1Z},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Large language models long context interpretability attention
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/767110cc2949ceea3d981df1b0c7cb39a0601614.pdf
foundation or frontier models, including LLMs
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Retrieval Head Mechanistically Explains Long-Context Factuality
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission