Retrieval Head Mechanistically Explains Long-Context Factuality

## Abstract

Abstract Despite the recent progress in long-context large language models (LLMs), it remains elusive how these transformer-based language models acquire the capability to retrieve relevant information from arbitrary locations within the long context. This paper aims to address this question.
Our systematic investigation across 4 model families, 6 model scales, and 3 types of finetuning reveals that a special type of attention heads are largely responsible for retrieving relevant information from long context, which we dub retrieval heads .
We identify important and 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 to 32-128K 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 on retrieval heads foster future research on reducing hallucination, improving reasoning, and compressing the KV cache.

## 1 Introduction

This work studies the internal mechanism of how long-context language models can utilize information at arbitrary locations within the input.
Recent advances in long-context language modeling [, , ] show inspiring results, particularly on the Needle-in-a-Haystack test [], which asks the model to precisely retrieve the information of a short sentence (the needle) within a long context (the haystack).
Such capability is the basis of more advanced long-context tasks, which usually interleaves retrieval and reasoning in a multi-step fashion [].
Based on extensive experiments across 4 model families, 6 model scales, and 3 types of finetuning,
we show that across the models’ attention layers, there exist a small number of retrieval heads that search the information being asked, and redirect the relevant tokens from the input to the output.
Activation of retrieval heads explains whether the output is factual or hallucinated.
When such heads are activated, the model behaves faithful to the input document.
When they are not activated, or intentionally masked out in controlled experiments (Fig. 1), the model cannot find the relevant information and hallucinate instead.

Figure: Figure 1: Retrieval heads are the heads that redirect information from the input to the output. Left: masking out the top retrieval heads of LLaMA 2 7B 80K, its Needle-in-a-Haystack performance drops significantly, and the model hallucinates during decoding. Right: masking out random non-retrieval heads does not influence the model’s needle-in-a-haystack behavior. We further note that the retrieval head only influence factuality but not the language capability: when they are masked, the model hallucinates by saying “go to the beach”, which is still a fluent sentence, but not factual (as the correct answer is “eating a sandwich at Dolores park.”).

The discovery of the retrieval head is motivated by the question of what the attention mechanism is doing when the model can or cannot find the given needle.
We take important inspiration from two existing works: the CopyNet [] and the Induction Head [].
The CopyNet is essentially a single-layer, single-head attention mechanism in the age of RNNs that copy-paste tokens from the input to the output.
Induction Heads [] are special heads within a multi-layer, multi-head attention network that implements an implicit program induction algorithm.
Combining the observation from the two works, we natually hypothesize that, just like induction heads are accountable for in-context learning, there might exist special heads that are accountable for information retrieval and implement a conditional copy-paste algorithm.

We design algorithms to detect retrieval heads within the transformer architecture (Sec. 2), and conduct large-scale experiments to demonstrate important properties of them (Sec. 3):
(1) retrieval heads are universal and sparse: for any model family (LLaMA [], Yi [], QWen [] and Mistral []), at any scale (6B, 14B, and 34B and 8$\times$7B), either base or chat, either dense or MoE, as long as the model can precisely recite the input information, they have a small number of retrieval heads (Fig. 1);
(2) they are intrinsic: the base model (e.g., LLaMA2 base) already contains retrieval heads (as a consequence of large-scale pretraining). Subsequent derivations, such as the long-context continue pretraining (LLaMA2 7B 80K), chat fine-tuning (Qwen Chat), or even sparse upcycling [, ] uses the same retrieval heads as the base model (Fig. 5);
(3) they are dynamically activated according to the context: the strongest retrieval heads (e.g., 13 for LLaMA 2 7B) are always activated no matter what the required information is, while weaker retrieval heads are activated on different parts of the required information; consequently these heads compensate each other’s functionality: removing a subset of the heads, the model at least retrieve part of the required information;
(4) the retrieval heads are causal:
say we put a needle "the best thing to do in San Francisco is to eat a sandwich in Dolores Park on a sunny day",
completely masking out retrieval heads, the model hallucinates (by saying the best thing is to visit Golden Gate bridge);
partially masking out the heads, the model retrieves part of the needle (e.g., it gets the sandwich but forget the Dolores Park);
masking out random non-retrieval heads, the model still find full needle;
when we do not mask the head yet the model still hallucinate in some cases,
the retrieval heads are not activated.
We further note that chain-of-thought reasoning also heavily relies on retrieval heads because the model needs to refer back the input information, indicating a complex relationship between the model’s retrieval and reasoning capability.

The discovery of retrieval head has profound implications on long-context modeling:
(1) it marks a significant step forward in the field of mechanistic interpretability [, ] because for the first time we pin point a particular subnet implementing the conditional retrieval algorithm;
(2) it explains why certain context-compression methods fail to keep factuality (because they removes the retrieval head, e.g., in Xiao et al. ), and suggests future research on KV cache compression [, ], a key problem for deploying long-context models, should consider the influence of retrieval heads.

## 2 Detecting Retrieval Head

Figure: Figure 2: An attention head that performs a copy-paste operation if the token that the head attend to is the same the token being generated. The retrieval score of a head is defined as the frequency of this head’s copy-paste behavior when answering quesions that asks for raw information from the input.

**Table 1: We consider a wide range of language model families and show that the basic properties of retrieval heads are universal and consistent across all language models we study.**
| Base Model | Variant | Variation Type |
| --- | --- | --- |
| Llama-2-7B | Llama-2-7B-80K | Length Extension via Continue Pretrain |
|  | Llama-2-13B-64K | Model Scaling and Length Extension |
| Mistral-7B-v0.2 | Mistral-7B-Instruct-v0.2 | SFT and RLHF |
|  | Mixtral-8x7B-v0.1 | Sparse Upcycling to Mixture of Experts |
| Yi-6B | Yi-6B-200K | Length Extension via Continue Pretrain |
|  | Yi-34B-200K | Model Scaling and Length Extension |
| Qwen1.5-14B | Qwen1.5-14B-Chat | SFT and RLHF |

To detect which head is implementing the retrieval algorithm, we introduct a retrieval score to measures the frequency of a head’s copy-paste behavior during autoregressive decoding.
An attention head with high retrieval score suggests that statistically across various contexts, this head is frequently copying the input tokens from the input to the output.

Needle-in-a-Haystack  Our retrieval head detection algorithm roots from the needle-in-a-Haystack test, which asks the model to copy-paste the input tokens to the output.
Given a question ${\bm{q}}$ and its corresponding answer ${\bm{k}}$ (the needle), we insert ${\bm{k}}$ in a given context ${\bm{x}}$ (the haystack) at a random position index range ${\bm{i}}_{q}$.
The language model is then tasked with answering ${\bm{q}}$ based on the haystack with the inserted needle.
We set ${\bm{q}}$ and ${\bm{k}}$ unique and irrelevant with the given long context,
ensuring that if an answer is correctly generated, it is indeed copied from the context, not from the model’s internal knowledge.

Retrieval Score for Attention Heads  We define the retrieval score as the frequency of a head’s copy-paste operations.
Specifically,
during auto-regressive decoding (we use greedy decoding by default),
denote the current token being generated as $w$ and
the attention scores of a head as $\bm{a}\in\mathcal{R}^{|{\bm{x}}|}$.
As demonstrated in Fig. 2, we say an attention head $h$ copies and pastes a token from the needle to the output sentence if it follows two criteria:
(1) $w\in{\bm{k}}$, i.e., $w$ is a token within the needle sentence. (2)
${\bm{x}}_{j}=w,j=\arg\max(\bm{a}),j\in{\bm{i}}_{q}$, i.e., the input token that receives the most attention probability mass by this head is a token within the needle and is the same token as the currently generated token.
Let ${\bm{g}}_{h}$ be the set containing all tokens copy and pasted by a given head $h$, we define:

$$ $\text{Retrieval score for head}\;\;h=\frac{|{\bm{g}}_{h}\cap{\bm{k}}|}{|{\bm{k }}|},$ (1) $$

Intuitively, retrieval score represents a token-level recall rate of the most attended tokens by an attention head.
For example, when retrieving a needle of 10 tokens, a retrieval score of 0.9 indicates that the attention head has copies and pasted 9 tokens in the 10-token target answer.

Retrieval Head Detection Algorithm  We calculate the retrieval score for all attention heads under a diverse set of input contexts.
For each language model we consider, we compile three sets of Needle-in-a-Haystack samples, each consisting of a unique tuple $({\bm{q}},{\bm{k}},{\bm{x}})$.
For each sample, we make sure $({\bm{q}},{\bm{k}})$ is semantically irrelevant with ${\bm{x}}$ and that ${\bm{q}}$ cannot be answered using the model’s existing knowledge by manually inspecting the model output.
Then for each $({\bm{q}},{\bm{k}},{\bm{x}})$ sample, we perform Needle-in-a-Haystack on 20 different length values uniformly sampled from 1K-50K, where in each length, ${\bm{q}}$ is inserted in 10 different depth uniformly ranging from the start to the end of ${\bm{x}}$.
We note that this scale of tests gives stable outputs as the average retrieval score converges after just a few samples.
In total, each language model is subjected to approximately 600 instances of retrieval testing.
We calculate the retrieval score for each attention head in each test and use the average of these scores as the head’s final retrieval score.
The attention heads with relatively larger retrieve score can be considered as retrieval head.
In our case (Fig. 3), we set the threshold as 0.1, meaning that as long as the head performs copy-paste in 10% of the times, we consider it a retrieval head.

## 3 Basic Properties of Retrieval Heads

Figure: Figure 3: Retrieval heads are universal and sparse across model family and scale. For all models we consider, less than 5% of the attention heads are activated more than 50% of the time (with a retrieval score higher than 0.5) when retrieval is required.

This section discusses important properties of retrieval heads:
(1) universal and sparse: any model that exhibits long-context capability has a small set of retrieval heads;
(2) dynamic: most of retrieval heads are activated under different contexts;
(3) intrinsic: retrieval heads are already within the base model as a consequence of large-scale pretraining. Subsequent models reuse the same set of heads.
Our results are supported by extensive experiments on a large spectrum of models (Table 1).
To investigate the influence of continued pretraining for context length extension, we compare Llama-2-7B 4K to Llama-2-7B-80K and Llama-2-13B-60K [].
To examine the effect of alignment, we have study Mistral-7B-Instruct-v0.2 and Qwen-1.5-14B-Chat [] and compare them to their base versions.
We further choose Mixtral-8x7B-v0.1 [], a mixture of expert versions derived from Mistral-7B-v0.2, presumably via sparse upcycling [], to study retrieval heads in different architectures.

### 3.1 Universal and Sparse

Figure 3 demonstrate that a sparse set of retrieval heads exist in all models we consider, regardless of the various pertraining, fine-tuning recipes and the underlying architecture.
Between 25% and 52% of the heads exhibit copy-paste behaviors at a very low frequency, with a score of between 0 and 0.1.
Approximately 45% to 73% of attention heads have 0 retrieval score, meaning that they have other functionality than retrieval.
Only about 3% to 6% of the attention heads have a retrieval score larger than 0.1 (recall that a retrieval score 0.1 means a head retrieves at least 10% of the tokens being asked).
It is also intriguing that although all model parameters, particularly the total number to attention heads, are largely different from each other, their ratio of retrieval heads lays in the same interval (about 5%).

Figure: Figure 4: Retrieval score (blue): average number of activated tokens; Activation Frequency (green): frequency of at least activated on one token. The gap between the two curves shows context-sensitivity: a head of high activation frequency but low retrieval score means it is only activated on certain tokens and contexts; a head of high retrieval score means it is activated on almost any context. For both LLaMA and Yi, there exist strong heads that are not sensitive to context and always activated.

### 3.2 Dynamically Activated Based on Tokens and Contexts

Now we study how sensitive a retrieval head is to its input context, i.e.,
whether a head is consistently activated no matter what the context is, or if a head is activated only on specific contexts.
For the needle sentences "the best thing to do in San Francisco is eating a sandwich in Dolores park in a sunny day", some heads are activated on the full sentence, whereas other heads only activated on certain tokens like “eating a sandwich” or “in Dolores park’.
We define activation frequency, the frequency of a head being activated on at least one token (v.s., the retrieval score measures the average number of activated tokens).
A head of high activation frequency but low retrieval score means it is only activated on certain tokens and contexts.
As is shown in Fig. 4,
Llama-2-7B-80K and Yi-6B-200K have 12 and 36 strongest retrieval heads, respectively, that are always activated (activation frequency equal to 1) under all the contexts we consider.
Weaker heads only activate on certain tokens and contexts.

### 3.3 Intrinsic

Figure: Figure 5: Retrieval head is intrinsic and already within the base model. The subsequent model derivations, either by continue pretraining (LLaMA 2 7B 80K) or chat finetuning (Qwen 1.5 14B Chat) or sparse upcycling (Mixtral 8$\times$7B), use the same set of retrieval head as the base model, as demonstrated by a high level of similarity between the heatmap patterns.

We show that the retrieval heads, thus the ability of utilizing information at arbitrary location of the input, is an intrinsic property [] of the base model as a consequence of large-scale pretraining, with subsequent small-scale training exerting only minor alterations to these head activation patterns.
In Figure 5, we present the retrieval score distributions for a range of base models in the initial row, followed by their corresponding variants in the subsequent row. We see that regardless of the models being continuously pre-trained, chat fine-tuned, or sparsely upcycled, there is a notable consistency in their retrieval scores heatmaps.
Figure 7 offers a more direct and strict examination, where we compute the statistical correlations between different models.
The data reveal a high degree of correlation in the retrieval score distributions between base models and their respective variants, with a Pearson correlation coefficient exceeding 0.8.
Models from different families exhibit a correlation coefficient of less than 0.1, indicative of their distinct pretraining recipes.

Figure: Figure 6: The retrieval heads of models of the same family are strongly correlated, i.e., the chat model and base model typically uses the same set of retrieval heads. The retrieval heads of models of different families are clearly different.

## 4 Influence on Downstream Tasks

This section examines how retrieval heads influence downstream tasks.
Across the experiments we use Mistrial-7B-Instruct-v0.2 [] as it is a popular and strong open language model with 32K context length.
We first show that retrieval heads explains the factuality of Needle-in-a-Haystack test.
When the model can retrieve the needle, retrieval heads are always activated.
When the model cannot retrieve the needle and hallucinate instead, retrieval heads are either partially activated or not activated.
Then we show that retrieval heads significantly influence question answering that requires extracting the information from the input, but does not strongly influence tasks where the model directly produce answers based on its internal knowledge.
We further explore how retrieval heads influence more sophisticated reasoning behaviors like chain-of-thought [].

### 4.1 Retrieval Heads Explains Factuality in Needle-in-a-Haystack

We start with a closer look at Needle-in-a-Haystack and
construct an additional set of needle tests that are different from the three sets used for retrieval head detection.
We gradually mask out the number of retrieval/ random heads and see how the model’s behavior changes.
As is shown in Fig. 7, masking out retrieval heads severely damages the model’s Needle-in-a-Haystack performance, while masking out random heads shows much smaller performance impact.
Notably, when increasing the number of masked heads K to 50 (about 5% of the full number of heads), all models’ needle test performance drop to below 50, showing that the top retrieval heads are responsible for most of the needle retrieval behavior.

Figure: Figure 8: Masking out retrieval heads severely damages ExtractQA performance.

We observe three types of error: (1) Incomplete retrieval, where models only captured partial information of the target and omitted key details; (2) Hallucination, where models generated hallucinated sentences; and (3) Wrong extraction, where models incorrectly retrieved irrelevant content from the haystack.
Without any masking, instances of wrong extraction occurred when retrieval heads were active but focused on incorrect sections. During hallucinated generation, retrieval heads predominantly attended to the initial token of the input, which is often known as "attention sink" [], presumably dummy tokens that provide less meaningful signal.

As we increase the number of masked heads, initially,
a small subset of the most robust retrieval heads are masked, and incomplete retrieval began to appear.
In the absence of the strongest retrieval heads, the remaining weaker heads only managed to retrieve a fraction of the target information.
Metaphorically, each retrieval head holds a small piece of the "needle," yet these pieces cannot form a complete one, resulting in an incomplete final output.
This phenomenon typically begins when the mask out heads of retrieval score larger than 0.4.
As we further increase the number of mask,
hallucinations become more prevalent, signaling a complete failure of the retrieval capability.

### 4.2 Influence on Extractive QA

Now we study how retrieval heads influence more realistic tasks beyond Needle-in-a-Haystack.
We use extractive QA as a test bed, as common usecase of long-context model where the user typically upload a pdf (research papers, finantial reports, legal documents, .etc) and ask questions about specific information within the document.
To make sure the knowledge being asked does not exist in the model’s internal knowledge, we synthesize an extractive QA dataset by selecting a set of up-to-date news articles, extract a paragraph from it, and asking GPT-4 to produce a question-answer pair based on the extracted paragraph, similar to the evaluation conducted in Anthropic [].
As illustrated in Figure 9,
randomly masking out non-retrieval heads demonstrated no significant impact on performance.
Masking out retrieval heads led to a substantial decrease in F1 scores, with reductions of 9.2% and 23.1%.
These observations demonstrate that real-world document QA tasks heavily rely on the functionality of retrieval heads.

### 4.3 Chain-of-Thought Reasoning also Requires Retrieval Heads

Figure: Figure 10: Retrieval heads significantly influence tasks that require chain-of-thought reasoning. This is because typically in a reasoning chain, the next step reasoning requires the model to refer to previous information. See Fig. 11 for examples.

Figure: Figure 11: When we mask out retrieval heads, the model becomes “blind” to important information in the question description resulting in incorrect reasoning chains.

We test Mistrial-7B-Instruct-v0.2’s performance on MMLU [], MuSiQue and GSM8K [], with and without chain-of-thought reasoning.
As is shown in Fig. 10,
if we use answer-only prompting (without CoT), masking out either retrieval or random heads do not really influence the performance, presumably because the model’s generation is based on its internal knowledge primarily stored in the FFN layers [].
For CoT styled reasoning, masking out retrieval heads signifantly influence the model’s performance.
Upon inspecting typically error cases (Fig. 11), we find out the model becomes “blind” to important input information and hallucinate instead.
We find the relationship between CoT and retrieval heads particularly intriguing as it may offers deeper insights into model’s complex reasoning performance.
We leave more in-depth studies to future research.

## 5 Discussions

General Functionalities of Attention Heads  For transformer language models, we tend to view the functionality of FNNs layers to be the place for storing knowledge [], and the attention layers to be the place for implementing algorithms [].
The induction head discussed in  Olsson et al. [] typically searches repeated patterns of the input, which is at a certain level similar to the retrieval heads (as it also searches and repeats information).
Different than the induction heads, the retrieval heads are typically responsible for redirecting the information according to the context, but do not for inferring programs.
We tend to believe that there exist more algorithm and functionalities implemented by other types of attention heads to be discovered by future research.

Relationship to Local and Linear Attention and State-Space Models  Although there exist numerous works about local [] / linear [] attention, state space models [], and hybrid architectures [] achieving inspiring efficiency in long-context modeling, so far there is no linear attention / SSM architecture that passes the Needle-in-a-Haystack test to the best of our knowledge, suggesting that the full attention might be a must for long-context information retrieval.
One example is that the Mistral v0.1 [] uses sliding window attention but cannot pass needle-in-a-haystack, and their authors changes the attention to full in v0.2 [], then it can pass the needle test.
Our results showing strong evidence why full attention is a must.
For the model to precisely utilize input information at arbitrary location, it is crutial for the retrieval heads to work on the full KV cache.

Applications to KV Cache Compression  The problem that the KV cache is too large and occupies a large chunk of the GPU memory severely hinders the deployment of long-context models.
For example, for LLaMA 2 7B, the KV cache of 100K tokens requires more than 50GB memory, while 2K context requires less than 1GB memory.
If we serve this model on one 80G A100, then the concurrency of 100K context can be 50 times less than 2K context queries, which is prohibitively expensive.
The results from this work indicates that we might be possible to radically prune out the KV cache corresponding to the non-retrieval heads (recall in Fig. 3 shows only 5% of the heads are retrieval) and significantly reducing the deployment cost of long-context models.
We leave this study to future research.

## 6 Conclusions

This paper discovers retrieval heads, a special set of attention heads that are responsible for implementing the conditional copy algorithm and redirect information from the input to the output.
Retrieval heads are the primarily reason why a successful long-context model can pass the Needle-in-a-Haystack test, and their activation explains why a language model is faithful to the input or hallucinate.
Compared to non-retrieval heads, retrieval heads have a stronger influence on downstream tasks that require the model to precisely recall the input information, either in extractive question answering or chain-of-thought reasoning.
We believe this work will foster future research on reducing hallucination,
improving reasoning, and compressing the KV cache.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: Please check the list of minor points in "weaknesses" section.
soundness: 4
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.
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.
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.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 3
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.
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.
weaknesses: Nothing major needs to be addressed. Please address some discussion questions in question sections.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
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
soundness: 4
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.
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.
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.

Meta-review: 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.

Decision: Accept (Oral)