Induction Heads

Present a model with a repeated random sequence. At the second occurrence of a token it predicts what followed the first occurrence — with tokens it has never seen in training, so this is not memorisation but an algorithm implemented in the weights.

The circuit

Two heads in different layers, composing through the residual stream.

The previous-token head, in an early layer, attends from each position to the one before it and writes that token’s embedding into a subspace of the residual stream. The induction head, later, forms its query from the current token and matches it against those written-back embeddings — so it attends to the position after an earlier copy of the current token, and its OV circuit copies that position’s token to the output.

The match therefore fires when $x_i \approx x_{j-1}$ — the current token equals the token before position $j$ — and attention lands on $j$, whose content is then copied.

The phase change

Induction heads do not appear gradually. They form in a narrow window of training, and that window coincides with a visible bump in the loss curve and with the model’s abrupt acquisition of in-context learning. Ablating the heads after training removes most of the in-context learning gain.

A model that uses context well predicts the 500th token of a document much better than the 50th; the gap is the score, and it jumps when the heads form.

The strong version of the claim — that induction heads are the mechanism behind in-context learning generally, including few-shot task performance — is supported by ablation and correlation, not proved. The weak version, that they implement literal copying, is not in doubt.

import torch
from torch import Tensor


@torch.no_grad()
def induction_score(model, seq_len: int = 64, vocab: int = 1000) -> Tensor:
    """Mean attention paid to the token after the current token's
    previous occurrence. Returns [layers, heads]."""
    half = torch.randint(0, vocab, (1, seq_len))
    tokens = torch.cat([half, half], dim=1)      # exact repeat

    _, cache = model.run_with_cache(tokens)
    scores = []

    for layer in range(model.cfg.n_layers):
        # [batch, head, query, key]
        pattern = cache["pattern", layer]
        # from position i in the second copy, the induction target is
        # i - seq_len + 1: one after where this token appeared before
        diagonal = pattern.diagonal(offset=-(seq_len - 1), dim1=-2, dim2=-1)
        scores.append(diagonal.mean(dim=-1)[0])

    return torch.stack(scores)

The offset is the whole detector: -(seq_len - 1) selects, for each query in the repeated half, the key one position after the matching token in the first half. Off by one and you are measuring a previous-token head instead, which scores high on a different diagonal and is easy to mistake for a positive result.