Attention Heads Are Learned: Watching Circuits Form in a Tiny Transformer
An attention head is a learned lookup: a query vector at one position is compared against key vectors at every earlier position, the scores are softmaxed into a distribution, and the head returns the corresponding weighted average of value vectors. Nothing in that definition says what a head should look at. The matrices that produce queries and keys start as random noise, and every recognisable pattern a trained transformer displays — attending to the previous token, to the matching token, to position zero — is a fit to the training task, not a design decision.
This note assumes the mechanism itself — queries, keys, values, the 1/√d scaling, causal masking, multi-head splitting. Those are covered in Self-Attention: The Content-Addressed Lookup at the Heart of Transformers. What follows is the other half: what those matrices become once gradient descent is finished with them.
🔬 The model on the table
The transformer below is a real decoder-only model — token and positional embeddings, a stack of pre-LayerNorm blocks each holding multi-head causal self-attention and a GELU feed-forward layer, then a final LayerNorm and an unembedding into logits. It is trained with Adam on a cross-entropy next-token loss. The only unusual thing about it is the size: about 18,000 parameters at the default settings, which is small enough that the forward pass, the backward pass, and the optimiser all run in a few milliseconds of plain JavaScript.
Two structural facts drive everything below. First, attention is the only operation that moves information between positions — the feed-forward layers and LayerNorms act on each position independently. Second, a layer-0 key can only depend on its own token and its own position, because nothing has written to the residual stream yet. Any pattern that requires knowing what sits next to a key position therefore cannot exist in layer 0, and needs a second layer to read what the first one deposited. That constraint turns out to be visible in the demo.
🎮 The playground
Pick a task, pick a shape, and press run. Each square below is one head's attention matrix on a single held-out example: row i is the distribution the query at position i puts over key positions j, so brighter means more attention and the black upper triangle is the causal mask. The matrices are redrawn as training proceeds, and each head is labelled by whichever canonical pattern its attention mass best matches.
🗺️ Reading a head
A head's label is not a judgement call: each canonical pattern is a specific set of cells in the attention matrix, and the score is simply how much probability mass lands there, averaged over query positions and over the batch. The head is named after whichever score wins, provided the winner clears 0.4.
| Pattern | Cells scored | What it computes |
|---|---|---|
| previous-token | A[i, i−1] | Copies each position's predecessor into the residual stream. |
| induction | A[i, m+1] | Where m is the last earlier position holding the same token — so the head reads what followed that token last time. |
| matching | A[i, m] | Finds the earlier copy itself rather than its successor. |
| mirror | A[i, T−1−i] | The anti-diagonal — position i reads the position reflected about the centre. |
| attention sink | A[i, 0] | Dumps mass on position 0, a way of attending to nothing in particular. |
| offset −δ | A[i, i−δ] | A fixed stride, for the best δ. |
| position-j | A[·, j] | Always looks at one absolute position, whatever the content. |
| uniform | H(A[i, ·]) | Normalised entropy near 1 — a flat average over everything visible. |
🔁 Why induction needs two layers
The induction task is a sequence of distinct tokens followed by the same tokens again in a rotated order, so every token in the second half has exactly one earlier copy, and the answer at each position is whatever followed that copy the first time. Because the rotation is random the answer sits at no fixed position and no fixed offset — content matching is the only route.
Solving it requires composing two heads. A query at position i wants the position m+1, where m holds the same token as i. But matching on content finds m, not m+1, and the key at m+1 has no idea which token precedes it. So a head in layer 0 first copies each position's predecessor one step forward — the previous-token head. Now the key at m+1 carries m's identity, and a layer-1 head can match position i's token against it directly: the induction head.
Ablation makes the composition concrete. Switching off a single head with the ⊘ button zeroes that head's contribution while leaving everything else trained and intact. On the converged default model — two layers, two heads, seed 1, 2500 steps:
| Head switched off | Label | Held-out accuracy |
|---|---|---|
| — | intact | 99.7% |
| L0 H0 | previous-token head 0.42 | 22.8% |
| L1 H0 | induction head 0.81 | 22.4% |
| L0 H1 | no clear pattern | 35.7% |
| L1 H1 | attention sink 0.53 | 99.7% |
Both links in the chain are load-bearing and the sink is free. This is the same prev-token-plus-induction circuit that Elhage et al. (2021) identified in two-layer attention-only transformers, and that Olsson et al. (2022) tied to the sharp phase change in in-context learning that real language models show early in training. Getting it to appear here takes about eighteen thousand parameters.
📍 Position is an input you can remove
Self-attention is permutation-equivariant: a softmax over dot products has no notion of order, and position enters only because it is added to the embeddings. Turning positional embeddings off makes that dependence measurable, and different tasks respond very differently.
Accuracy after 2000 steps, averaged over eight seeds, with the per-seed values shown because the spread is the point:
| Task | With positions | Without | Per-seed, without |
|---|---|---|---|
| Reverse | 100% | 73% | 91 · 42 · 94 · 91 · 93 · 87 · 43 · 45 |
| Majority | 100% | 100% | 100 · 100 · 100 · 100 · 100 |
Majority does not notice, which is what a bag-of-tokens task should do. Reverse is the interesting one, and not because it collapses — it does not. The outcome is bimodal: some seeds recover to around 90%, others stall in the mid-40s, and the average of 73% describes no individual run. Even the bad runs sit far above the 1-in-6 chance floor.
Position is not actually gone when the embeddings are removed, because the causal mask leaks it. Position i can attend to exactly i+1 keys, so a uniform head returns an average carrying weight 1/(i+1) — a usable positional signal, and a uniform head is exactly what the model grows here (its uniformity score pins at 0.99). Removing positional embeddings makes order expensive and unreliable to recover rather than unavailable. Switching the causal mask off as well takes the last of it away.
🎯 Consequences worth remembering
- Heads are fitted, not specified. The same architecture grows a mirror head on reverse, a uniform counting head on majority, and a prev-token/induction pair on the copy task. Only the data changed.
- Run it more than once. Single runs here were misleading in both directions — the induction circuit appears in roughly a third of seeds, and removing positional embeddings gives anything from 42% to 94% on the same task. Every number in this note that comes from one run is labelled as such.
- Layer 0 keys are content-blind about their neighbours. Any pattern that needs to know what is adjacent to a key position costs an extra layer, which is why induction is a two-layer circuit and not a one-layer one.
- Ablation is the test that matters. A pretty attention picture is a hypothesis. Deleting the head and watching accuracy fall from 99.9% to 15.8% is evidence.
- Toy tasks leak. Three of this note's task designs were solvable by a shortcut that looked like the real circuit in the heatmap. The diagnostic that saved it was making a fixed-position detector compete with the induction score.
- Redundancy is normal. On majority, ablating any single head changes nothing — several heads compute the same average, so no one of them is necessary.
- Interpretability is not free. Every seed here learns the task; only some learn it in a form that can be read off a heatmap. A model being understandable is a property of the particular solution it landed in, not of the architecture.