Retrieval Head Mechanistically Explains Long-Context Factuality
Authors: Wenhao Wu
Yizhong Wang
Guangxuan Xiao
Hao Peng
Yao Fu
Venue: ICLR 2025 Oral (2025)
Keywords: Large language models
long context
interpretability
attention
Primary area: foundation or frontier models, including LLMs
TL;DR: We study retrieval head, a special type of attention head that mechanistically explains long-context factuality

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.

$ University of Washington "MIT *UIUC #University of Edinburgh Published as a conference paper at ICLR 2025
RETRIEVAL HEAD MECHANISTICALLY EXPLAINS LONG-CONTEXT FACTUALITY
Wenhao Wu²
Yizhong Wang
Guangxuan Xiao
Hao Peng
Yao Fu"
à *Peking University
wa waynewu@pku.co edu. en hacpengg3i1immois yao. fußed. ac uk
ABSTRACT
Despite the recent progress in long-context language models, it remains elusive how transformer-buse 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 wide spectrum of models reveals that a special type of attention heads are largely responsible for retrieving informa- tion (either copy-paste paraphrase), which we dub retrieval heads. We identify intriguing properties retrieval heads:(1) universal: all the explored models with long-context capability have set of retrieval heads; (2) sparse: only 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, 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 serviously-generatee 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 to- kens. We believe our insights will foster future research on reducing hallucination, improving reasoning, and compressing the KV (Key-Value) cache.
INTRODUCTION
Transformer- bbased language models have demonstrated strong capabilities in processing long con- text (Anthropic) 2023 Reid et al., 2024 Fu al.) 2024), such as accurately retrieving relevant information from extremely-long and using it effectively in solving complex tasks (Kamradi) 2023 Hsieh et al. 2024). This capability lays the foundation for many other downstream tasks, which often require both retrieval relevant information and performing multistep reasoning (Kuratov et al. 2024).
We inquire: how do these models acquire such long- cong-context capabilities?
In this work, we investigate the internal mechanisms that allow large language models to leverage information from arbitrary positions within their input. Our comprehensive experiments, conducted across four model families, six model scales, and three types of post-training variants, reveal key insights into the internal mechanics of these models. We identify a special class of attention heads, which we term retrieval heads, as the primary contributors to retrieving relevant information from long contexts. Inspired by prior works like CopyNet (Gu et al. 2016) and Induction Heads (Olsson hypothesize that, similar to induction heads' role in in-context learning, retrieval
et al. (022),, heads are responsible for conditional information retrieval, executing a pastee or paraphrase algorithm that enables long-context understanding.

---

Published as a conference paper at ICLR 2025
Masking out top 20 retrieval heads, accuracy 63.6
Masking out 20 random heads, accuracy 94.7
0.1
0.5
5K
50K
100K
5K
50K
100K
Figure 1: Retrieval heads are specialized attention heads responsible for redirecting relevant infor- mation from the input to the output. Left: When the top retrieval heads in LLAMA 7B 80K are masked, the model's performance to hallucinations during decoding. Right: In contrast, masking random non-retrieval heads has no significant impact the model's ability retrieve correct information. Furthermore, retrieval heads primarily affect factuality, not language fluency. When masked, the model generates a fluent but incorrect sentence, such "go the beach," instead of the factual response "eat a sandwich Dolores Park."
the Needle-in-a-Hasttane
task deteriorates sharply, leading
To validate this hypothesis, we develop methods to detect retrieval heads in the transformer architec- ture (. and conduct extensive experiments that reveal four key properties of these heads (Sec (1) Universality and Sparsity: Retrieval heads are sparsely presented across all model families we tested, including LLAMA (Touvron al.) 2023). Yi Young et al., 2024), Qwen (Bai al. (023), and Mistral (Jiang et al. 2023), at various scales (6B, 14B, 34B, and 8x 7B). This universality holds for both base and chat models. whether dense mixture-of-experr (MoE). (2) Intrinsic Nature: These heads emerge naturally from large-scale pretraining. Even models like LLAMA 2, which have never been explicitly trained with long contexts, exhibit retrieval heads. Subsequent derivations, such as the long-context continual pretraining (LLAMA2 7B 80K), chat fine-tuning or RLHF (Qwen Chat), even sparse upcycling (Komatsuzaki et 2022; Jiang al. 2024) do modify the core retrieval heads but rather exploit the patterns already present from pretraining (Fig, 5); (3) Dynamic Activation: Retrieval heads activate based on context, with some heads consistently responsible for recalling specific information (e.g., head 19 in layer 16 in LLAMA A 2 7B), while others are selectively triggered by different content. This dynamic activation allows retrieval heads to complement one another; even when a subset of them is pruned, the model can still partially retrieve information. (4) Causality: Retrieval heads directly influence the model's capability to recall specific information. For instance, introducing the phrase "the best thing to do in San Francisco is to eat a sandwich in Dolores Park on a sunny day" and pruning all core retrieval heads causes the model to hallucinate things like "visiting the Golden Gate Bridge". Partial pruning results in partial recall (e.g., mentioning the sandwich but omitting Dolores Park). In contrast, pruning non-retrieval heads does not disrupt the retrieval of the full phrase. We further note that the functionality of retrieval heads goes beyond copy-paste. Specifically, we explore how retrieval heads contribute to chain-of thain-of-thought reasoning, process that inherently requires recalling earlier inputs. We find that retrieval heads play crucial role in redirecting the previous intermediate results to next reasoning steps, suggesting strong connection between memory retrieval and reasoning capabilities in large language models.
Retrieval heads are closely connected to induction heads in classical mechanistic interpretability literature (Bricken al., 2023: Olsson et al. 2022) as we both start from the copy-paste behavior (which further trace back to the earlier works like Gu et al. 2016). Yet our work may not be viewed as "rediscovery" of what already stated Olsson al. (202)), as we list the following notable differences: (1) scale: Olsson et al. (2022) focuses relatively small scale transformers (less than IB) while we consider a wide range of much larger models (from 7B to 34B); (2) focused capability: Olsson et al. (2022) focuses in-context learning, while most of our settings zero-shot. Our focuses are long-context factuality, question-answereee answering, and chain-of ofotthuugt reasoning - all of them important topics in today's context but not studied in Olsson et al. (2022).
We believe that the discovery of retrieval heads has profound implications for practical applications in long-context modeling: (1) it explains why certain compext-compressionn methods fail to maintain
2

---

Published as a conference paper at ICLR 2025
Attention distribution of ratrieval
What the the best thing
A: eall
Once upan a time
The best thing to do San Francisco is to ant sandwich
live happlly aver aftor
100K context
Neecle sentence
100K context
Figure An attention head is considered perform a copy-paste operation when the token it attends to matches the token being generated. The retrieval score for head is defined as the frequency of this copy-paste behavior during tasks that require retrieving raw information from the input.
Table 1:
We consider 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-80K
Length Extension via Continue Pretrain
-lama-2-7B
Llama-2-13B-
Model Scaling and Length Extension
SFT and RLHF
Mistral-7B-v0.2
Mixtral-8x7B-v00
Sparse Upeycling to Mixture of Experts
Yi-6B-2000
Length Extension via Continual Pretraining
Yi-6B
Yi-34B-200K
Model Scaling and Length Extension
Qwen1.5-144
Qwen1.5-14B- 15-14B-Chat
SFT and RL.FF
factual accuracy, as they completely remove the KV cache corresponding to the retrieval heads (iiao et al. )023); (2) suggests that future research on KV cache compression, a critical issue for deploying long-conexx: models, should seek sparsity in the head dimension instead of the token dimension. That
say, instead of pruning out the entire KV cache for less important tokens (Ge et al. 2023 Kang et al. (2024), one may seek pruning out the KV cache for less important heads (while keep all the tokens).
DETECTING RETRIEVAL HEAD
To identify which attention head is implementing the retrieval mechanism, we introduce a retrieval score, which measures the frequency of a head's copy-paste behavior during autoregressive decoding. An attention head with a high retrieval score indicates that, across various contexts, the head frequently copies tokens from the input to the output.
Needle-in-a- Haystakk (NIAH) Our retrieval head detection algorithm roots from the Needle-in-a- Haystack test (NIAH), which asks the model to copy-paste the input tokens to the output. Given question q and its corresponding answer (the "needle"), we insert & into a context x (the "haystack") randomly chosen position indexed by ig. The language model is tasked with answering q based on the haystack with the inserted needle. We make sure that is sufficiently unique so that it can only be resolved by referring the content within k. and not by drawing upon other content in or the model's parametric knowledge. This is to say, the needle is semantically irrelevant to the context æ (see Fig. 2 as an example). Token-level recall is applied to evaluate whether the model retrieves the answer by covering the salient information in k. The final NIAH performance score for a given test is reported as the model's success rate in retrieving k.
Retrieval Score for Attention Heads copy-paste operations. Specifically, during regressegees decoding (we use greedy decoding by default), denote the current token being generated and the attention scores of head IR [R
We define the retrieval score as the frequency of head's

---

Published as a conference paper at ICLR 2025
As demonstrated in Fig. 2 say attention head á copy-paste operation for token from the needle if it satisfies two conditions::
1. Token Inclusion:
The token w must belong the needle sentence, i.e., w E
2. Maximal Attention: The token must correspond the input position j. i.e., æjj w where j arg max(a) and j E dq. meaning that the highest attention score aligns with w.
Let 9h represent the set of tokens copied from the needle according to the criteria above. The retrieval score for head h is then defined as:
Retrieval score for head h 1g l/A/]
Intuitively, retrieval score represents a token-level recall rate of the most attended tokens by an attention head. For instance, in retrieving a needle of 10 tokens, a retrieval score of 0.9 means the attention head correctly copied 9 of the 10 tokens, suggesting that the head is specialized in retrieving information from long contexts. We further note that although we detect retrieval heads by copy, in practice, their functionality goes beyond copy-paste, as we observe that they are activated during paraphrasing, ceestion-ansseringg and chain-of-thougee reasoning.
Retrieval Head Detection Algorithm We compute the retrieval score for all attention heads across a broad range of test cases, We construct three sets of NIAH samples. Each sample is defined as a unique tuple (q, k 2), where (g. k) is query-key pair that is intentionally designed to be semantically irrelevant to æ (the context). We manually verify that cannot be answered from the model's prior knowledge alone, ensuring that retrieval relies solely on the context æ. For each (q. k, æ) sample, we conduct the NIAH test by evaluating the model's behavior over 20 different sequence lengths uniformly sampled between IK and 50K tokens. At each length, q inserted at 10 evenly distributed positions, from the start to the end of æ. This allows us to evaluate the model's retrieval capabilities at varying depths and in diverse contexts.
Our experiments show that the retrieval score stabilizes quickly, often converging after just a few samples. In total, each model undergoes approximately 600 retrieval testing instances. For each test, compute the retrieval score for every attention head, then average these scores to obtain final retrieval score for each head. To identify retrieval heads, we apply a threshold criterion. In retrieval head if its average retrieval score exceeds
our experiments (Fig. 3 a head is classified 0.1, meaning that it successfully performs a copy-paste operation in at least 10% of the test cases. This threshold reflects the minimal level of retrieval activity necessary for head to be considered specialized for retrieval tasks.
3 BASIC PROPERTIES OF RETRIEVAL HEADS
This section discusses important properties of retrieval heads discovered from retrieval head detection algorithm. Our results are supported by extensive experiments on a large spectrum of models (Table []. To investigate the influence of continued pretraining for context length extension, we compare LLA A-A-M----BB 4K to examine the effect of alignment, we have study Mistral-BB-Instrucce
and LLAMA---115 13B-60K (Fu et al. 2024). To
and Qwen- 1.veen-.5-14B-. T.5-1BBBhhh (Bai et al. 2023) and compare them to their base versions. We further choose Mixtral-8x7B-v0l (Jiang al. 2024), mixture of expert versions derived from Mistral-7B-vol presumably via sparse upeycling (Komatsuzaki et al. 2022), to examine the behavior of retrieval heads in distinct architectures.
Universality and Sparsity Figure 3 highlights the presence of sparse set of retrieval heads across all models studied, regardless of variations pretraining or fine-tuning methods, as well as underlying architectural differences. Between 25% and 52% of attention heads exhibit copy-paste behavior at low frequencies, with retrieval scores between and O. Additionally, approximately 45% to 73% of attention heads have retrieval score of 0, indicating that they serve functions other than retrieval. Notably, only about 3% to 6% of attention heads achieve a retrieval score above 0.1, meaning they retrieve least 10% of the target tokens. Of these, only 0.1 % to 0.8% of heads exhibit a retrieval score higher than 0.5, indicating frequent engagement in copy-paste operations. Interestingly, despite the substantial variation in model size and the total number of attention heads, the proportion of
4

---

Published as a conference paper at ICLR 2025
Llama-2-78-80K
YI-6B-200K
Qwen1.5-14B
Mistral-7B-vo.2
75.2%
50.7%
53.5%
05
0.5%
4,5%
3.2%
25.8%
43.7%
02.5%
50.8%
32 x 32 heads
32 32 heads
40 x 40 heads
32 x32 heads
Llama2-13B-64K
Yi-34B-200K
Qwen1.5-14B-Chat
Mixtral-6x7B-v0..1
15.2%
71.15
04%
0.1%
3.1%
4.2%
MI%
$1.8%
D-0.1
0.1-0.5
0.5 1
40 heads
60 heads
40 heads
heads
Figure 3: Distribution of retrieval scores for different models. Blue means weaker retrieval scores where heads do not retrieve, and red means heads that retrieve. In all models analyzed, fewer than 1% of attention heads are activated more than 50% of the time, with retrieval score exceeding 0.5, when retrieval tasks are required. For example, for Llama-2-7B-800 70% of the heads never do retrieval, and only about (2.6 + 0.6)% heads retrieve more than 50% of the time, meaning that retrieval heads are sparse within the architecture.
1amma---8-80K
Yi-6B-200K
1.0
0.
0.8
0.6
0.0
174 142 276
20.0
18
13.10
10:2
0.4
64
0.2
Retrieval Score Activated Prequency
Layer-Head
LayerHead
Figure 4: Retrieval Score (blue): Represents the average portion of tokens that are activated across dif- ferent contexts. Activation Frequency (green): Indicates the proportion of instances where least one token is activated. The divergence between the blue and green curves illustrates the sensiiiiey of the model's heads. A significant gap suggests that a particular head is activated frequently but only under specific token and contextual conditions, indicating high context-dependence. Conversely, small gap or overlap indicates a head with a broad activation pattern, suggesting low context- sensitivity. Both LLAMA and Yi models exhibit heads that are consistently activated across various contexts, demonstrating robustness their activation patterns.
retrieval heads remains consistent across models, hovering around 5%. We further note that the sparsity ratio may depend on the task: for tasks that are -haayy, one may expect a higher level sparsity, as what we see here. Yet for tasks that may heavily involve the model's internal knowledge and reasoning, one may not necessarily observe the same level of high sparsity (e.g., Gi et al. 2023 observes about 50% sparsity on general chat).
Dynamically Activated Based on Tokens and Contexts We next explore the sensitivity of retrieval heads to input
they are consistently activated across contexts or only in response
5

---

Published as a conference paper at ICLR 2025
Liame 2-78
Qwen1-5-148
Nintral-7B.v0.22
Layer
Unia 2-78- 20k
Qpen11--18B--hae
Mixtral-bx7--v...
Head
Figure 5: Distribution of retrieval heads within the model architecture across layers. Red and orange squares denotes retrieval heads. The retrieval head is intrinsic to the base model and remains consistent across model variants, including continued pretraining (LLaMA 2 7B 80K). chat fine-tuning (Qwen 1.5 14B Chat), and sparse upcycling (Mixtral 8x7B). This is evidenced by the high similarity in heatmap patterns, indicating that the retrieval mechanism is preserved across these transformations.
to specific content. For instance, in the sentence "the best thing to do San Francisco is eating a sandwich in Dolores park sunny day," certain heads activated across the entire sentence, while others focus only on specific phrases such as "eating a sandwich" or "in Dolores park." To capture this behavior, we define activation frequency- the frequency with which head is activated on least one token (as opposed to the retrieval score, which measures the average number of tokens activated). A head with high activation frequency but low retrieval score indicates selective activation based specific tokens and contexts. As shown in Fig. 4 LLAM M---MAA-----00 and Yi-6B-200K respectively have 12 and 36 strongest dedicated retrieval heads that are always activated (activation frequency equal to 1) under all the contexts we consider. Less dedicated retrieval heads only activate on certain tokens and contexts.
Intrinsic Nature We find that retrieval heads-and the capacity to retrieve information from arbitrary positions within the input- are intrinsic properties of base models, emerging naturally during large-scale pretraining (Fu et al.] 2024). These heads exist even in models that have not been explicitly trained on long-context tasks, with task-specific fine-tuning leading to only minimal changes in their activation patterns. In Figure we visualize the distribution of retrieval scores across a range of base models (first row) and their corresponding variants (second row). The heatmaps reveal striking consistency in retrieval patterns, regardless of continued pretraining, chat fine-tuning, or sparse upcycling. Figure O further supports this observation, showing Spearman correlations between the retrieval scores of different models. Base models and their fine-tuned counterparts exhibit strong positive correlations (with Pearson coefficients greater than 0,8), while models from different families display much weaker correlations (less than 0. I), reflecting their distinct pretraining methods.
INFLUENCE ON DOWNSTREAM TASKS
This section analyzes the impact of retrieval heads on downstream tasks, focusing on experiments conducted using -istral--B-IIsstruut--dd 7B--Instruct-v0.2 (Mistral 2024). Specifically, retrieval heads are consistently activated when the model retrieves the "needle." In contrast, when the model fails to retrieve the needle and hallucinates, the retrieval heads are either only partially activated remain inactive. We then show that retrieval heads significantly affect extractive question-answerr information extraction from the input, but have less influence tasks where the model generates answers based on its internal knowledge. Finally, we explore how retrieval heads contribute to more sophisticated reasoning behaviors, such thain-of-thought reasoning (Wei al. 2022
tasks that require

---

Published as a conference paper at ICLR 2025
100
7.8
9,922
2,00
2M22
80-
Average Score
Mers 3Bx02
002007
0.120
0.644
Uanu 2-b--tk
10:21
0.00121
0.022
0.0005
7.4
Qwm1).: 5:40-320 Che
3.222
3.43
3.11
3,023
0:003
(estral--x7--00.. Yn 6B-200K
20
Oven 140
Nooz SxTS.0
Random Head
Masking
Figure 7: NIAH scores when masking out top- K retrieval heads versus K random heads: For all models considered, removing retrieval heads significantly degrades formance. In contrast, removing equal number of non-retrieval heads has a much smaller impact.
Figure 6: The retrieval heads of models of the same family are strongly correlated, i.c., the chat and base model typically utilizing the same set of retrieval heads. contrast, retrieval heads across different model families show clear distinctions.
Gosdle-in-a-Haysteee
per-
ExtractQA
Incomplete
56.7
55.7
55.4
Go to Dolorss park on a sunny day
Golden Gate Bridge
Attention top Retrieval Heads:
Attention Retrieval Heads:
47.5
No Masking Retrieval Head
16-18
11-15
17-22
19-15
16-19
17-22
Random head
Go to Dalores park on sunny day and eat sandwich [Begin of Input) Once upon time...
32.3
Haystack
50
100
Mask Num.
Two types of typical errors when the model fails to retrieve needles: (1) Incomplete retrieval, where the retrieval heads miss part of the information "eat sandwich": (2) Hallucination, where the retrieval heads attend to the initial tokens.
Figure 9:
Figure 8: Masking out re- trieval heads severely damages ExtractQA performance.
4.1 RETRIEVAL HEADS AND FACTUALITY IN NEEDLE-IN-A-HASSIACE
We begin with detailed investigation of the NIAH test, constructing additional evaluation tests using (q. kt x)) tuples. We gradually prune retrieval heads and observe the resulting performance changes. The pruning strategy follows higher lower retrieval score orders. Specifically, we prune the top-K retrieval heads by masking out attention heads with the highest retrieval scores. As baseline, we compare this the pruning of equal number of random attention heads. As shown in Fig. 7 masking out retrieval heads severely impacts NIAH performance, whereas pruning random heads has far less effect. Notably, pruning more than 50 retrieval heads approximaaax 5% of all attention heaus-results in performance dropping below 50%, indicating that the top retrieval heads play critical role in needle retrieval.
We identify three types of errors: (1) incomplete retrieval, (left in Fig. 9), where the model retrieves only part the required information, omitting crucial details; (2) Hallucination (right in Fig. 9 where the model generates fabricated information: and (3) Wrong extraction, where irrelevant content retrieved from the haystack. Without masking, wrong extractions occur when retrieval heads focus on incorrect sections. In cases of hallucination, retrieval heads tend to attend primarily to the input's initial tokens, often termed an "attention sink" (Xiao et al. 2023). which contributes little to the final output.
7

---

Published as a conference paper at ICLR 2025
GSMBK (w/ COT)
MMLU (w/ COT)
MuSiQue (w/ COT)
30
29.3
22.1
10
9:3
11.4
2.0
100
100
100
Mask Num.
Mask Num.
Mask Num.
GSMBK (w/o COT)
MMLU (w/o COT)
MuSiQue (w/o COT)
30
40
No Masking
25
20
60
55.6.
57.257.0
54:2
54.0
Masking Retrieval
Masking Ra Random
50
15
Acc
16.0 17.16
10
40
12.5
5-
2.5
2.3
2.3
5
0
100
100
100
Mask Num.
Mask Num.
Mask Num.
Figure 10: Retrieval heads significantly influence tasks that require chain-of-thought reasoning. This because typically in a reasoning chain, the next step reasoning requires the model to refer to previous information. See Fig. II for examples.
As we increase the number of masked heads, incomplete retrievals emerge. This occurs because, without the most effective retrieval heads, the remaining weaker heads retrieve only partial information. This effect typically begins when retrieval heads with scores greater than 0.4 are masked. As masking continues, hallucinations become more frequent, ultimately leading complete retrieval failures. Intuitively, each retrieval head holds a small piece the "needle," yet these pieces cannot form a complete one, resulting in partial retrievals. 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 frequent, leading complete failures of retrievals.
4.2 IMPACT EXTRACTIVE QUESTION ANSWERING
Next, we examine how retrieval heads affect other downstream tasks, focusing extractive QA, common use case for long-context models where users input large documents (e.g., research papers, financial reports, legal documents) and pose questions requiring information extraction.
To ensure the relevant knowledge is absent from the model's internal parametric knowledge, we construct an extractive QA dataset using recent news articles. We extract paragraphs from these articles and have GPTT-4 generate corresponding question-answer pairs. This methodology mirrors the approach of Anthropii (2023). As illustrated in Figure randomly masking out non-retrieval heads demonstrates no significant impact on the models' performance. However, masking out retrieval heads leads substantial decrease in Fl scores, with reductions of 9.2% and 23.1 %. These observations demonstrate that retrieval heads are crucial for real-world long-context QA tasks.
4.3 SHAIN-FF-HHUGGH REASONING ALSO REQUIRES RETRIEVAL HEADS
To examine how retrieval heads effect reasoning tasks, we test Mistrial- 7B- 77--Insruct--0.. on MMLU (Hendrycks al. 2020), MuSiQue and GSM8K (Cobbe al., 2021). with and with- out oh (CoT) reasoning. MMLU primarily assesses a model's parametric knowledge and requires minimal reasoning, thus offering limited benefits from CoT reasoning. In contrast,
8

---

Published as a conference paper at ICLR 2025
Question 1:
Question 2:
Mishia bought pairs shorts, of pants, One pair of shorts
deckdes sprints times He meters nach sprint. many total meters does he run week
$16.50. One pair pants costs $22.50 and one
pair of shoes costs $42.
Masking Top 100 Retrieval Heads:
How many dollans did Mishka spend on all the clothing heem?
Let's think through this step by stop:
Masking Top 100 Retrieval Heads:
James runs 3 times a week for each sprint session.
Lot's calculate the cost of each type of clothing for Mishkaa:
During each session, covers distance meters (3 times the longth of standard football field gool post).
1. For ehorts, there are 3 pairs, so the total cost $16.5 $69.10. For pents, the
same, 3. In total, Mishka spent $39 a 568 $125.
pairs cost $222 $66.
Therefore, the total distance James covers in week is: session)
(sessions week) x1 (number
So, the answer $125.
60jmeters in a sprintj= 360 meters.
Misunderstood Information
So. James runs a total of 380 meters every week.
When we mask out retrieval heads, the model ignores important information in the question description resulting in incorrect reasoning chains.
Figure 11:
both MuSiQue (multi-hop QA) and GSM8K (math problem-solving) demand complex, multi-step reasoning, where CoT prompting has been shown to substantially enhance performance.
As illustrated in Figure T0 using an answer-only prompt without CoT, the model's performance remains largely unaffected by masking either retrieval or random heads. This suggests that, in these cases, the model's generation relies primarily its internal parametric knowledge, likely stored in the feed-forward network (FFN) layers, as proposed by Geva et al. (2020). However, when using CoT reasoning, masking retrieval heads significantly degrades performance. A closer examination of common failure cases (Figure [T] reveals that when retrieval heads are masked, the model often fails to fully comprehend key input details, leading to hallucinations. CoT reasoning. which involves decomposing complex tasks into smaller steps, heavily depends on accurately retrieving detailed information from the input. Without effective retrieval, the model "loses sight" of important conditions, resulting in flawed reasoning. For example, in the left of For example, in the case shown on the left side of Figure [T the model fails to retrieve the input condition about the costs of pants and shoes, instead fabricating the values during CoT reasoning. Similarly, in the case on the right, the model misses the condition "3 sprints. 3 times a week," " and hallucinates new rules to calculate the total distance. These findings highlight the critical role retrieval heads play in enabling effective CoT reasoning. We believe that further in-depth exploration of this relationship could offer significant insights into the mechanisms underpinning language models' reasoning capabilities. However, leave these broader investigations to future work.
DISCUSSIONS
General Functionalities of Attention Heads For transformer language models, FFN layers are generally understood to store knowledge, as suggested by Geva al. (2020), while attention layers implement dynamic algorithms (Olsson et al. (2022). The induction heads introduced in Olsson et al. (2022) search for repeated patterns in the input, which bears some similarity the role of retrieval heads, as both mechanisms involve retrieving and repeating information from the context. However, unlike induction heads, retrieval heads focus on redirecting information based on the context without directly executing inference programs. We believe that future research will uncover additional functionalities and algorithms implemented by other types of attention heads, further expanding our understanding of transformers' internal mechanisms. We tend to believe that there exist more algorithm and functionalitics 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 (Xiao et al. (023) / linear ( (Wang et al.) 2020) attention, state space models (Gu & Dao 2023), and hybrid architectures (De et al. (024) achieving inspiring efficiency in long-context modeling, many of these architectures, despite their efficiency, perform poorly on tasks requiring long-context understanding, such as the NIAH test. For instance, Mistral v0.1 et al. 2023) implemented sliding window attention, which failed to pass the NIAH test. However, when the authors switched to full attention in Mistral v0.2 (Mistral 2024), the model successfully passed the test. Our findings provide compelling evidence that full attention is crucial for effective long-context

---

Published as a conference paper at ICLR 2025
information retrieval. Specifically, retrieval heads rely access to the entire key-value (KV) cache to precisely utilize input information from arbitrary positions. Without full attention, retrieval heads lose the capacity to fully retrieve contextually relevant information, leading performance degradation in complex tasks requiring fine-grained information retrieval.
Applications KV Cache Compression major challenge deploying long-context models is the significant memory overhead caused by the large KV cache. For example. LLAMA 2 7B requires more than 50GB of memory maintain a 100K-token KV cache, compared to less than IGB for 2K context. This discrepancy drastically reduces the concurrency of 100K-token queries, making deployment on systems like an 80GB A100 GPU prohibitively expensive. Our findings indicate that it may be possible to prune KV cache entries associated with non-retrieval heads, as Figure 3 demonstrates that only 5% of the attention heads function as retrieval heads. This could significantly lower the deployment costs of long-context models. We leave further exploration of KV cache compression for future work.
CONCLUSIONS
This paper discovers retrieval heads, a special set of attention heads that are responsible for im- plementing the conditional copy algorithm and redirect information from the input the output. Retrieval heads are the primarily reason why successful long-context model can pass the NIAH test, and their activation explains why language model is faithful the input or hallucinate. Compared to non-retrieval heads, retrieval heads have stronger influence on downstream tasks that require the model precisely recall the input information, either in extractive question answering or chain-of- thought reasoning. We hope this work foster future research on reducing hallucination, improving reasoning, and compressing the KV cache.
10

---

Published as a conference paper at ICLR 2025
7 APPENDIX
7.1 RESULTS ON JAMBA
100
80
jamba v0.
60
Liama-2: 2-BB-OKK
Average
68-6000K
Retneval Head
Random Head
3-
02. 10
100
Masking Heads
Figure 12:
Masking out retrieval heads in Jamba, a hybrid state-space and attention model, signifi- cantly reduces its NIAH performance than masking random heads.
7.2 DETAILS ON RETRIEVAL HEADS DETECTION
Ranking Changes Over Iterations
16-19
Retrieval Score Rankings
4-
8:26
6-9
7-12
17:22
11-2
19-15
21-30
10
20
30
50
Detection Iteration
Figure 13: Illustrate how Ranking of top retrieval heads of Llama2-7B-80K changes when detection iteration increase.
From the figure above, we observe that at the initial stages of retrieval head detection, the rankings of the top retrieval heads (ranked by the current average retrieval score) fluctuate significantly. However, as the number of detection iterations increases, the rankings of most heads stabilize.
REFERENCES
Anthropic. Model card and evaluations for claude models, July 2023. URL https: www
anthropic wwtp:roductttctcccoooooooooocco
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023,
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly. Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec,
11

---

Published as a conference paper at ICLR 2025
0.0%
0.0%
=
0-0.1
0.1-0.5
0.5-1.0
0.7
10 needles
100 needles
Figure 14: Increasing the needle used for retrieval head detection from 10 to 100 does not change the detection results. This is to say, small set of needle is sufficient for identifying retrieval heads. Experiments performed on Mistral 7B. Also note the figure here aligns with the distribution in Fig.
Nicholas Schiefer, Tim Maxwell, Nicholas Joseph. Zac Satfield-Dodde Alex Tamkin, Karina Nguyen, Brayden Mel McLean, Josiah Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. features/index.html
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichhro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian- Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 2904229722999.... 2024.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context. arXiv preprint arXiv:2402.10177, 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310001801, 2023.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint 2rXh::2012.14913 2020.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint ar arXiv:2312.007522 2023.
Jiatao Gu, Zhengdong Lu, Hang Li, and Victor O.K. Li. Incorporating copying mechanism in sequence-to-sequence learning. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 1631-1640, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10.18653/v1/P16- 1154. URL / aereccaanhht ant wooology.ooo org..ogg.o...oooooo.o 61544
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2025
Paraphrasing
Input: The Whispering Forest of Lunthera, home bioluminescent insects, is famed for its murmuring trees and ancient legend about lost travelers finding their way home.
Mask roo retrieval heads; Found in the Lunthara Woods, murmurous shapes say colossal tales and have left their mark in historical memory.
Mask 10o random heads: Lunthera, situated among whispering trees, is widely renowned for the legends of lost travelers finding their way home.
Paraphrasing
Input: The crystal-clear waters of Virelia's Crescent Lake, surrounded by silver-barked trees, said to heal wounds and reflect the constellations even during the day.
Mask 10o retrieval heads; Vire's Crest, a silver hued cascade, cascades over a plateau that dangles from the sides of Vanderite's plateaus and overhangs in the valley.
Mask 100 random heads; The limpid waters of Virelia's Crescent Lake, bathed suee trees, are believed to help wounds and reflect the constellations even during the day.
Question Answering
Input: Here are some interesting facts about the world: The towering cliffs of Eldoria, reaching over 2,000 feet, are known for their shimmering crimson hues sunset and the rare Starfire flowers that grow their edges Q: what is towering cliffs of Eldoria famous for?
Mask roo retrieval heads; The brilliant flowers that cover cliffs are harvested and used in traditional healing potions.
Mask 100 random heads: Their shimmering crimson hues at sunset.
Question Answering
Here are some interesting facts about the world: The glowing sands Duskveil Desert, enriched with rare minerals, shimmer in the dark and are said to hold heat for days after the sun sets Q: can you tell me the something about Duskveil Desert?
Input:
Mask retrieval heads; Many gnomes in the forests surrounding Ealdwalleria believe that the sand Derva cavern region origin a key to D Vixen.
Mask too random heads; Glowing sands, heat conservation, and rare minerals.
Figure 15: Examples of masking out retrieval heads for paraphrasing and question answering
Cheng-Ping Hsich, Simeng Sun. Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language
13

---

Published as a conference paper at ICLR 2025
models?, 2024. BBL//tttss//aaaxx..rr///////22.0
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825. 2023.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint XXiv:2441.00088 2024.
Greg Kamradt. Needle in a haystack pressure testing Ilms. https://ittub.com
gkamradt Ltaaattttttteteeteeeeee Www.NeeeleInAHHaaaaaaa InAHaystackk 2023.
Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo Zhao. Gear: An efficient kv cache compression recipefor near-lossless generative inference of Ilm. arXiv preprint arXiv:2403.05522 2024.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of- experts from dense checkpoints. arXiv preprint arXiv:2212.05055. 2022.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. In search of needles in 10m haystack: Recurrent memory finds what Ilms miss. arXiv preprint M.Xiv:2402.10790.. 2024.
Mistral. Model card for misttal--7b-instrutt--0222 April 2024. URL https://huggingfaae
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209 1895, 2022.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint 2rXiv:2403.05530 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09280 2023.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04760 2020.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoe V. Le, and Denny Zhou, Chain-of- prompting elicits reasoning in large language models. In Sanmi Koyejo. S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http: http::://appp.eeee cetpp//er.r..i.ps.ppppee.eeeee wcwppper.ppppppppeeeeeeee files lates/ppppee http//ppprr///2/ (dd560963522444
Cow.:eacctt--oooeeneee.... html
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. /0099/77444 /7453, 2023.
Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Heng Li, Jiangeheng Zhu, Jianqun Chen, Jing Chang, et al. Yi: Open foundation models by 01, ai. arXiv preprint arXiv:2403. 04652, 2024.
14


Reviews:
[{"review_id": "r93bv6e3Jw", "reviewer": "Reviewer_Br4p", "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\n- (1) exist in all the explored models,\n- (2) are only 5% of the attention heads,\n- (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,\n- (4) are dynamically activated given different contexts, and\n- (5) will cause degradation in retrieval abilities or chain-of-thought abilities if pruned.", "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.\n- 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%.\n- 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?\n- L203: What do you mean by “stabilize” and “converge”? Please either provide definitions or plots to illustrate.\n- Figure 7: Could be nice to include the dataset name in the caption.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n- 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.\n- 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.\n- 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.\n- 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?\n- The paper misses experimental details, such as prompts used, links to existing assets used, etc.\n- The questions below need to be addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DNCrZ4uhx1", "reviewer": "Reviewer_ZPPp", "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: \n1) retrieval heads are present across a diverse set of models,\n2) only a small subset of attention heads function as retrieval heads,\n3) these retrieval heads exist even in models pretrained with limited context, suggesting they are an intrinsic artifact of pretraining,\n4) they are dynamically activated rather than continuously active, and\n5) there is a causal link between retrieval heads and the model’s capability to retrieve relevant information.", "questions": "Please check the list of minor points in \"weaknesses\" section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 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.", "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.\n\n\nAdditional minor points:\n1) 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.\n2) 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.\n3) 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.\n4) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mWdKPkQwfZ", "reviewer": "Reviewer_F46b", "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: \n\n1. **Universality** – retrieval heads are present across all explored models.\n2. **Sparsity** – only a small subset of attention heads serve this retrieval function.\n3. **Intrinsic nature** – these heads exist in pretrained models, even those trained on short contexts.\n4. **Dynamic activation** – their activation varies depending on specific tokens and contexts.\n5. **Causality** – pruning these heads leads to significant performance drops.\n\nBy examining the influence of retrieval heads across various tasks, these findings shed light on which internal model components actively seek information from input tokens.", "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?\n\n2. 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?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Provides in-depth analysis and extensive experiments on various properties of retrieval heads.\n\n2. Well-written, with clear graphics, easy-to-follow explanations, and a well-organized structure.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3z2EneKbdE", "reviewer": "Reviewer_xfgX", "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.", "questions": "- It would be better to show that similar findings transfer to paraphrasing tasks. \n- More appendix figures for evaluation results of Retrieval Head Detection Algorithm on various settings (line 199)\n\n- Minor fixes\nline 247-250 repetition\nFigure 9 caption typo: needels", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper shows originality by exploring the retrieval capabilities of transformer language models, a topic that has not been extensively studied.\n- 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.\n- Figures are clear and intuitive to comprehend  \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)