# Contextualize-then-Aggregate: Circuits for In-Context Learning in Gemma-2 2B

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 37 references that are in the map.
- First posted on arXiv 31 Mar 2025.

## Builds on (most cited first, 25 of 25)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (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)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Causal Abstractions of Neural Networks](https://mlsanity.com/paper/2106.02997/) (2021)
- [In-Context Learning Creates Task Vectors](https://mlsanity.com/paper/2310.15916/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning](https://mlsanity.com/paper/2305.14160/) (2023)
- [Attribution Patching Outperforms Automated Circuit Discovery](https://mlsanity.com/paper/2310.10348/) (2023)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Have Faith in Faithfulness: Going Beyond Circuit Overlap When Finding Model Mechanisms](https://mlsanity.com/paper/2403.17806/) (2024)
- [The Learnability of In-Context Learning](https://mlsanity.com/paper/2303.07895/) (2023)
- [Information Flow Routes: Automatically Interpreting Language Models at Scale](https://mlsanity.com/paper/2403.00824/) (2024)
- [A Theory of Emergent In-Context Learning as Implicit Structure Induction](https://mlsanity.com/paper/2303.07971/) (2023)
- [Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scale](https://mlsanity.com/paper/2212.09095/) (2022)
- [Induction Heads as an Essential Mechanism for Pattern Matching in In-context Learning](https://mlsanity.com/paper/2407.07011/) (2024)
- [Linear Transformers are Versatile In-Context Learners](https://mlsanity.com/paper/2402.14180/) (2024)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (2024)

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

- [Just-in-time and distributed task representations in language models](https://mlsanity.com/paper/2509.04466/) (2025)
- [Understanding In-context Learning of Addition via Activation Subspaces](https://mlsanity.com/paper/2505.05145/) (2025)
- [Decomposing Representation Space into Interpretable Subspaces with Unsupervised Learning](https://mlsanity.com/paper/2508.01916/) (2025)
- [Causality neq Invariance: Function and Concept Vectors in LLMs](https://mlsanity.com/paper/2602.22424/) (2026)
- [Born a Transformer -- Always a Transformer? On the Effect of Pretraining on Architectural Abilities](https://mlsanity.com/paper/2505.21785/) (2025)
- [Discursive Circuits: How Do Language Models Understand Discourse Relations?](https://mlsanity.com/paper/2510.11210/) (2025)
- [Latent Concept Disentanglement in Transformer-based Language Models](https://mlsanity.com/paper/2506.16975/) (2025)
- [Emergent Structured Representations Support Flexible In-Context Inference in Large Language Models](https://mlsanity.com/paper/2602.07794/) (2026)
- [Symmetry Reveals Layerwise Dynamics: How Transformers Perform In-Context Classification](https://mlsanity.com/paper/2604.11613/) (2026)
- [Are In-Context Images Worth 10 Dimensions?](https://mlsanity.com/paper/2609.37659/) (2026)

## Abstract (arXiv)

In-Context Learning (ICL) is an intriguing ability of large language models (LLMs). Despite a substantial amount of work on its behavioral aspects and how it emerges in miniature setups, it remains unclear which mechanism assembles task information from the individual examples in a fewshot prompt. We use causal interventions to identify information flow in Gemma-2 2B for five naturalistic ICL tasks. We find that the model infers task information using a two-step strategy we call contextualize-then-aggregate: In the lower layers, the model builds up representations of individual fewshot examples, which are contextualized by preceding examples through connections between fewshot input and output tokens across the sequence. In the higher layers, these representations are aggregated to identify the task and prepare prediction of the next output. The importance of the contextualization step differs between tasks, and it may become more important in the presence of ambiguous examples. Overall, by providing rigorous causal analysis, our results shed light on the mechanisms through which ICL happens in language models.