week 06 / 12
Opening the box: induction heads
Real transformers do their work via attention-head circuits.
works through arena 1.2 · 1.2 Intro to Mech Interp
new terms this week · 4
Interpretability foundations
- Mechanistic interpretability
- The project of explaining model behavior by identifying the internal algorithms and components that cause it.
Tools
- TransformerLens
- A library that exposes transformer activations and hooks for mechanistic interpretability work.
Circuit analysis
- Induction head
- An attention head that implements a "copy what followed this token last time" pattern.
Causal methods
- Ablation
- Removing or zeroing a component to test whether behavior depends on it.
The puzzle: a model that learns from gibberish
Take fifty random tokens (genuine noise sampled uniformly from the vocabulary) and paste the same fifty tokens again. Then inspect the per-token loss. The random first copy offers no pattern and the model performs poorly. On the second copy, the loss falls off a cliff. The model predicts newly generated gibberish with near-perfect accuracy.
You invented the sequence thirty seconds ago, which rules out memorization. At inference time, GPT-2 noticed "this stretch of tokens happened before, and last time it was followed by that" and used the prompt as a lookup table. You can open the model, find the two specific attention heads responsible, and read their algorithm from the attention patterns. The mechanism is called an induction head, and finding it is the week's project.
From building models to reading models
Mechanistic interpretability asks: which internal components caused a completion, and what algorithm are they implementing? Week 6 is the course pivot from construction to investigation.
Separate behavioral from mechanistic questions. "Does GPT-2 complete repeated text correctly?" is behavioral; inputs and outputs answer it while the model stays a black box. "Which heads implement that behavior, and what algorithm do they compute?" is mechanistic; weights and activations answer it. Mechanistic interpretability is the decompiler dream: a trained network's weights are a compiled program nobody wrote, and the goal is to recover its source code.
Last week you built GPT-2, so you know its inventory: 12 layers × 12 heads = 144 attention heads, plus 12 MLPs, all reading and writing the residual stream. A circuit is a small team of components that implements a legible algorithm. Your first target is an induction circuit.
TransformerLens gives every activation a name
A bare PyTorch forward pass discards its intermediate tensors. TransformerLens wraps a transformer with hooks and an activation cache, letting you keep those tensors alongside the final logits.
# no-run
# 01: run a prompt, keep every activation
from transformer_lens import HookedTransformer
model = HookedTransformer.from_pretrained("gpt2-small")
logits, cache = model.run_with_cache("The cat sat on the mat")
print(cache["pattern", 0].shape)
# torch.Size([1, 12, 7, 7]): batch, heads, destination token, source token
# (7 = the 6 prompt tokens plus the BOS token TransformerLens prepends)
Think of the cache as a dictionary whose keys are stable probe points inside GPT-2. cache["pattern", 0] holds the attention patterns of every head in layer 0. Other keys expose the residual stream around each block, each head's output, and each MLP's activations. Two verbs cover the workflow: run_with_cache records tensors for inspection, while hooks change a tensor during the forward pass.
The cell's shape ends in 7, 7 for a 6-token prompt because TransformerLens prepends a beginning-of-sequence (BOS) token. Off-by-one confusion with BOS is this week's version of week 5's shape bugs.
Name both attention axes before indexing. In cache["pattern", layer][batch, head, destination, source], destination is the token doing the reading and source is the token being read. Swapping them transposes the story. Pair the matrix with token labels whenever you plot it; row 5 and column 5 mean little until both point back to words in the prompt.
Reading attention patterns
An attention pattern is a heatmap: each row is the token doing the reading, and each column is the token being read. Row values are the softmax weights from week 5. Each row sums to 1, and the causal mask zeroes everything above the diagonal, producing a lower-triangular heatmap.
Plot a few with the notebook's interactive circuitsvis widgets and a recurring zoo emerges. Previous-token heads put weight on the token immediately before, forming a crisp stripe below the diagonal. First-token heads park unused attention at the start of the sequence; the leading token provides a harmless "nothing to do here" target. Other heads attend to themselves or resist a simple label. A pattern suggests a hypothesis about where a head reads. Its output and an intervention must establish what the head does.
The notebook turns visual patterns into scores. A previous-token score averages the attention weight one step below the main diagonal. An induction score averages a different diagonal on a repeated sequence: each destination in the second copy should attend to the source immediately after its matching token in the first copy. Ranking heads by these scores narrows 144 heatmaps to a short list you can inspect and ablate.
The induction-head algorithm
Induction behavior appears on repeated text. If the sequence contains ... A B ... A, the model becomes better at predicting B. In Harry Potter went … Harry, the model at the final Harry should guess Potter. The prompt itself supplies that answer; book knowledge is unnecessary. A simplified two-head story follows:
- A previous-token head writes "what came before me" into each token's residual stream.
- An induction head at the later
Asearches for earlier positions whose previous-token information matchesA. - It attends to the earlier
Bposition and copies information that raises the logit forB.
Pure Python expresses the algorithm:
# 02: induction as eight lines of pure Python
tokens = ["Harry", "Potter", "went", "to", "class", ".", "Harry"]
current = tokens[-1]
for pos in range(len(tokens) - 1):
if tokens[pos] == current: # 1. match
prediction = tokens[pos + 1] # 2. shift
print(f"match at position {pos} → predict: {prediction}") # 3. copy
# match at position 0 → predict: Potter
This loop exposes the algorithm and hides the transformer's soft decisions. If a token appeared several times, the loop would print several candidates. An attention head can distribute weight across those matches, and its OV circuit can push several candidate logits by different amounts. The clean repeated-token experiment removes that ambiguity, making the diagonal stripe easy to measure before you return to natural text.
The transformer must implement each line in a particular place. A naive query from the final Harry matches keys made from Harry-like tokens and finds the earlier Harry. The prediction needs the following Potter. The model bridges that one-token shift by composing heads across layers:
The previous-token head in layer 0 stamps every position with a note saying "the token before me was X." The induction head in layer 1 builds its keys from those notes. When the final Harry asks "who has Harry in their before-me slot?", Potter answers. The match is already shifted one step forward, so the head attends to the token after the previous occurrence and raises that token's logit. Building keys from an earlier head's output is K-composition. The residual stream acts as a communication bus: layer 0 leaves a message, and layer 1 reads it.
That behavior is a seed of in-context learning: the model uses prompt patterns alongside memorized facts. In-context learning means improving at a task over one prompt through few-shot examples, copied formats, or names introduced and reused. Anthropic's induction-heads paper (linked below) argues that these copy heads drive a major part of the effect. They appear in a sudden phase change during training as in-context learning ability jumps.
Hooks are breakpoints you can write to
Everything so far is observation. To move from "this head looks at the right place" toward "this head causes the behavior", you need to intervene. A hook is middleware for a forward pass: a function TransformerLens calls at a named probe point, with the activation tensor as its argument. You can inspect an activation, replace it, or perform Ablation by zeroing a component and watching the loss change.
# no-run
# 03: a hook that deletes one head's work
# for the per-head output hook point "z", act has shape (batch, seq, head, d_head)
def zero_head_output(act, hook):
act[:, :, 0, :] = 0.0 # remove head 0's contribution
return act
Attach this hook to a suspected induction head and rerun the repeated-token prompt. A jump in second-half loss gives causal evidence beyond the suggestive heatmap. Watch the axis order: at hook points such as "pattern", the head axis comes directly after batch. Check the shape before indexing.
Measure three quantities: the untouched second-half loss, the loss after ablating a candidate head, and the loss after ablating a comparison head with a low induction score. A useful candidate causes a larger, targeted increase. Zero ablation is a blunt intervention because an all-zero activation may be unusual for the model. The simple test makes this first causal distinction; later weeks replace components with activations from matched prompts.
QK and OV: where to look vs what to copy
The notebook's final section reverse-engineers the circuit from weights using two linear maps. The QK circuit, built from query and key weights, decides where the head attends by scoring each (destination, source) pair. The OV circuit, built from value and output weights, decides what gets copied and how it changes logits. For an induction head, QK implements "match my token against every position's before-me note," and OV implements "copy the attended token toward the output." Learn cached activations and hooks first; revisit the QK/OV algebra after the session.
Why this matters
By the end of the notebook, you will have tested the claim that trained transformers contain discoverable circuits: a two-layer model, two named heads, a legible algorithm, and interventions that break it on demand. Later weeks scale this recipe. Week 9 turns single-head ablation into systematic activation patching on a 26-head circuit, and weeks 10–11 read and write the residual-stream messages that heads pass to each other.
Pair-session guide
Core work is ARENA 1.2 sections 1 and 2: use TransformerLens and find induction heads in a small two-layer attention-only model that ARENA loads for you. The hooks section later repeats the search on GPT-2-small. Section 3 supplies the intervention method required in week 9. Reverse-engineering QK/OV circuits is demanding stretch work.
What you should see
You should see an induction stripe: on repeated random tokens, an attention head forms a diagonal band at offset seq_len - 1, one less than the repeat length, because each repeated token attends to the token immediately after its previous occurrence. Per-token loss should also drop on the second half of the sequence. That is your first circuit-level "I found the mechanism" moment.
Where to go next
- Callum McDougall's Induction heads: illustrated is a diagram-dense post that makes the two-head algorithm click; ARENA recommends it alongside the induction circuits section of Neel Nanda's glossary.
- Anthropic's A Mathematical Framework for Transformer Circuits is the paper behind the QK/OV decomposition and the notebook's section 4. Read it after the session.
- Anthropic's In-context Learning and Induction Heads lays out the evidence that induction heads drive in-context learning, including the striking phase change during training.
this week's pair session
core
- 1-2: TransformerLens and finding induction heads
stretch
- Hooks
- Reverse-engineering QK/OV circuits