PaperlayerView exact revision on arXiv

Source-linked AI summary

Retrieval Head Mechanistically Explains Long-Context Factuality

Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, Yao Fu

arXiv:2404.15574v1cs.CL

TL;DR

How long-context transformers retrieve relevant information from arbitrary input locations remains unclear, despite retrieval supporting more advanced tasks that combine retrieval and reasoning. This paper detects and studies retrieval heads, finding that their activation tracks faithful input retrieval and that they more strongly influence tasks requiring input recall than answers based on internal knowledge.

  • Problem

    How transformer-based long-context models retrieve relevant information from arbitrary input locations remains unclear, despite this capability underpinning more advanced tasks that interleave retrieval and reasoning.

    Source: §1 · p. 1

  • Method

    The paper detects heads using a score for copy-paste behavior during autoregressive decoding, then tests their properties across model families, scales, and finetuning settings.

    Source: §1 · p. 2, §2 · p. 3

  • Results

    Retrieval-head activation tracks faithful input retrieval versus hallucination, and these heads more strongly influence input-recall tasks than answers based on internal knowledge.

    Source: §1 · p. 2, §4 · p. 6, §6 · p. 9

  • Takeaways & Limitations

    The authors argue that future KV-cache compression research should account for retrieval heads' influence.

    Source: §1 · p. 2

    The paper leaves more in-depth studies of how retrieval heads relate to complex reasoning performance to future research.

    Source: §4.3 · p. 7

Abstract

from arXiv · show

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.

1 Introduction

The paper identifies sparse retrieval heads that locate requested information in long contexts and redirect relevant input tokens to the output. These heads help explain when models remain faithful to the input or hallucinate, and matter for reasoning and context compression.

Introduction

  • The authors frame retrieval heads as a conditional copy-paste mechanism, drawing on CopyNet [10] and Induction Heads [19] to motivate their hypothesis.They propose that specialized heads implement information retrieval in a way analogous to induction heads supporting in-context learning.Source: §1 · p. 2
  • Retrieval heads are universal, sparse, and intrinsic: models across families and configurations use a small set that persists through continued pretraining, chat fine-tuning, and sparse upcycling.The base model already contains these heads, and later derivations use the same ones.Source: §1 · p. 2
  • Masking retrieval heads can cause hallucination, while masking random non-retrieval heads leaves retrieval intact; partial masking can preserve only part of the requested information.The heads redirect relevant input tokens to the output; when they are not activated, or are masked, the model may fail to find the information and hallucinate.Source: §1 · p. 2, Figure 1 · p. 2
  • Retrieval heads activate dynamically: the strongest heads remain active across different requested information, while weaker heads activate in different contexts and can compensate for one another.The introduction gives 13 as an example of the strongest heads in LLaMA 2 7B.Source: §1 · p. 2
  • Chain-of-thought reasoning relies heavily on retrieval heads because it requires referring back to input information, whereas answering from intrinsic knowledge is less affected by masking them.The reduced impact on intrinsic-knowledge answers is stated in the paper context.Source: §1 · p. 2
  • Identifying retrieval heads pinpoints a subnet for conditional retrieval and suggests context-compression methods may harm factuality when they remove these heads.The authors argue that KV-cache compression research should account for retrieval heads.Source: §1 · p. 2

2 Detecting Retrieval Head

The section detects retrieval heads by measuring how often they copy tokens from a context into generated answers. It defines a retrieval score through needle-in-a-haystack tests and averages scores across varied contexts.

Needle-in-a-Haystack

  • Needle-in-a-haystack tests insert a unique answer at a random context position and ask a question designed to require copying it from the context.The question and answer are chosen to be semantically irrelevant to the context, and the answer cannot be supplied from the model’s existing knowledge.Source: §2 · p. 3

Retrieval Score for Attention Heads

  • A head copies a needle token when it generates that token and its highest-attended input token is the same token within the needle.The retrieval score is the frequency of such copy-paste operations; it acts as a token-level recall rate, so 0.9 on a 10-token needle means 9 tokens copied.Source: §2 · p. 3

Retrieval Head Detection Algorithm

  • Across explored models, fewer than 5% of attention heads have retrieval scores above 0.5 more than half the time retrieval is required.The figure characterizes retrieval heads as universal across the explored model family and scale, yet sparse.Source: Figure 3 · p. 4
  • Each model is tested on approximately 600 instances across 20 context lengths from 1K–50K and 10 insertion depths per length; head scores are averaged across tests.A head is classified as a retrieval head when its average retrieval score exceeds the 0.1 threshold.Source: §2 · p. 4

3 Basic Properties of Retrieval Heads

Across the models studied, retrieval is performed by a small set of heads, whose activation varies with context and whose patterns are largely present in base models and persist through subsequent training.

3.1 Universal and Sparse

  • Only about 3% to 6% of attention heads have a retrieval score larger than 0.1 across the models studied, with retrieval-head ratios clustering around 5%.A retrieval score of 0.1 means a head retrieves at least 10% of the requested tokens.Source: §3.1 · p. 4

3.2 Dynamically Activated Based on Tokens and Contexts

  • Llama-2-7B-80K and Yi-6B-200K have 12 and 36 strongest retrieval heads, respectively, that activate under all considered contexts, while weaker heads activate only on certain tokens and contexts.Activation frequency measures whether a head activates on at least one token; retrieval score measures its average number of activated tokens.Source: §3.2 · p. 5, Figure 4 · p. 5

3.3 Intrinsic

  • Retrieval-head patterns persist across continued pretraining, chat fine-tuning, and sparse upcycling, with Pearson correlations exceeding 0.8 for base models and their variants but below 0.1 across different model families.The findings indicate that retrieval heads are already present in base models and are only modestly altered by subsequent small-scale training.Source: §3.3 · pp. 5–6, Figure 5 · p. 5, Figure 6 · p. 6

4 Influence on Downstream Tasks

Retrieval heads support long-context factuality and extractive QA, while their influence on reasoning depends on whether answers are generated directly or through chain-of-thought. Masking them reveals failure patterns in needle retrieval and impairs CoT reasoning.

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

  • Masking retrieval heads severely damaged Needle-in-a-Haystack performance, while masking random heads had much smaller effects; masking 50 heads reduced all models below 50.Successful needle retrieval coincided with retrieval-head activation, while failed retrieval and hallucination coincided with partial or absent activation. As stronger heads were masked, incomplete retrieval emerged before hallucinations became more prevalent.Source: §4.1 · pp. 6–7, §4 · p. 6
  • Needle-retrieval failures included incomplete retrieval, hallucination, and wrong extraction; hallucinations often involved retrieval heads attending to the input’s initial token.Wrong extraction could occur even when retrieval heads were active if they focused on incorrect sections. Incomplete retrieval began as the strongest heads were masked, typically at retrieval scores above 0.4; further masking made hallucinations more prevalent.Source: §4.1 · pp. 6–7

4.2 Influence on Extractive QA

  • Masking retrieval heads reduced extractive-QA F1 scores by 9.2% and 23.1%, whereas masking non-retrieval heads had no significant impact.The test used questions about specific information in up-to-date news articles, designed to assess knowledge not present in the model’s internal knowledge.Source: Figure 8 · p. 7, §4.2 · p. 7

4.3 Chain-of-Thought Reasoning also Requires Retrieval Heads

  • For CoT reasoning, masking retrieval heads significantly impaired performance and led the model to miss important input information and hallucinate.With answer-only prompting, masking retrieval or random heads did not meaningfully affect performance; the authors leave deeper investigation of the CoT–retrieval-head relationship to future work.Source: §4.3 · p. 7

5 Discussions

The discussion characterizes retrieval heads as context-dependent information routers, not program-inference mechanisms, and connects their role to full attention and potential KV-cache compression. It also notes their importance for chain-of-thought reasoning.

General Functionalities of Attention Heads

  • Retrieval heads significantly influence chain-of-thought tasks, where each reasoning step may need information from earlier context.Masking them can make the model “blind” to important information in the question description and lead to incorrect reasoning chains.Source: Figure 10 · p. 8, Figure 11 · p. 8
  • Retrieval heads redirect information according to context, unlike induction heads that search repeated input patterns, and do not infer programs.The authors suggest other attention-head algorithms and functions remain to be discovered.Source: §5 · pp. 7–8

Relationship to Local and Linear Attention and State-Space Models

  • Full attention appears necessary for precise retrieval from arbitrary context locations because retrieval heads must access the full KV cache.The authors note that, to their knowledge, no linear-attention or state-space architecture passes Needle-in-a-Haystack; Mistral v0.1’s sliding-window attention failed, while v0.2’s full attention passed.Source: §5 · p. 8

Applications to KV Cache Compression

  • 100K tokens require more than 50GB of KV cache for LLaMA 2 7B, versus less than 1GB for 2K context; 100K-context concurrency can be 50 times lower.The authors propose that pruning KV cache associated with non-retrieval heads could reduce deployment costs, leaving this for future research.Source: §5 · pp. 8–9

6 Conclusions

The paper identifies retrieval heads as attention heads that implement conditional copying, redirecting information from input to output. Their activity supports long-context retrieval and input faithfulness, especially in tasks requiring precise recall.

6 Conclusions

  • Retrieval heads implement a conditional copy algorithm that redirects information from input tokens to output tokens.The paper identifies these heads as a special set of attention heads.Source: §6 · p. 9
  • Retrieval heads are presented as the primary reason successful long-context models pass the Needle-in-a-Haystack test.Their activation is also linked to whether a language model is faithful to input or hallucinates.Source: §6 · p. 9
  • Retrieval heads influence downstream tasks requiring precise input recall more strongly than non-retrieval heads, including extractive question answering and chain-of-thought reasoning.Source: §6 · p. 9
  • The authors expect this work to foster research on reducing hallucination, improving reasoning, and compressing the KV cache.Source: §6 · p. 9

Related papers

Loading 2404.15574v1…