3.2 A Mathematical Framework & Induction Heads

Learn

A transformer is a sum, not a stack

Elhage et al.'s Mathematical Framework is the paper that turned attention from “a mechanism” into “an object you can do algebra on”. Its first move costs nothing and changes everything: stop thinking of the residual stream as a value that gets transformed, and start thinking of it as a running sum that every component writes into.

For an attention-only transformer with layers $\ell$ and heads $h$ , the final residual stream at a position is exactly

x_{\text{final}} = \underbrace{W_E t}_{\text{embedding}} + \sum_{\ell}\sum_{h} \underbrace{h^{\ell,h}(x)}_{\text{one head's write}}

and the logits are $W_U x_{\text{final}}$ . Because the unembedding is linear and the sum is a sum, the logits decompose into one term per component. Ask “how much did head 4 in layer 2 contribute to the logit for this token?” and there is a literal answer: apply $W_U$ to that head's output. That is direct logit attribution, and it exists only because of additivity.

Every head reads the whole stream and adds its output back. Heads in the same layer cannot see each other; heads in later layers see everything earlier ones wrote. The stream is a shared bus with no arbiter.

Key idea

Heads within a layer are independent: they read the same input, compute in parallel, and their outputs add. Nothing forces them to coordinate. This is why “head 5.1 does X” is even a coherent sentence — and why you can ablate one head and expect the others to keep working.

What the framework leaves out

Elhage et al. analyse attention-only models: zero, one and two layers, no MLPs. That is not modesty, it is strategy — MLPs are nonlinear and resist this algebra, and roughly two-thirds of a real transformer's parameters live in them. Everything in this module is exactly true of attention-only models and approximately, usefully true of real ones. Keep the asterisk.

Learn

QK and OV: where to look, what to bring

A head has four weight matrices — $W_Q, W_K, W_V, W_O$ — and the framework's second move is to notice they only ever appear in two pairs.

The attention scores use $W_Q$ and $W_K$ only through the product

\text{score}(i,j) = \frac{(W_Q x_i)^\top (W_K x_j)}{\sqrt{d_{\text{head}}}} = \frac{x_i^\top \big(W_Q^\top W_K\big) x_j}{\sqrt{d_{\text{head}}}}, \qquad W_{QK} \equiv W_Q^\top W_K

and the head's output uses $W_V$ and $W_O$ only through the product $W_{OV} \equiv W_O W_V$ . So the individual matrices are not the meaningful objects; the products are. This is not a notational nicety — it means the head's learned content lives in two matrices that map residual stream to residual stream, and you can look at them directly.

QK circuit

W_{QK} = W_Q^\top W_K

, a bilinear form on the residual stream. It answers “where should this position attend?”. Read as

W_E^\top W_{QK} W_E

it becomes a vocabulary × vocabulary matrix: which source tokens does this destination token want.

OV circuit

W_{OV} = W_O W_V

. It answers “what gets written to the destination, given that we attended there?”. Read as

W_U W_{OV} W_E

it is a vocabulary × vocabulary matrix: attending to token

j

raises the logit of token

i

by this much.

One head, two circuits, no interaction between them. The QK circuit reads the whole context to build a pattern; the OV circuit never sees the pattern and never knows which position it is moving.

Key idea

Freeze the attention pattern and a head becomes a linear map. All of the head's nonlinearity is in the softmax that produces

A

. That is why so much of interpretability is “stare at the pattern, then stare at

W_U W_{OV} W_E

” — between them they are the whole head.

Both circuits are severely low rank. In GPT-2 small, $d_{\text{model}} = 768$ but $d_{\text{head}} = 64$ , so $W_{QK}$ and $W_{OV}$ are 768×768 matrices of rank at most 64. A head cannot read or write arbitrary things; it gets a 64-dimensional slice of the stream. That constraint is what makes heads specialize, and it is a large part of why they are legible at all.

The copying test

If

W_U W_{OV} W_E

has large positive diagonal entries — attending to token

t

raises the logit of token

t

— the head is a copying head. Elhage et al. check this in a basis-free way with the eigenvalues of

W_U W_{OV} W_E

: mostly-positive eigenvalues mean the map tends to push in the direction of whatever it is given. You will run this test on a real model in the problem set.

Learn

Composition, and heads that do not exist

One layer of attention can only do so much. Its QK circuit reads token embeddings, so it can only decide where to look based on what the tokens are. Its OV circuit writes token information straight to the logits. The result is a skip-trigram model: patterns of the form “… $A$ … $B$ → $C$ ”, with a characteristic failure mode the paper enjoys pointing out — the head cannot make $C$ depend on $A$ and $B$ jointly, so a head that learns “keep… in mind” also fires “keep… in mind” on the wrong $A$ . Elhage et al. call these skip-trigram bugs, and they are direct evidence about the algorithm rather than the behaviour.

Two layers change the picture, because layer 1 reads a residual stream that layer 0 has already written into. There are exactly three places that can happen, one per input of the head:

Q-composition

H1's queries are computed from H0's output. Where H1 looks from depends on earlier processing, not only on the current token.

K-composition

H1's keys are computed from H0's output. A source position advertises something H0 computed about it rather than its own identity.

V-composition

H1's values are computed from H0's output. What gets moved is the result of an earlier movement.

Key idea

Only V-composition creates a virtual attention head: a composite with OV matrix

W_{OV}^{H1}W_{OV}^{H0}

and attention pattern

A^{H1}A^{H0}

, which behaves in the path expansion exactly like a single head that does not physically exist. Q- and K-composition do not add OV terms; they change where attention goes. Both matter enormously — the induction circuit is built from K-composition — but they are different kinds of thing.

Expand the two-layer attention-only model and every term is one of these paths:

\text{logits} = \underbrace{W_U W_E}_{\text{direct}} + \sum_{h}\underbrace{W_U W_{OV}^{h} W_E}_{\text{individual heads}} + \sum_{h_0, h_1}\underbrace{W_U W_{OV}^{h_1} W_{OV}^{h_0} W_E}_{\text{virtual heads}}

With $n$ heads per layer and 2 layers you get $n^2$ virtual heads on top of $2n$ real ones. This is where the combinatorics get frightening: the count of paths grows exponentially in depth, so “enumerate all circuits” stops being a plan almost immediately. In practice most virtual heads carry negligible weight, and finding the few that matter is the job.

Composition is a matter of degree

“Is there K-composition between these heads?” is not a yes/no question about wiring — every head reads the entire stream. The real question is how much of H1's key subspace overlaps with H0's output subspace, which the paper measures with a Frobenius-norm ratio of the composed matrices against the product of their norms. The toggles in the widget are a cartoon of a continuous quantity.

Learn

Induction heads: the circuit worth memorizing

Here is what two layers buy you, and it is the single most important concrete result in mechanistic interpretability so far.

An induction head implements: find an earlier occurrence of the current token, and predict whatever followed it. Given $[A][B] \ldots [A]$ it predicts $[B]$ . It takes two heads working together:

A previous-token head in layer 0 attends from each position to the one before it and copies that token's identity into the residual stream. Useless alone.
An induction head in layer 1 whose keys are computed from that written subspace — K-composition. Its key at position $p$ now means “the token before me was $T$ ”; its query at the destination means “my current token is $T$ ”. The match lands attention on the position after the earlier occurrence, and a copying OV circuit writes that token to the output.

This is why a one-layer attention-only model cannot do induction and a two-layer one can — a clean, falsifiable capability boundary that falls straight out of the algebra.

Key idea

Induction heads are the field's best evidence for universality. They appear in essentially every autoregressive transformer anyone has looked at, from two-layer toys to frontier models, and they appear at a sharply defined moment in training.

The induction bump: a small window early in training where loss drops faster than the surrounding trend and in-context learning ability appears. Olsson et al. show it coincides with induction heads forming. Schematic — the real curves are noisier and the bump is clearest when you plot the derivative or the in-context learning score.

Olsson et al. build a case, across six lines of evidence, that induction heads are the main source of in-context learning in transformers: the phase change in the loss curve coincides with induction-head formation; perturbing the architecture so induction heads form earlier or later moves the bump with them; ablating induction heads in small models removes most of the in-context learning; and per-head in-context-learning scores concentrate on induction heads. The authors are careful about the limits, and so should you be: the mechanistic story is demonstrated in small attention-only models and argued by correlation and analogy in large ones. Their own summary is that the evidence is strong but not conclusive at scale.

The heads found in large models are also not the crisp toy circuit. They do fuzzy matching: paraphrases, translations, and abstract pattern completion, not just literal token repeats. Whether that is “the same circuit, generalized” or a family of related mechanisms is unsettled.

in-context learning as fast weights

An induction head builds an associative memory at inference time out of the context: keys are “what preceded this”, values are “what came next”. Reading it is a lookup. That is a learned mapping, constructed on the fly, used once, and discarded — the same idea as the 1990s fast weights literature, implemented in attention rather than in a separate weight matrix.

Your on-the-fly learning thread starts here

You came to this course partly to understand models that learn without weight updates. This is the first mechanism in the course that actually does it — and the honest summary is that in-context learning is not one thing. Induction is the well-understood floor; above it sit “fuzzy” induction, task-vector effects, and (contested) claims that in-context learning implements gradient descent in the forward pass. Module 5.2 picks the thread up at the other end, where you edit the weights directly.

A learning channel nobody audits

If a model can acquire a behaviour from its context, then everything you established about the weights has a runtime escape hatch. Many-shot jailbreaking is the blunt demonstration: fill a long context with hundreds of examples of the model complying with harmful requests and the refusal training gives way — with effectiveness that grows smoothly with the number of examples, and grows faster in larger models. That is in-context learning working exactly as designed, pointed at your safety training. Weight audits do not see it; context-aware monitoring is a different and largely unsolved problem.

Explore

Build the circuit

First, walk the induction circuit end to end and watch the two attention patterns do their separate jobs. Then take the diagram apart: decide which of H1's three reads see H0's output, and watch which terms of the path expansion blink into existence.

The induction circuit, one step at a time

Two heads in two layers. Neither does anything interesting alone; composed through the residual stream they implement “repeat what followed this token last time”.

sequence

0

Mr

1

Dursley

2

was

3

the

4

director

5

of

6

Grunnings

7

,

8

said

9

Mr

→

?

layer 0 · previous-token head

layer 1 · induction head

step 1 / 5 · The task

The final token has appeared before. Somewhere earlier in the context is the answer to “what came next last time?”. Nothing in the weights knows this sequence — the pattern only exists in the context.

The random-token sequence is the control that matters. Those tokens never co-occur in training data, so no bigram statistic stored in the weights can produce the answer. The circuit works anyway, because it reads the pattern out of the context.

Composition builder: which circuits exist?

H1 reads the residual stream three times — once for queries, once for keys, once for values. Each read either sees only the embeddings, or sees the embeddings plus whatever H0 wrote. Those three independent choices are Q-, K- and V-composition.

Direct path

W_U W_E

present

Bigram statistics straight from the embedding. Always present, and in small models it does a surprising amount of the work.

H0's own OV path

W_U\, W_{OV}^{H0}\, W_E

present

H0 moves token information and it lands directly on the logits. A skip-trigram: “…A… B → C”.

H1's own OV path

W_U\, W_{OV}^{H1}\, W_E

present

Same, one layer up. Note it reads the embeddings, not H0's output — that is what V-composition would change.

Virtual attention head H0 → H1

W_U\, W_{OV}^{H1} W_{OV}^{H0}\, W_E \quad\text{with pattern } A^{H1}A^{H0}

absent

H1's values are computed from what H0 wrote, so the composite behaves like a single head whose OV matrix is the product and whose attention pattern is the product of the two patterns. This is the only kind of composition that creates a new head.

Q-composition

\big(W_{OV}^{H0}W_E\big)^{\!\top} W_{QK}^{H1}\, W_E

absent

H1's queries depend on H0's output: where H1 looks from now depends on earlier processing, not just on the current token.

K-composition

W_E^{\top}\, W_{QK}^{H1} \big(W_{OV}^{H0}W_E\big)

absent

H1's keys depend on H0's output: a position advertises not what it is, but something H0 computed about it.

With no composition, this is effectively two independent one-layer models added together. Every term is a bigram or skip-trigram; nothing here can look up a pattern in the context.

Virtual attention heads in this configuration: 0. Only V-composition makes one. Q- and K-composition are just as important — they are how the induction circuit works — but they change where H1 looks, not what it moves, so they do not appear as a new OV term in the path expansion.

Things to try: (1) Switch the induction visualizer to random tokens and step through again — the mechanism is identical, which is the proof that it is reading the context rather than recalling a bigram. (2) In the composition builder, hit Induction preset: K-composition only. Note that no virtual head appears, and yet this is the configuration that produces the field's canonical circuit — a useful antidote to the assumption that virtual heads are where the action is. (3) Turn on all three at once and count the terms; then imagine 12 heads per layer and 12 layers, and you will understand why circuit discovery needs automation.

Practice

Problem set

The pencil problems are the ones that make the paper readable; do them before you open it. The code problems are the field's standard first experiment — by the end you will have found induction heads in a real model and proved they matter by breaking them.

1.Shapes, ranks, and what that buys youpencil & paper

GPT-2 small: $d_{\text{model}} = 768$ , $d_{\text{head}} = 64$ , 12 heads per layer.

Give the shapes and maximum ranks of $W_{QK}$ and $W_{OV}$ for one head.
What are the shapes of $W_E^\top W_{QK} W_E$ and $W_U W_{OV} W_E$ , and what does each entry mean in words?
Why does it not matter that $W_Q$ and $W_K$ individually are not identifiable?

2.Counting pathspencil & paper

In a 2-layer attention-only model with $n$ heads per layer, the logits expand into a sum of terms.

How many terms are there in total, as a function of $n$ ? Break them down by type.
Generalize to $L$ layers: give the number of paths and say in one sentence why enumeration is not a research strategy.
A one-layer model has how many terms? What does that tell you about the class of functions it can express?

3.Design the induction headpencil & paper

Suppose the residual stream has two orthogonal subspaces: $S_{\text{tok}}$ holding the current token's embedding, and $S_{\text{prev}}$ , empty at layer 0 and written by the previous-token head.

Specify, in words and in matrix terms, what $W_{QK}$ and $W_{OV}$ of the layer-1 induction head must do. Then say what breaks if the previous-token head writes into $S_{\text{tok}}$ instead of a separate subspace.

4.Find the induction headscode

Load gpt2-small (or the 2-layer attn-only-2l model from the TransformerLens demos) and build a repeated-random-token sequence: pick 50 random token ids, concatenate the sequence with itself, prepend BOS. Run run_with_cache.

For every head, compute the induction score: the mean attention weight on the diagonal offset by seq_len - 1 in the second half of the sequence. Plot a layer × head heatmap. Separately compute a previous-token score (mean weight on the offset-1 diagonal).

Success check: a small number of heads have induction scores far above the rest, and at least one earlier-layer head has a high previous-token score. In GPT-2 small, heads 5.5 and 6.9 are commonly reported as strong induction heads and 4.11 as a strong previous-token head — treat those as a sanity check on your indexing, not as the answer, and trust your own numbers if they disagree.

5.Verify K-composition, then break itcode

Take your best induction head $H_1$ and your best previous-token head $H_0$ . Two experiments:

(a) Measure the composition. Compute the K-composition score $\|W_{QK}^{H_1\top} W_{OV}^{H_0}\|_F \,/\, (\|W_{QK}^{H_1}\|_F \|W_{OV}^{H_0}\|_F)$ and compare it against the same score for 20 randomly chosen earlier-layer heads.

(b) Break it. Mean-ablate $H_0$ 's output (replace it with its mean over a batch of prompts) and re-measure $H_1$ 's induction score and the model's loss on the repeated sequence.

Success check: the composition score for the real pair is a clear outlier, and ablating $H_0$ collapses $H_1$ 's induction score while ablating a random earlier head does not.

6.Measure in-context learning directlycode

Olsson et al. define an in-context learning score as the difference in loss between the 500th token of a context and the 50th, averaged over documents: how much better does the model predict once it has seen more of the document?

Compute it for gpt2-small on ~50 documents of ≥600 tokens. Then recompute it with your top induction head mean-ablated.

Success check: the base score is clearly negative (loss at token 500 is lower than at token 50), and ablating induction heads shrinks the magnitude measurably more than ablating a random head of the same layer.

0 of 6 problems marked done

Check

Check yourself

1.

Why do interpretability papers talk about

W_{QK} = W_Q^\top W_K

rather than about

W_Q

and

W_K

separately?

2.

Which kind of composition creates a virtual attention head?

3.

A one-layer attention-only model is shown [A][B] … [A] with tokens it has never seen adjacent in training. Can it predict [B]?

4.

In the induction circuit, what does the previous-token head contribute?

5.

Olsson et al. report a “phase change” early in training. What is the most defensible summary of the finding?

6.

Both

W_{QK}

and

W_{OV}

have rank at most

d_{\text{head}} = 64

in a 768-dimensional stream. What is the interpretability consequence?

7.

You want to claim head 5.5 is an induction head. Which single piece of evidence is strongest?

Answer all 7 questions to submit.

Submit your answers to complete this check.

Go deeper

The Mathematical Framework is dense and worth three sittings. Read it with the composition widget open. Nanda's walkthrough is the single best study aid for it — treat it as the lecture that accompanies the text.

EssentialA Mathematical Framework for Transformer Circuitspaper

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, et al. (Anthropic) · 2021 · 3 sittings

Sitting 1: the summary, 'Transformer Overview', and the zero-layer and one-layer sections — stop after the skip-trigram bugs, they are the best intuition pump in the paper. Sitting 2: two-layer models, composition, and induction heads. Sitting 3: the 'Summarizing OV/QK matrices' and 'Virtual weights' subsections, plus the appendix on notation. Skip the detailed model-specific results on a first pass. If the tensor-product notation slows you down, ignore it — every claim is also stated in ordinary matrix terms.

EssentialA Walkthrough of A Mathematical Framework for Transformer Circuitsvideo

Neel Nanda · 2022 · 3h (skimmable)

A co-author reading the paper aloud and explaining what each part is actually for, including which parts he thinks are over-engineered. Watch it alongside your second sitting, at 1.5× speed, pausing at the composition section. Worth more than a third re-read of the text.

EssentialIn-context Learning and Induction Headspaper

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, et al. (Anthropic) · 2022 · 2h

Read §1–4 as the curriculum suggests: the definitions, the phase change, and arguments 1–3. Then jump to the 'What we are not claiming' discussion, which is where the paper's intellectual honesty lives and where you learn how to calibrate a claim like this. The per-argument confidence table is a model for how to present uncertain evidence.

Transformer Circuits: exercisescourse

Anthropic · 2021 · 2h

Short exercises written to accompany the framework, including several on QK/OV algebra and composition. Do these instead of re-reading if the algebra has not landed — they are quick and they target exactly the confusions the paper produces.