Retrieval Heads are Dynamic

Recent studies have identified “retrieval heads” in Large Language Models (LLMs) responsible for extracting information from input contexts. However, prior works largely rely on static statistics aggregated across datasets, identifying heads that perform retrieval on average. This perspective overlooks the fine-grained temporal dynamics of autoregressive generation. In this paper, we investigate retrieval heads from a dynamic perspective. Through extensive analysis, we establish three core claims: (1) Dynamism: Retrieval heads vary dynamically across timesteps; (2) Irreplaceability: Dynamic retrieval heads are specific at each timestep and cannot be effectively replaced by static retrieval heads; and (3) Correlation: The model’s hidden state encodes a predictive signal for future retrieval head patterns, indicating an internal planning mechanism. We validate these findings on the Needle-in-a-Haystack task and a multi-hop QA task, and quantify the differences on the utility of dynamic and static retrieval heads in a Dynamic Retrieval-Augmented Generation framework. Our study provides new insights into the internal mechanisms of LLMs.

💡 Research Summary

The paper investigates “retrieval heads” – attention heads that copy tokens directly from the input context – in large language models (LLMs) from a temporal, dynamic perspective. Prior work identified a static set of such heads by aggregating attention statistics over many samples, implicitly assuming that the same heads are responsible for retrieval at every generation step. The authors argue that this view ignores the fine‑grained dynamics of autoregressive generation.

Definition and Task
A retrieval head is defined as an attention head that, at a given timestep, assigns the highest attention weight to a token that (a) lies inside a pre‑designated “needle” region of the input and (b) matches the token the model actually generates. When both conditions hold, the head receives a binary “copy‑paste” score of 1. The Needle‑in‑a‑Haystack (NIAH) benchmark, where a short “needle” is hidden inside a long distractor document, is used to probe this behavior. A multi‑hop QA dataset is later added to test more complex reasoning.

Claim 1 – Dynamism
The authors track the copy‑paste score of each head across the entire generation sequence. Visualizations (Figure 1) show that heads frequently turn on and off, with high variance across timesteps. Quantitatively, three metrics are reported:

• Jaccard w/ Static (overlap between dynamic heads at a step and the top‑20 static heads) ranges from 0.18 to 0.46, indicating limited overlap.
• Adjacent Jaccard (overlap between consecutive steps) is low (0.28–0.51), confirming rapid turnover.
• Entropy of the head distribution exceeds 3.0 (up to 4.9), far above the entropy of a uniform selection over 20 heads (≈2.99), implying that at least 20 different heads become “retrieval” heads throughout a single generation. These results demonstrate that a static set cannot capture the real‑time retrieval behavior.

Claim 2 – Irreplaceability
Two ablation studies test whether static heads can substitute dynamic ones.

1. Step‑wise masking: For each token, the identified dynamic heads are masked, and the model re‑generates the token. The same number of heads is masked from (a) the top‑20 static heads and (b) a random set. On Llama‑3.1‑8B, masking dynamic heads causes a dramatic drop in exact‑match accuracy and ROUGE‑L, far exceeding the effect of masking static or random heads (Figure 2).
2. Progressive masking: Increasing numbers of dynamic heads (k) are masked, and the model’s “compensated” heads—those that become retrieval heads after masking—are counted. While many compensated heads belong to the static top‑20 (the blue line in Figure 3 rises to 4–7), overall retrieval performance (red line) continues to decline sharply. This shows that static heads cannot fully replace the context‑specific dynamic heads.

Claim 3 – Correlation / Planning
The authors hypothesize that the model’s hidden state contains a signal about which heads will be needed next. They train a simple linear classifier that takes the final hidden vector at timestep t and predicts the set of retrieval heads at timestep t + 1. The classifier achieves substantially higher accuracy than a random baseline (≈2×), indicating that the hidden state encodes a predictive plan for future head usage. This suggests an internal planning mechanism rather than a purely reactive process.

Dynamic Retrieval‑Augmented Generation (D‑RAG)
Leveraging the above insights, the authors build a Dynamic Retrieval‑Augmented Generation pipeline. At each generation step, the model’s hidden state selects the most probable retrieval heads, extracts the tokens they attend to, and feeds them to an external knowledge base for augmentation. Compared to a static‑head RAG baseline, D‑RAG improves both NIAH accuracy and multi‑hop QA F1 by roughly 7 percentage points, demonstrating practical utility.

Related Work and Positioning
The paper situates itself within mechanistic interpretability (induction heads, attention sinks, “lost in the middle”) and prior retrieval‑head studies (Wu et al., 2024; Zhang et al., 2025; Fu et al., 2024). It highlights that earlier works focused on static circuit identification, whereas this work emphasizes temporal modulation of attention circuits.

Limitations and Future Directions
The study concentrates on exact copy‑paste retrieval; extending the analysis to semantic retrieval, multimodal inputs, or more sophisticated planning models (e.g., recurrent networks, reinforcement learning) is left for future work. Additionally, the linear predictor is a proof‑of‑concept; richer models could yield finer‑grained planning insights.

Conclusion
The authors empirically validate three core claims: (1) retrieval heads are highly dynamic across generation timesteps; (2) these dynamic heads cannot be substituted by a static set without severe performance loss; (3) the model’s hidden state predicts future retrieval‑head patterns, indicating an internal planning mechanism. Dynamic head selection improves Retrieval‑Augmented Generation, offering new avenues for interpretability, head pruning, and state‑aware interventions in LLMs.