# The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation

arXiv: https://arxiv.org/abs/2505.15807 · HTML: https://mlsanity.com/paper/2505.15807/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 77 of its 89 references that are in the map.
- First posted on arXiv 21 May 2025.

## Builds on (most cited first, 30 of 77)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps](https://mlsanity.com/paper/1312.6034/) (2013)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Language Models as Knowledge Bases?](https://mlsanity.com/paper/1909.01066/) (2019)
- [A Call for Clarity in Reporting BLEU Scores](https://mlsanity.com/paper/1804.08771/) (2018)
- [Inference-Time Intervention: Eliciting Truthful Answers from a Language Model](https://mlsanity.com/paper/2306.03341/) (2023)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Attention is not Explanation](https://mlsanity.com/paper/1902.10186/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [DialoGPT: Large-Scale Generative Pre-training for Conversational Response Generation](https://mlsanity.com/paper/1911.00536/) (2019)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [In-Context Retrieval-Augmented Language Models](https://mlsanity.com/paper/2302.00083/) (2023)
- [Retrieval Augmentation Reduces Hallucination in Conversation](https://mlsanity.com/paper/2104.07567/) (2021)

## Built on by (most cited first, 11 of 11)

- [RoPE Distinguishes Neither Positions Nor Tokens in Long Contexts, Provably](https://mlsanity.com/paper/2605.15514/) (2026)
- [Large Vision-Language Models Get Lost in Attention](https://mlsanity.com/paper/2605.05668/) (2026)
- [Task Vectors, Learned Not Extracted: Performance Gains and Mechanistic Insight](https://mlsanity.com/paper/2509.24169/) (2025)
- [Interpreting Language Models Through Concept Descriptions: A Survey](https://mlsanity.com/paper/2510.01048/) (2025)
- [Atlas-Alignment: Making Interpretability Transferable Across Language Models](https://mlsanity.com/paper/2510.27413/) (2025)
- [Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis](https://mlsanity.com/paper/2509.24164/) (2025)
- [Fast & Faithful Function Vectors](https://mlsanity.com/paper/2606.05079/) (2026)
- [Mechanistic Attention Guidance for Agent Memory Refinement](https://mlsanity.com/paper/2607.17621/) (2026)
- [ResLRP: The Role of Residual Cancellation in Attribution Instability in Vision Transformers](https://mlsanity.com/paper/2609.17152/) (2026)
- [The Commit-Abstain Circuit: Why Language Models Hallucinate Instead of Abstaining](https://mlsanity.com/paper/2609.32964/) (2026)
- [Beyond the Shadows of Plato's Cave: Evaluating False Memory in Autonomous Agents via Counterfactual Reasoning](https://mlsanity.com/paper/2609.39473/) (2026)

## Abstract (arXiv)

Large language models are able to exploit in-context learning to access external knowledge beyond their training data through retrieval-augmentation. While promising, its inner workings remain unclear. In this work, we shed light on the mechanism of in-context retrieval augmentation for question answering by viewing a prompt as a composition of informational components. We propose an attribution-based method to identify specialized attention heads, revealing in-context heads that comprehend instructions and retrieve relevant contextual information, and parametric heads that store entities' relational knowledge. To better understand their roles, we extract function vectors and modify their attention weights to show how they can influence the answer generation process. Finally, we leverage the gained insights to trace the sources of knowledge used during inference, paving the way towards more safe and transparent language models.