Sam Babichenko

Reading notes / In-context learning

Induction heads

Olsson et al., 2022

Context. Language models predict later tokens in a context better than early ones. The earlier circuits paper found induction heads in two-layer attention-only transformers, a two-head circuit that completes [A][B]…[A]→[B][A][B]\dots[A]\to[B]. This paper asks whether that circuit is the mechanism behind most in-context learning in transformers of every size, and answers with six lines of evidence of very different strength, from 34 models tracked through training and more than 50,000 head ablations. Early in training every multi-layer transformer goes through a short phase change in which induction heads form, in-context learning jumps to its final level, and the loss curve shows its only bump. An induction head is an in-context nearest-neighbor rule (find an earlier position whose prefix matches the present one and copy what came next), which, applied to abstract features, can plausibly cover much more than literal repetition, a lookup quite unlike the regression estimators of the theory papers in this section.

Setting. The models are small attention-only and small MLP models with 1 to 6 layers (200 snapshots each), full-scale models from 13M to 13B parameters (14 or 15 snapshots), and "smeared-key" 1- and 2-layer models. In-context learning is measured on all text rather than on few-shot tasks. The in-context learning score is the loss at token 50 minus the loss at token 500. A head is an induction head if, on a random token sequence repeated four times, it attends from each token to the token that followed an earlier copy of it (prefix matching) and its direct effect raises the logit of the attended token (copying). In the small models the mechanism is a previous-token head in an earlier layer whose output the induction head's keys read (K-composition), together with a copying OV circuit.

Main findings. In every model with more than one layer, within a short window early in training (about 2.5 to 5 billion tokens, 1 to 3 billion for the small models), the in-context learning score rises from under 0.15 to about 0.4 nats, with more than 75% of its final value forming in the window, and then stays there for the rest of training and across model sizes. In the same window induction heads form and the loss curve has its only non-convex stretch; one-layer models show none of this. The authors grade their evidence that induction heads drive in-context learning by model class: strong and causal for small attention-only models; for small models with MLPs, strong and causal for "some" of it and medium for "most"; for large models, medium and correlational (and mechanistic for "some").

Step 1

The events coincide in training, which is suggestive but leaves room for a common cause. The score, the prefix-matching scores, the loss bump and a turn in the models' per-token loss trajectories move together in every multi-layer model, and in no one-layer model, which cannot form induction heads because a single head's attention score cannot depend on the token before the attended one. The weight-decay change comes after the window, but the warmup, which ends at 1.5 billion tokens, falls inside the small models' window; a books-only dataset shows the same transition. The open alternative is a common cause: the phase change might be when models learn to compose heads across layers at all, enabling induction heads and other composed mechanisms together. The large models also have only 15 snapshots, which makes their timing evidence weaker.

Step 2

Making the prerequisite available in one layer moves the phase change. Smeared keys, kj←σ(α)kj+(1−σ(α))kj−1k_j\leftarrow\sigma(\alpha)k_j+(1-\sigma(\alpha))k_{j-1} with a trainable α\alpha per head, let a single head see the previous token, so prefix matching no longer needs an earlier layer. One-layer models then acquire in-context learning and two-layer models acquire it earlier. This intervention rules out explanations tied to depth as such, but it was run only for 1- and 2-layer models and says nothing about whether other mechanisms take over later or at scale.

Step 3

Knocking out induction heads removes almost all in-context learning in small attention-only models. Each head in each snapshot is ablated in turn with every attention pattern frozen to the unablated run, so only downstream values change. In attention-only models the logits are, up to LayerNorm, a sum of per-head terms plus a direct path, so this is a principled attribution, and it assigns almost all in-context learning to the induction heads from the start of the phase change to the end of training; ablating most other heads raises the score, by hurting ordinary prediction. With MLPs the attribution is murkier, since heads interact nonlinearly with MLPs and redundant heads mask each other's marginal effect. No ablations were run on the full-scale models.

Step 4

In a 13B model, heads that meet the narrow definition also do abstract in-context tasks. Some heads selected as induction heads on repeated random tokens do more than copy: one repeats names and whole sentences, one attends, during English-to-French or German translation, to the earlier word that comes next semantically, and one maps lines like "(month) (animal)" to one of four labels by attending to an earlier line of the same template (about 65% of its attention on correct positions). This is [A∗][B∗]…[A]→[B][A^*][B^*]\dots[A]\to[B] with A∗≈AA^*\approx A, B∗≈BB^*\approx B in some feature space. Three hand-picked examples show capability and say nothing about what fraction of in-context learning such heads carry.

Step 5

The mechanism explains why induction heads must help. A head that finds earlier positions whose prefix matches the present one and copies the next token is a nearest-neighbor predictor whose data set is the context, so loss must fall with position wherever it applies. In two-layer attention-only models the parameters can be read directly, a previous-token term in the QK circuit and a copying OV circuit, both with positive eigenvalues, and no rival mechanism is visible. In a replication, an idealized induction pattern recovers about 65% of the ablation effect. The generalization to abstract features is a hypothesis, and GPT-2 shows another route to induction through positional embeddings, so the circuit is not unique.

Step 6

Everything measured looks the same from small to large models, which supports the hypothesis only by analogy. From a 2-layer attention-only model to the 13B model, every multi-layer model goes through the phase change, gains the same 0.4 nats and forms induction heads. The equal score is less strange than it looks: large models take most of their advantage in the first ten or so tokens, and after that both lower their loss by about the same amount, which for a large model is the harder reduction from a lower baseline. The authors name the strongest alternative: larger models have room for many other composed heads, and if they all form in the same window, non-induction composition could account for most of the gain above some size.

Consequences and limits. The paper ties a macroscopic event in training, the loss bump, to a specific circuit, and for small attention-only models the ablations make it close to established that the circuit carries most of in-context learning, defined as lower loss later in context. The limits are that for large models the case is correlational plus analogy, with no ablations; that the in-context learning measured is next-token loss on natural text, which correlates with few-shot task performance but is not the same thing; that the abstract behaviors are anecdotes; and that the paper does not address whether models also run something like the gradient-descent or Bayesian procedures that the theory papers construct.

Terms

In-context learning score
loss at token 50 minus loss at token 500; jumps from under 0.15 to about 0.4 nats
Induction head
head that completes [A][B]…[A]→[B][A][B]\dots[A]\to[B]: high prefix-matching and copying scores
Prefix matching
on a repeated random sequence, attends to the token after an earlier copy of the present token
Copying
head's direct effect raises the logit of the token it attends to (OV circuit)
Previous-token head
earlier head writing the preceding token into each position, read by the induction head's keys
Phase change
short early window (billions of tokens) where induction heads form, the score jumps and the loss bumps
Smeared key
kj←σ(α)kj+(1−σ(α))kj−1k_j\leftarrow\sigma(\alpha)k_j+(1-\sigma(\alpha))k_{j-1}: lets a one-layer model form induction heads
Pattern-preserving ablation
zero a head's output with attention patterns frozen; only downstream values change
Fuzzy induction
[A∗][B∗]…[A]→[B][A^*][B^*]\dots[A]\to[B] with A∗≈AA^*\approx A: in-context nearest neighbor on abstract features
Evidence by scale
small attention-only: causal; small with MLPs: causal for some; large: correlational and analogy

Quiz

  1. Why can a one-layer attention-only transformer not implement an induction head?
  2. Which alternative explanation do the authors consider the most serious for large models?
  3. Why do the abstract behaviors in the 13B model (translation, template matching) support the hypothesis only weakly?

Sources