← All papers  /  Nov 3, 2026
Training

A Circuit for In-Context Learning

First page
A Circuit for In-Context Learning
The curator’s take

Andrew Gao and William Gao (Stanford) trace the full circuit for few-shot classification with arbitrary labels in three 3B to 4B transformers, showing that in-context learning here is a gated key-value lookup over label positions.

Ask this paper

Key points
01

Write step. Each label token's residual stream encodes the label identity. Middle-layer heads write the class of the preceding demonstration into it as a key.

02

Read step. At the query, upper-layer heads match the query's class against stored keys and copy out the matching label.

03

Gate and suppression. An early-layer signal for whether any label has repeated gates the readout. Without it, readers attend to the start of the prompt. A separate head set suppresses the most recent label.

04

Evidence. Each step is established with activation patching, head ablations against random-head baselines and steering. A two-constant-per-head softmax model fit on one prompt distribution predicts another.

05

Generality. The circuit is shared across the three models, appears from 0.6B to 8B in one family, survives instruction tuning, and a transformer trained from scratch develops the same parts. Failed predictions are reported.

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.

Every Monday
Get next week’s papers.
Subscribe on Substack