The Mechanism of In-Context Learning in Large Models Cracked: Gated Key-Value Memory Circuit

Introduction: The Mechanistic Puzzle of ICL

In-Context Learning (ICL) enables pretrained language models to perform new tasks from examples in the prompt alone, without updating parameters. When labels are arbitrary, the model must acquire a mapping it could not possibly have memorized, and learning happens within a single forward pass. Prior work identified candidate components but did not explain how they combine. Recent research presents the complete circuit for few-shot sentence classification under arbitrary labels, revealing how ICL works.

Residual Stream is the information backbone running through all layers of a Transformer. At each layer, attention heads and MLP sublayers read from the residual stream, compute, and add their results back to it, so it acts as a shared "working memory." Researchers often refer to the representation of a concept (such as "label identity") in the residual stream as the encoding of that concept. Understanding the residual stream helps grasp the statement in this article that "the residual-stream encoding of a label token encodes label identity": it means label information exists in this channel in a linearly readable form that later layers can directly exploit.

The Core of the Circuit: Gated Key-Value Memory

The researchers discovered this circuit in three Transformers with 3-4B parameters. The models implement key-value memory at label token positions: the residual stream of each label token encodes label identity, and middle-layer attention heads write the category of the preceding demonstration as a key. At query time, upper-layer heads match the query category against the stored keys and copy out the corresponding label.

An early-layer signal encodes whether labels repeat and gates this readout process; without this signal, the readout heads attend to the beginning of the prompt and no label is preferred. Another group of heads suppresses recently appeared labels. Each step is established through activation patching, head ablation, and guided causality.

Activation Patching is a causal intervention method: first run the model on the original input and record an activation at some location, then, after changing the input (such as replacing demonstration labels), substitute that activation back and observe whether the output is restored. If restored, the activation has a causal role in the behavior. Head Ablation directly removes or zeroes out an attention head to test its necessity. Guided Causal usually refers to combining gradients or directional guidance to localize and verify causal pathways. Together, the three form a chain of evidence moving from correlation to causality.

Causal Verification and Generalization

The study verified the mechanism through multiple causal interventions. A softmax readout model with two constants, fitted on one prompt distribution, can predict another distribution. The computation is conserved across the three models, name heads can reproduce the behavior, and Transformers trained from scratch also grow the same components. The circuit transfers to semantic concepts and meaningful labels, and the causal gating appears across model families from 0.6B to 8B parameters and is retained after instruction tuning.

Softmax readout model here refers to a simplified computational model: the readout head applies softmax normalization to keys and then weighted-copies values, with two constants (such as temperature or bias) characterizing the sharpness of matching. The researchers fit this model to behavior on one prompt distribution and then predict another distribution, thereby testing the generalizing power of the mechanistic hypothesis. If the simplified model can predict across distributions, it indicates that the real circuit is functionally equivalent to this readout form, supporting the "gated associative lookup" interpretation.

Significance and Outlook

In this setting, ICL is a gated associative lookup built from a small number of identifiable components, with causal tracing from demonstrations to predictions. This provides a foundation for understanding the capabilities of large models and may inspire more efficient architecture designs.

One-line Summary

Research reveals that ICL under arbitrary labels is implemented by a gated key-value memory circuit and is causally verifiable.

Paragraph Summary

This study discovered the complete ICL circuit for few-shot sentence classification under arbitrary labels in three 3-4B-parameter Transformers. The core mechanism is gated key-value memory: middle-layer attention heads write demonstration categories as keys, upper-layer heads match query categories and copy labels; an early-layer signal encodes whether labels repeat and gates the readout, while another group of heads suppresses recent labels. The study verified causality through activation patching, head ablation, and guided causality, and showed that a simplified softmax readout model can predict behavior across prompt distributions. The circuit is conserved across the three models, transfers to semantic concepts, appears across model families from 0.6B to 8B parameters, and is retained after instruction tuning.