Transformer Math

QK Circuit — Pattern Matching

The attention score between current position and past position determines whether the induction head "notices" a match. The query comes from the current token embedding; the key comes from the previous-token head's output — which encodes token[j−1]:

where is derived from token[i]'s embedding and encodes token[j−1]'s identity. So is large exactly when token[j−1] = token[i] — i.e., when position j is the token that follows the last occurrence of the current token.

OV Circuit — Completion Copying

Once the induction head attends to position j, the OV (Output-Value) circuit reads token[j]'s content and writes it to the output. Since j is the position after the previous occurrence of the current token, this copies the completion:

The matrix is the full OV circuit. Elhage et al. (2021)^ⁱ showed this can be analyzed directly as a single matrix that maps "what's at position j" to "what gets written into position i."

Layer Composition via Residual Stream

The two heads communicate through the residual stream. Layer 0 writes its output additively; layer 1 reads the updated stream:

The induction head in sees both the original token embedding andthe previous-token head's output in . This composition of two single-layer heads creates a two-layer circuit — the simplest possible form of inter-head communication.

The Induction Score — Detecting the Circuit

To identify induction heads, Olsson et al. feed a repeated random sequence into the model and inspect each head's attention matrix. An induction head will place most of its attention at the diagonal offset by — position attends to position , exactly where the current token last appeared:

where is the positions in the second copy and is head 's attention matrix. A score close to 1 means the head is attending precisely at the expected offset — a strong induction signature. In GPT-2 Small, ~5 heads score above 0.4^ⁱ, and those are the heads responsible for in-context learning.

PyTorch: Detecting Induction Heads via Attention Scores

import torch
import torch.nn.functional as F

def induction_score(attn_pattern: torch.Tensor) -> float:
    """
    Measure how strongly a head shows induction behavior.

    attn_pattern: (seq_len, seq_len) attention weight matrix on a
    repeated random sequence of length seq_len//2.

    An induction head attends at the [seq_len//2 - 1] diagonal:
    position i attends to position i - (seq_len//2 - 1), the spot
    where the current token appeared in the first copy.
    """
    seq_len = attn_pattern.shape[0]
    half = seq_len // 2

    # Extract the diagonal offset = -(half - 1)
    # i.e., for position i in second copy, attend to position i - half + 1
    offset = -(half - 1)
    diag = torch.diagonal(attn_pattern, offset=offset)
    return diag.mean().item()


def find_induction_heads(
    model,
    seq_len: int = 50,
    threshold: float = 0.4,
    device: str = "cpu"
) -> list[tuple[int, int]]:
    """
    Run a repeated random sequence through the model and return all
    (layer, head) pairs with induction score above threshold.
    """
    vocab_size = model.config.vocab_size
    n_layers = model.config.num_hidden_layers
    n_heads = model.config.num_attention_heads

    # Build a repeated random sequence: [A B C ... A B C ...]
    rand_tokens = torch.randint(1, vocab_size, (1, seq_len), device=device)
    tokens = torch.cat([rand_tokens, rand_tokens], dim=1)  # (1, 2*seq_len)

    with torch.no_grad():
        outputs = model(tokens, output_attentions=True)

    induction_heads = []
    for layer_idx, layer_attn in enumerate(outputs.attentions):
        # layer_attn: (batch, n_heads, seq, seq)
        for head_idx in range(n_heads):
            pattern = layer_attn[0, head_idx]  # (2*seq_len, 2*seq_len)
            score = induction_score(pattern)
            if score > threshold:
                induction_heads.append((layer_idx, head_idx))
                print(f"Layer {layer_idx}, Head {head_idx}: score={score:.3f}")

    return induction_heads