The paper identifies a small subset of attention heads, termed 'retrieval heads,' which are primarily responsible for retrieving relevant information (copy-paste or paraphrase) from long contexts. Across all tested model families (Llama, Yi, Qwen, Mistral, Mixtral) and variants, only about 5% of attention heads are retrieval heads, and they exist intrinsically in base models. Causal intervention by masking retrieval heads sharply degrades performance on the needle-in-a-haystack task, while masking random non-retrieval heads has little effect. The study further shows that these heads are dynamically activated based on context and that their role extends to extractive QA and chain-of-thought reasoning.
Evidence
interventional
Key metric
Accuracy drops from 94.7% (20 random heads masked) to 63.6% (top 20 retrieval heads masked) on NIAH for Llama-2-7b-80k; retrieval heads are <5% of all heads; performance drops below 50% when masking >50 retrieval heads.
Caveat
The retrieval score threshold of 0.1 and the focus on copy-paste behavior for detection may not capture all retrieval functions; tasks heavily reliant on parametric knowledge (e.g., MMLU without CoT) are less affected.