A Circuit for In-Context Learning
Abstract
In-context learning (ICL) lets a pretrained language model perform a task specified by demonstrations in its prompt, with no parameter updates. With arbitrary labels the model must acquire a mapping it cannot have memorized, so the learning happens within a single forward pass. Prior work found candidate parts but not how they compose. We give the composition for few-shot sentence classification with arbitrary labels in three 3-4B-parameter transformers. The model implements a key-value memory over label-token positions: each label token's residual stream encodes the label's identity, and middle-layer attention heads write into it the class of the preceding demonstration as its key. At the query, upper-layer heads match the query's class against the stored keys and copy out the matching labels. An early-layer signal encoding whether any label has repeated gates this readout; in its absence the readers attend to the start of the prompt and no label is favored. A separate set of heads suppresses the most recently seen label. Each step is established causally by activation patching, head ablations against random-head baselines, and steering with the model's own signals. A softmax model of the readout with two constants per head, fit on one prompt distribution, predicts another. The computation is conserved across the three models, the named heads alone reproduce the behaviour, and a transformer trained from scratch grows the same parts. The circuit transfers to semantic concepts and to labels with meaning, and the same parts, with a causal gate, appear from 0.6B to 8B in one family and survive instruction tuning. ICL in this setting is a gated associative lookup built from a few identifiable component types, traced causally from demonstrations to prediction.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper provides a rigorous, causal decomposition of in-context learning in large language models, identifying a gated key-value memory mechanism composed of specific attention heads that write class keys to label tokens and read them out via a softmax-based lookup, validated across multiple model families and scales.
Methodology
The paper employs a rigorous mechanistic interpretability framework, utilizing activation patching, head ablations, and steering vectors to dissect the in-context learning (ICL) process. The methodology is exceptionally strong, relying on "twin prompts" (minimal pairs) to isolate causal effects and strict controls (random heads, random directions) to validate findings. The authors move beyond correlation to establish causality by showing that specific attention heads act as a key-value memory, writing class keys to label tokens and reading them out at the query. The introduction of a "softmax law" to model the readout mechanism is a novel theoretical contribution that allows for prediction of model behavior in unseen regimes (fewer slots per class). The analysis of the "gate" mechanism, which suppresses retrieval when labels are unique, is particularly insightful, explaining the model's default behavior via attention sink dynamics.
Results & evaluation
The experiments are comprehensive, spanning three different 3-4B parameter models (Qwen3, SmolLM3, Llama-3.2) to ensure the findings are not model-specific. The authors test the circuit's robustness by training a transformer from scratch and observing the emergence of the same components. They also extend the analysis to natural language labels and semantic concepts, showing that the circuit handles arbitrary labels via retrieval but relies on different mechanisms (brakes/MLPs) for word meaning. The inclusion of a "gain bound" proposition explains why the computation is distributed across many heads rather than concentrated in a few, providing a theoretical underpinning for the empirical observations. The evaluation of instruction-tuned models and varying model sizes (0.6B to 8B) further strengthens the generalizability of the results.
Reproducibility
The paper provides high reproducibility standards. It specifies the exact models used, the intervention techniques (mean-ablation, steering doses), and the metrics (margin, logit contributions). The authors explicitly discuss the choice of reference values for ablation tests and how it affects results, demonstrating a deep understanding of the experimental pitfalls. The code and data are not explicitly linked in the provided text, but the detailed procedural descriptions allow for replication. The use of standard HuggingFace implementations and bfloat16 precision on H100s makes the setup accessible.
Limitations
The study focuses primarily on few-shot sentence classification with arbitrary or simple labels. While it touches on semantic concepts, the complexity of real-world tasks (long-form generation, multi-step reasoning) is outside the scope. The "softmax law" breaks down for longer prompts, indicating that the model's attention budget management changes with context length, a limitation the authors acknowledge but do not fully resolve. The analysis is limited to base models and instruction-tuned variants of specific architectures; it may not generalize to MoE or non-transformer architectures.
Broader impact
This paper significantly advances the field of mechanistic interpretability by providing a complete, causal account of a core LLM capability. It shifts the understanding of ICL from a "black box" phenomenon to a structured algorithm of gated associative lookup. This insight has implications for model safety (understanding how models acquire tasks from prompts), efficiency (potential for pruning or steering specific heads), and theory (validating the key-value memory hypothesis). The identification of "brakes" and "gates" offers new levers for controlling model behavior. The work serves as a benchmark for future interpretability studies, setting a high bar for causal validation. The paper provides a rigorous, causal decomposition of in-context learning in large language models, identifying a gated key-value memory mechanism composed of specific attention heads that write class keys to label tokens and read them out via a softmax-based lookup, validated across multiple model families and scales.
Author of this paper? Send a correction.