Research map / Large language models / Retrieval and knowledge
The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Uses AttnLRP-based relevance attribution to identify in-context and parametric attention heads in instruction-tuned LLMs, separates task and retrieval heads, and shows that patching their function vectors or attention weights steers QA behavior and enables a probe for knowledge provenance.
Problem
The mechanism by which LLMs use in-context retrieval augmentation (choosing between parametric and contextual knowledge) for question answering is not well understood, and attention-weight analyses are unreliable.
Approach
Apply AttnLRP relevance propagation to attention heads to score each head's contribution in open-book (counterfactual) versus closed-book settings, yielding in-context and parametric heads; further split in-context heads by relevance to question tokens (task heads) versus answer tokens (retrieval heads), then use their outputs as function vectors or reweight their attention to steer generation, an
Key results
- Ablating 20 identified in-context heads reduces open-book recall by 13.86%-63.84% across configurations, with comparable drop to ablating 100 AWR heads.
- Adding task-head function vectors yields 78.75-point average improvement over random heads on biography intensional-frame tasks; parametric FVs recover original attributes with a 30.41-point gain over random.
- Retrieval-head probe achieves ROC AUC of at least 94% and top-1 localization accuracy of at least 84% across Llama 3.1, Mistral v0.3, and Gemma 2.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Ablate identified in-context heads (ablation) | Llama 3.1 · Open-book QA across configurations | recall: 13.86%-63.84% | Hurt |
| Ablate identified parametric heads (ablation) | Llama 3.1 · Closed-book QA across configurations | recall: 44.26%-68.66% | Hurt |
| Ablate randomly selected heads (ablation) | Llama 3.1 · QA | recall | No clear effect |
| Ablate parametric heads (ablation) | Llama 3.1 · NQ-Swap counterfactual open-book QA | recall | Helped |
| Ablate parametric heads (ablation) | Llama 3.1 · NQ-Swap counterfactual open-book QA | recall for counterfactual answers | Helped |
| Ablate in-context heads (ablation) | Llama 3.1 · NQ-Swap counterfactual QA | recall for original answers | Helped |
| Use 20 identified in-context heads instead of 100 AWR heads (ablation) | Llama 3.1 · Open-book QA | recall reduction after ablation | Mixed |
| Extract and patch task-head function vectors rather than random-head vectors | Biography QA · Biography QA | recall: +78.75 points | Helped |
| Extract and patch parametric-head function vectors rather than random-head vectors | Biography QA · Biography QA | recall: 30.41 points | Helped |
| Increase retrieval-head attention weights on needle tokens rather than use random heads | Biography NIAH · Needle-in-a-haystack retrieval | recall: 83.15 points | Helped |
| Select only task heads that are also strong in-context heads, instead of the top 40 task heads (ablation) | Biography QA · At which date was he born? (Gemma 2) | recall: 99% -> 22% | Hurt |
| Select only task heads that are also strong in-context heads, instead of the top 40 task heads (ablation) | Biography QA · At which institution was he educated? (Llama 3.1) | recall: 98% -> 96% | No clear effect |
Surprising findings
- Ablating in-context heads in the closed-book setting also causes a non-negligible performance drop, indicating they also process the question/instruction.
- Gemma 2 exhibits parametric task heads, so task heads are not exclusively in-context.
Limitations
- Focuses on attention heads and does not study interaction with MLP components such as knowledge neurons.
- Analyses are mainly on QA with only a preliminary machine translation experiment.
- Some heads show inhibitory effects and some intermediate heads are not exclusively in-context or parametric.
Experimental setup
Datasets: NQ-Swap TriviaQA DPR Wikipedia dump (Dec 2018) Wikidata biography dataset
Benchmarks: NQ-Swap TriviaQA
Metrics: recall exact match BERT matching (BEM) K-Precision ROC AUC top-1 localization accuracy
Baselines: random heads Attention Weight Recall (AWR) heads
Builds on: AttnLRP function vectors retrieval heads logit lens
Closest papers by meaning
- Retrieval Head Mechanistically Explains Long-Context Factuality (ICLR 2025 · similarity 0.90)
- Retrieval from Within: An Intrinsic Capability of Attention-Based Models (NeurIPS 2026 · similarity 0.89)
- Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning (NeurIPS 2025 · similarity 0.88)
- Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads (NeurIPS 2026 · similarity 0.88)
- Making Retrieval-Augmented Language Models Robust to Irrelevant Context (ICLR 2024 · similarity 0.88)
- Massive Values in Self-Attention Modules are the Key to Contextual Knowledge Understanding (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.