Source-linked AI summary

In-context Learning and Induction Heads

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, Chris Olah

arXiv:2209.11895v1cs.LG

TL;DR

The paper asks whether induction heads explain the majority of in-context learning in transformers, especially given limited evidence connecting small and large models. It analyzes the phase change and induction-head circuits across model scales, finding strong causal evidence in small attention-only models and weaker correlational evidence in larger models. The authors conclude that induction heads may be the primary mechanism, while acknowledging that extrapolation to large models is uncertain.

  • Problem

    The paper investigates how confidently evidence from small models can establish that induction heads drive in-context learning in larger transformers.

  • Method

    The paper combines circuit analysis, training-dynamics comparisons, and cross-model evidence to relate induction-head formation to in-context learning.

  • Results

    Induction heads appear to be responsible for most in-context learning in small transformers, while the evidence is weaker for large transformers.

  • Takeaways & Limitations

    The induction-head phase change provides a bridge between mechanistic circuit formation and macroscopic learning dynamics.

  • Takeaways & Limitations

    Inferring large-model mechanisms from small models requires caution because larger models may use additional composition mechanisms.

Abstract

from arXiv · show

"Induction heads" are attention heads that implement a simple algorithm to complete token sequences like [A][B] ... [A] -> [B]. In this work, we present preliminary and indirect evidence for a hypothesis that induction heads might constitute the mechanism for the majority of all "in-context learning" in large transformer models (i.e. decreasing loss at increasing token indices). We find that induction heads develop at precisely the same point as a sudden sharp increase in in-context learning ability, visible as a bump in the training loss. We present six complementary lines of evidence, arguing that induction heads may be the mechanistic source of general in-context learning in transformer models of any size. For small attention-only models, we present strong, causal evidence; for larger models with MLPs, we present correlational evidence.

W ​V O

The paper argues that induction heads may explain most in-context learning, with strong causal evidence in small attention-only models and weaker correlational evidence in larger models. Their formation coincides with a sharp phase change in in-context learning, while alternative composition mechanisms remain possible.

  • Evidence across model sizes: The authors extrapolate from small models to argue that induction heads account for most in-context learning in larger transformers, but acknowledge weaker evidence there.The extrapolation is explicitly presented as an extended inference from the preceding arguments.
  • Limitations: Non-induction composition heads could become more important in sufficiently large models, so induction heads may not explain all phase-change improvements.Larger models have more capacity for composition mechanisms unavailable to small models.
  • Phase change: After the phase change, in-context learning varies only slightly across models, from tiny two-layer models to 13-billion-parameter models.The reported comparison emphasizes whether a model has undergone the phase change, rather than model size or additional training duration.
  • Broader significance: The phase change links macroscopic loss dynamics with circuit-level mechanistic explanations and is compared with grokking as a discontinuous training phenomenon.The paper presents this as a bridge between scaling laws and mechanistic interpretability.
  • Induction-head mechanism: Induction heads implement pattern copying, including [A][B]...[A]→[B] and longer repeated sequences, through QK and OV circuits.Their OV circuit copies tokens, while QK composition uses information from a previous-token head.

Change Log

The publication record documents corrections to prefix-matching definitions and clarifies how the copying-head evaluator measures effects on token logits.

  • Corrections: The change log corrects prefix-matching wording in two places, replacing an incorrect “preceded by the current token” formulation.The revised wording uses “preceded” or “were followed by”.
  • Clarifications: The change log adds that copying-head effects on raising logits are calculated using a ReLU.

Footnotes

The footnotes define induction-head behavior, explain why composition is required, and qualify the evidence and measurement procedures. They also discuss ablation interpretation and unresolved questions about learning dynamics.

  • Definitions and mechanism: Induction heads require at least two attention layers because their attended position depends on information written by an earlier head.A one-layer head cannot make its attention target depend on the preceding token of the attended token.
  • Measurement caveats: Induction-head properties are continuous in practice, although a subset of heads exhibits them substantially more often than chance.
  • Definitions and mechanism: Induction heads are identified using repeated random sequences to distinguish induction from simple previous-token copying.The behavioral definition supports testing whether heads rely on induction rather than a heuristic copying strategy.
  • Open questions: The paper notes that the balance between increasing in-context-learning capacity and the difficulty of later bits remains unexplained.
  • Measurement caveats: Ablating many non-induction heads can increase measured in-context learning because it may remove competing normal-prediction pathways.
  • Measurement caveats: In attention-only models, attention-head ablations provide a principled contribution measure because in-context learning originates in attention heads and logits are nearly additive.The direct path depends only on the present token and cannot contribute to in-context learning.
  • Measurement caveats: Ablation interpretation is harder with MLPs because in-context learning depends nonlinearly on which heads are removed.Ablations can shift MLP statistics or disrupt mechanisms that use multiple heads.

Citation Information

The work is cited as Olsson et al.’s “In-context Learning and Induction Heads,” published in the Transformer Circuits Thread in 2022.

  • Citation: The citation identifies the work as “In-context Learning and Induction Heads” by Olsson and colleagues, from the Transformer Circuits Thread, 2022.

Where induction heads form

Induction heads form at different depths depending on model architecture and scale. They appear in the last layer of attention-only models but earlier in larger models.

  • Small attention-only models: In small attention-only models, induction heads form in the last layer.
  • Small models with MLPs: By the 5L and 6L sizes, small models with MLPs form induction heads more often in the second-to-last layer.
  • Full-scale models: In the full-scale sweep, induction heads form earlier than in the smaller models.
  • Full-scale models: In 24L and 40L models, the majority of induction heads form before the halfway point in model depth.

Distribution of scores

Head activation scores are more sharply distributed in small models than in full-size models, and prefix matching aligns with copying, especially among high-scoring heads.

  • The score distribution is sharper in small models than in full-size models.
  • Heads with positive prefix matching scores are more likely to have positive copying scores.
  • The correlation between prefix matching and copying is stronger among heads with the highest prefix matching scores.

Validating head activation evaluators

The heuristic head evaluators track mathematically based measures of attention-head circuits in small attention-only models. Copying aligns with OV structure, while prefix matching aligns with the previous-token QK circuit.

  • The head activation evaluators correlate well with mathematically based measures in small attention-only models.
  • The copying evaluator correlates with each head’s OV matrix through a normalized eigenvalue-based measure.
  • The prefix matching evaluator correlates with the trace of QK eigenvalues for the previous-token QK-circuit term.

Data Collection

The analysis measures model behavior across training snapshots and evaluates attention heads using losses, ablations, heuristic activation scores, and mathematical circuit measures. These measurements separate whole-model behavior from individual-head functions and assess copying, prefix matching, and previous-token attention.

  • Per-token losses: Per-token losses are measured on 10,000 fixed examples of 512 tokens for every model snapshot.
  • Per-token losses: Losses are summarized both from one random token per example and by averaging losses at each context index.
  • Head ablations: Individual-head effects are tested by successively zeroing each head’s contribution to the residual stream and comparing model behavior.
  • Head ablations: Pattern-preserving ablation removes a head’s result while preserving its contribution to later attention patterns, isolating downstream V-calculation effects.
  • Head activation evaluators: Three empirical head evaluators score copying, prefix matching, and previous-token attention from observed activations on example data.
  • Head activation evaluators: Copying depends on a head’s OV circuit, whereas prefix matching and previous-token attention depend on its QK circuit.

Analyses

The analyses examine per-token loss across training, context position, model trajectories, attention-head measurements, and ablations. They also compare phase-change behavior with induction-head activity and quantify whether ablating heads reverses the observed behavioral change.

  • Per-token loss views: Per-token losses are analyzed across training snapshots and token positions to compare model performance earlier versus later in the context.The 2D context-index plot represents loss by token position and training time; horizontal slices are training-time loss curves, while vertical slices fix a token index.
  • Per-token loss views: The in-context learning score is the average loss at token 500 minus the average loss at token 50 for length-512 contexts.This single statistic summarizes how much better the model predicts later rather than earlier tokens over training.
  • Model trajectories: PCA projects concatenated per-token loss data onto two principal components to trace each model’s trajectory in function space over training.The matrix contains 200 snapshots for each of 12 small models and 15 snapshots for each of 6 full-scale models.
  • Attention head measurements: Attention-head analyses track prefix-matching activation scores and the prefix-matching QK-circuit eigenvalue trace over training.All heads are shown for small models, while 100 selected heads are shown for full-scale models.
  • Attention head measurements: The QK-circuit trajectories show a weight-decay change near 5 billion tokens, but this bend occurs after the phase change and is unrelated to it.The authors identify the bend as a separate phenomenon rather than part of the induction-head-associated transition.
  • Ablation attribution to phase change: For small models, ablation attribution compares each head’s loss change with the normalized before-and-after vector describing phase-change behavior.Induction-head ablations have a negative dot product, indicating that removing them is like undoing the behavioral changes that occurred during the phase change.
  • Ablation attribution to phase change: For induction heads, the negative ablation attribution indicates that removing them reverses changes in model behavior associated with the phase change.The ablation change vector is compared with the before-and-after vector using a dot product.
Loading 2209.11895v1…