In-context Learning and Induction Heads

interpretabilitymechanistic interpretabilityinduction headsin-context learningphase changetraining dynamics

Olsson et al. track 34 transformers through training and find that, in every model with more than one layer, induction heads form in the same narrow window in which in-context learning abruptly improves; with co-occurrence, an architecture change, ablations, examples and mechanistic reasoning they build an explicitly indirect case that induction heads may be the main mechanism of in-context learning.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph et al. (Anthropic) · 2022-03-08 · original

Language models get better at predicting text the further they read into a document: the later tokens of a passage are easier to predict than the early ones, because the model can use what it has already seen. This ability, in-context learning, is what lets a prompt teach a model a task without changing its weights. What mechanism inside the model produces it? The previous paper on the Transformer Circuits Thread, A Mathematical Framework for Transformer Circuits, had found induction heads in two-layer attention-only models and shown that they were the main source of in-context learning there. This paper, published on 8 March 2022, asks whether the same is true of real language models with many layers and MLP layers, where reverse engineering from the weights is not yet possible. The authors describe their evidence as "preliminary and indirect evidence for a tantalizing hypothesis" (Introduction): that induction heads, in a looser, fuzzier form, might implement the majority of in-context learning in large transformers. The paper is unusual in how openly it grades its own evidence, with a table for each argument rating its strength for small and large models.

Terms used in this note

  • In-context learning score. The paper's measure of in-context learning: the average loss on the 500th token of a context minus the average loss on the 50th token (section "In-context Learning"). Loss is measured in nats (natural-log units). The paper follows Kaplan et al. (2020) in treating in-context learning as falling loss with position in the context, rather than as success on a set of few-shot tasks, because this is more general and does not depend on a choice of tasks.
  • Induction head. Defined in this paper by behavior, on a sequence of random tokens that repeats: a head that attends from the current token back to the token that followed the current token's earlier occurrence (prefix matching) and raises the logit of the token it attends to (copying) (section "Induction Heads"). In the earlier paper's terms, it completes "[A][B] … [A]" with "[B]". Mechanically, in small models it works together with a previous-token head in an earlier layer.
  • Phase change. The authors' name for an abrupt change during training that shows up both in large-scale measurements (loss, in-context learning) and in internal structure (induction heads), by loose analogy with ice melting.
  • Ablation. Removing one attention head's contribution at test time and measuring what changes, like a gene knockout. The paper mostly uses a pattern-preserving ablation that zeroes the head's output while keeping every head's attention pattern as it was in an unablated run (appendix "Per-token losses (attention head ablations)").
  • Per-token loss vector and PCA. For each saved snapshot of a model, the losses on a fixed set of 10,000 tokens, one from each of 10,000 examples, form a vector; principal component analysis of these vectors across snapshots and models shows how models' predictions move over training (section "Per-Token Loss Analysis").

What the paper does

The models

The analysis covers 34 decoder-only transformers, each trained once, with many snapshots saved during training (section "Model Analysis Table"):

  • Small attention-only models, 1 to 6 layers, without MLP layers;
  • Small models with MLPs, 1 to 6 layers;
  • Full-scale models, from 4 layers and 13 million non-embedding parameters to 40 layers and 13 billion, from the same family as the models in A General Language Assistant as a Laboratory for Alignment;
  • Smeared-key models, an architectural variant described below, at one and two layers, each with a matched control.

The small models have a residual stream of width 768, 12 heads per layer and an 8,192-token context, and were trained for 10,000 steps (about 10 billion tokens) with 200 snapshots, one every 50 steps (section "Model Details"). The full-scale models have only 14 or 15 snapshots each, spaced exponentially. The training data was filtered Common Crawl, internet books and smaller sources including about 10% Python code; a second set of small models was trained on internet books only, to check that the results do not depend on the dataset. Ablations were run only on the small models; a footnote notes that the cost of a full set of ablations grows faster than model size.

Measurements

For each snapshot the authors record per-token losses, the in-context learning score, and three scores per attention head computed on synthetic sequences: copying, prefix matching and previous-token attention (appendix "Head activation evaluators"). They also compute a weight-based measure of prefix matching drawn from the earlier paper, and validate the behavioral scores against weight-based ones in the small attention-only models.

Six arguments

The paper is organized as six arguments, each with its own strength table (section "Arguments that induction heads are the mechanism for the majority of in-context learning"):

  1. Co-occurrence: induction heads form at the same time in-context learning improves.
  2. Co-perturbation: changing the architecture so that induction heads form earlier, or at all in a one-layer model, moves the improvement in in-context learning with them.
  3. Ablation: removing induction heads in small models removes most in-context learning.
  4. Examples of generality: induction heads in a large model also perform more abstract tasks such as translation.
  5. Mechanistic plausibility: the mechanism understood in small models could naturally generalize.
  6. Continuity: measurements look the same from small to large models, so the simplest explanation is that the mechanism is the same.

What it finds

A phase change early in training (Argument 1)

In-context learning develops abruptly in a narrow window early in training, roughly 2.5 to 5 billion tokens, and then stays nearly constant: below 0.15 nats of in-context learning before the window, roughly 0.4 nats after, and about the same amount across model sizes, except for the one-layer model, which never develops much (section "Argument 1"). In exactly the same window:

  • induction heads form, as measured by prefix-matching scores, in every model with more than one layer and never in the one-layer model;
  • the training loss shows a visible "bump", the only point at which the loss curve is not convex;
  • the models' trajectories in per-token-loss space turn sharply.

The authors report that the window does not appear to correspond to a scheduled change in learning rate, warmup or weight decay, and that it appears in the models trained on books only (section "Looking at the Phase Change More Closely"). The section "Model Details" adds that, in the small models, weight decay was reduced at about 5 billion tokens, past the phase change, and that the learning-rate warm-up, which ramps up over the first 1.5 billion tokens, is the only other scheduled change within the range of the phase change. On a passage of text, most of the change during the phase concerns repeated tokens: the model gets better at predicting a sequence the second time it appears, and worse at a token that breaks an earlier pattern, such as the first "Mrs Potter" after two occurrences of "Mrs Dursley". The section "Model Details" gives a slightly different window for the small models, approximately 1 to 3 billion tokens, and Argument 1 estimates that for large models the window falls about 1 to 2% of the way through training.

The authors list the confounds themselves (section "Assessing the Evidence"). The full-scale models offer only 15 points in time. Something else, such as the general ability to compose heads across layers, could form at the same moment and drive both changes. And a constant score after the phase change does not mean the mechanism stays constant, since gaining the same 0.4 nats from a better starting point is probably harder. They conclude that the case is stronger for the moment of the phase change than for the end of training.

Moving the phase change by changing the architecture (Argument 2)

An induction head needs each key to carry information about the token before it, which a standard transformer can only arrange with two layers. The smeared-key variant gives every head a trainable parameter that mixes each token's key with the previous token's key, so that a single layer can express an induction head (section "Argument 2"). As the hypothesis predicts, a one-layer smeared-key model develops in-context learning, which the ordinary one-layer model does not, and a two-layer smeared-key model develops it earlier than its control. The authors rate this as medium evidence that induction heads contribute some in-context learning and weak evidence that they contribute the majority, and they note that it does not rule out other mechanisms in larger models.

Ablations in small models (Argument 3)

Knocking out induction heads in small models greatly reduces the in-context learning score; in small attention-only models, "almost all" in-context learning appears to come from them, from the start of the phase change to the end of training (section "Argument 3"). A footnote reports that ablating most other heads appears to increase the score, which the authors explain as the model leaning on in-context learning for tokens it can no longer predict in other ways. For models with MLP layers they treat the ablation evidence as suggestive only, since heads and MLP layers can interact and two similar heads can mask each other.

Induction heads doing more abstract work (Argument 4)

In the 40-layer, 13-billion-parameter model, the authors show heads that meet the strict definition of an induction head on random repeated tokens and also do something more abstract (section "Argument 4"):

  • a literal copying head (copying score 0.89, prefix-matching score 0.75) that repeats names and sentences from earlier in a passage;
  • a translation head in layer 7 (copying 0.20, prefix matching 0.85) that, given the same sentence in English, French and German, attends to the word in the earlier language that will come next in the current one;
  • a pattern-matching head (copying 0.69, prefix matching 0.94) that, on synthetic lines of the form "month animal: 0", "month fruit: 1", "color animal: 2", "color fruit: 3", attends back to an earlier line of the same category and puts about 65% of its attention from the colons on correct positions.

The scores are from the table at the start of the section. The authors present these as anecdotes that make the general claim plausible, not as measurements of how much in-context learning induction heads perform. They describe the abstract behavior as a "fuzzy" or "nearest neighbor" version of induction, completing [A*][B*] … [A] with [B], where A* resembles A and B* resembles B.

Mechanism and continuity (Arguments 5 and 6)

For small attention-only models the mechanism is known from the earlier paper: a copying OV circuit plus K-composition with a previous-token head. In two-layer attention-only models the authors see no other mechanism that could supply the in-context learning (section "Argument 5"). They also report a second mechanism in GPT-2, which builds induction heads from positional embeddings, unavailable to their own models. For large models they argue from continuity: every model with more than one layer goes through the same phase change, with the same amounts of in-context learning before and after (section "Argument 6"). The most compelling alternative they name is that other composition mechanisms form at the same time and, in larger models with more heads, account for more of the effect than induction heads do.

Curiosities

The section "Unexplained Curiosities" collects observations the authors cannot explain:

  • After the phase change, the in-context learning score is about the same for a small two-layer model and the 13-billion-parameter model. Large models are better at every position, but they gain most of their advantage within the first ten tokens of the context.
  • Plotted against the logarithm of training tokens, the rate at which loss falls orders the models by size one way before the phase change and the opposite way after it (section "Phase Change Effect on Loss Derivatives").
  • Individual oddities include a non-induction head in the 6-layer attention-only model whose ablation looks like reversing the phase change, and an induction head in the 6-layer MLP model whose ablation has the opposite effect on in-context learning.

Safety framing

The authors give two reasons the work bears on safety (sections "Introduction" and "Safety Implications"). Abrupt changes in capability, during training or with scale, make it harder to prepare for problems, and this phase change is one that can be studied up close. In-context learning means a model's behavior can change during use without further training. On the concern that in-context learning might hide an internal optimizer (mesa-optimization), they report that the main mechanism, at least in small models, is induction heads, and "We did not observe any evidence of mesa-optimizers" (section "Safety Implications"). They also present the phase change as the first case they know of in which a phase change in machine learning has a mechanistic explanation, linking scaling-law-style measurements to circuits.

Limitations and open questions

  • Indirect evidence for large models. The paper's summary table rates the evidence that induction heads contribute the majority of in-context learning in large models as "Medium, Correlational", against "Strong, Causal" for small attention-only models. No ablations were run on the full-scale models.
  • A narrow measure. The in-context learning score measures how much later tokens are easier to predict. It does not isolate specific few-shot tasks, and the paper's link to task-style in-context learning rests on the examples in Argument 4.
  • Few snapshots for large models. With 14 or 15 snapshots, the timing of changes in large models is less certain.
  • Shared training order. Apart from the books-only runs, all models saw the same data in the same order, which the authors flag in a footnote as a limit on the dataset check.
  • Definition changes. A change log records that in September 2022 the authors corrected the wording of the prefix-matching definition in two places and added detail on how the copying evaluator uses a ReLU; a footnote also records that the copying evaluator omitted a start-of-sequence token by mistake.
  • Figures that differ between passages. The opening of Argument 4 says the authors study models up to 12 billion parameters, while the same section and the model table give the largest as 13 billion; the text places the pattern-matching head at layer 26, while the table at the start of the section lists it at layer 8 of 40. This note gives neither layer for that head.
  • Replications. Two external comments are appended. Adam Scherlis of Redwood Research reports induction heads reliably appearing in two-layer attention-only models; replacing the previous-token head's attention with exact previous-token attention recovered 99% of the loss difference, and a simple idealized induction pattern recovered about 65%. Tom Lieberum reports candidate induction heads in GPT-2 and GPT-Neo, mostly from mid-depth (section "Comments & Replications").

Research questions

  1. Do induction heads account for most in-context learning in a large model under ablation? The paper had no ablations at full scale. Running pattern-preserving ablations of the highest-scoring induction heads in one large model, and measuring the change in the in-context learning score, would test the central claim directly.
  2. Is the phase change driven by induction heads or by composition in general? The authors name this as the main alternative. Training models in which K- and Q-composition are available but previous-token information is blocked, and checking whether the bump still appears, would separate the two.
  3. Why is the in-context learning score nearly the same across model sizes? Measuring how much of the late-context loss reduction in large models comes from induction heads, against other mechanisms, would show whether the constant score hides a change of mechanism.
  4. How common are abstract induction heads? Argument 4 shows three hand-picked heads. Scoring every induction head of a large model on translation and category-matching tasks would give a rate rather than examples.
  5. What role do induction heads play in task-style few-shot prompting? The paper's measure is loss by position. Ablating induction heads while measuring accuracy on few-shot tasks would connect the two conceptions of in-context learning the paper distinguishes.
  6. Do the unexplained heads matter? The non-induction head whose ablation reverses the phase change, and the "anti-copying prefix-search" heads found in full-scale models above 16 layers, are listed but not studied. Characterizing their OV and QK circuits would show whether they are part of the same mechanism.

How it connects

References

  • Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. (2022). In-context Learning and Induction Heads. Transformer Circuits Thread, 8 March 2022. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html (also arXiv:2209.11895). Section names in this note refer to the web version.
  • Anthropic. (2022, March 8). In-context Learning and Induction Heads [research page]. https://www.anthropic.com/research/in-context-learning-and-induction-heads
  • Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. (2020). Scaling Laws for Neural Language Models. arXiv:2001.08361. Cited here only as the source of the paper's definition of in-context learning, as the paper itself states.