%matplotlib inline
from d2l import torch as d2l
import torch
from torch.nn import functional as F10.6 What Attention Computes
The attention weights of a trained model look readable in a way no convolution kernel ever did: row \(i\) of the attention map says, in probabilities, where token \(i\) looked. The first generation of attention papers read the maps eagerly, and Section 10.3 already supplied a reason for caution — a diffuse row may mean “nothing matched” rather than “everything mattered”. This closing section asks the question behind the pictures: not where a model attends, but what the attending computes. For one family of models the question has a remarkably complete answer. In transformers stripped to attention alone (no feed-forward layers, no normalization), every head factors into two matrices that can be read off the checkpoint, information flow between layers can be traced term by term, and one particular two-layer circuit, the induction head, has been reverse-engineered end to end (Elhage et al. 2021). The TinyCharLM of Section 10.4 is exactly such a model. We first develop the vocabulary (the residual stream, the QK and OV circuits) and derive what each depth can and cannot express. Then we train the model on sequences of repeated random tokens and watch the induction circuit assemble itself: abruptly, visibly in the attention maps, and checkably in the weights. The mechanism we find is the same one that has been implicated in in-context learning in large language models (Olsson et al. 2022).
%matplotlib inline
from d2l import jax as d2l
from flax import nnx
import jax
from jax import numpy as jnp
import optax10.6.1 The Residual Stream
10.6.1.1 Tokens as Vectors in a Shared Workspace
Write \(\mathbf{e}_{x_i} \in \mathbb{R}^d\) for the embedding of token \(x_i\) and set \(\mathbf{h}_i^{(0)} = \mathbf{e}_{x_i}\). Unrolling TinyCharLM’s forward pass, block \(\ell\) of \(L\) computes, for every position \(i\),
\[ \mathbf{h}_i^{(\ell)} = \mathbf{h}_i^{(\ell-1)} + \sum_{h=1}^{H} \sum_{j \leq i} \alpha_{ij}^{\ell h}\, \mathbf{W}_{\mathrm{OV}}^{\ell h}\, \mathbf{h}_j^{(\ell-1)}, \tag{10.6.1}\]
and the tied output head scores the final vector against every embedding, \(\mathbf{z}_i = \mathbf{E}\, \mathbf{h}_i^{(L)}\). The sequence of vectors \(\mathbf{h}_i^{(0)}, \ldots, \mathbf{h}_i^{(L)}\) is the residual stream of position \(i\): a \(d\)-dimensional workspace that begins as the token’s embedding and ends as the vector the output head reads (Figure 10.6.1). Nothing along the way overwrites it — layers only add. Each attention head reads from the streams of earlier positions and adds what it read, linearly transformed, into the stream of the current position. The stream is therefore a communication channel running through depth: whatever a head wants later layers or the output head to know, it must write into the stream, and whatever it writes stays there for every later reader.
One structural fact makes this model unusually transparent: once the attention patterns \(\alpha\) are fixed, the map from embeddings to logits in Equation 10.6.1 is linear — a sum of products of weight matrices. Expanding the recursion writes the logits as a sum over paths through the network, each path a readable chain of matrices, which is what makes complete mechanistic analysis possible. This is the setting of the mathematical framework of Elhage et al. (2021), whose exact results hold for attention-only transformers without feed-forward layers or normalization, and it is why our specimen was built without them: TinyCharLM contains embeddings, attention, and nothing else, so everything it learns must pass through the mechanism we are studying. (Real transformers interleave both; we return to what survives the generalization at the end of the section.)
10.6.1.2 Where and What: The QK and OV Circuits
Each head’s term in Equation 10.6.1 splits into two independent pieces. The attention weights come from softmax over scores that depend on the query and key projections only through their product, and the payload depends only on the value and output projections:
\[ s_{ij} = \frac{\mathbf{h}_i^\top \mathbf{W}_{\mathrm{QK}}\, \mathbf{h}_j}{\sqrt{d_h}}, \qquad \mathbf{W}_{\mathrm{QK}} = \mathbf{W}_q^\top \mathbf{W}_k, \qquad \mathbf{W}_{\mathrm{OV}} = \mathbf{W}_o \mathbf{W}_v. \tag{10.6.2}\]
The QK circuit \(\mathbf{W}_{\mathrm{QK}} \in \mathbb{R}^{d \times d}\) is a bilinear form on the stream: it decides where the head attends, and nothing else. The OV circuit \(\mathbf{W}_{\mathrm{OV}} \in \mathbb{R}^{d
\times d}\) decides what an attended stream contributes to the destination, and nothing else. The split matters because the individual projections \(\mathbf{W}_q, \mathbf{W}_k, \mathbf{W}_v, \mathbf{W}_o\) are not identifiable (replacing \(\mathbf{W}_q\) by \(\mathbf{R}\mathbf{W}_q\) and \(\mathbf{W}_k\) by \(\mathbf{R}^{-\top}\mathbf{W}_k\) changes nothing the model computes), while the two products are what the head actually is. Both are \(d \times d\) but of rank at most \(d_h = d/H\): a head can test and move only a \(d_h\)-dimensional slice of the stream. With \(d = 128\) and four heads, each head works through a rank-32 bottleneck. In TinyCharLM the four projections live in two fused layers — qkv stacks the query, key, and value maps, and proj holds the per-head output maps side by side — and neither layer carries a bias, so extracting a head’s circuits is a matter of slicing, and the slices are the head, exactly:
torch.manual_seed(0)
model = d2l.TinyCharLM(vocab_size=64, pos='rope')
D = model.token_emb.weight.shape[1]
d_h = D // model.num_heads
W = model.blks[0]['qkv'].weight # Rows stack W_q, W_k, W_v
W_q, W_k, W_v = W[:D], W[D:2 * D], W[2 * D:]
W_o = model.blks[0]['proj'].weight # Columns split by head
h = 0
W_QK = W_q[h * d_h:(h + 1) * d_h].T @ W_k[h * d_h:(h + 1) * d_h]
W_OV = W_o[:, h * d_h:(h + 1) * d_h] @ W_v[h * d_h:(h + 1) * d_h]
for name, M in (('W_QK', W_QK), ('W_OV', W_OV)):
print(f'{name}: shape {tuple(M.shape)}, '
f'rank {torch.linalg.matrix_rank(M).item()}')W_QK: shape (128, 128), rank 32
W_OV: shape (128, 128), rank 32model = d2l.TinyCharLM(vocab_size=64, pos='rope', rngs=nnx.Rngs(0))
D = model.token_emb.embedding[...].shape[1]
d_h = D // model.num_heads
W = model.blks[0]['qkv'].kernel[...] # Columns stack W_q, W_k, W_v
W_q, W_k, W_v = W[:, :D], W[:, D:2 * D], W[:, 2 * D:]
W_o = model.blks[0]['proj'].kernel[...] # Rows split by head
h = 0
W_QK = W_q[:, h * d_h:(h + 1) * d_h] @ W_k[:, h * d_h:(h + 1) * d_h].T
W_OV = W_v[:, h * d_h:(h + 1) * d_h] @ W_o[h * d_h:(h + 1) * d_h]
for name, M in (('W_QK', W_QK), ('W_OV', W_OV)):
print(f'{name}: shape {tuple(M.shape)}, '
f'rank {jnp.linalg.matrix_rank(M)}')W_QK: shape (128, 128), rank 32
W_OV: shape (128, 128), rank 32Positions enter this picture in exactly one place. RoPE rotates queries and keys before they meet, so the score gains a relative-position dial inside the QK circuit — \(s_{ij}\) acquires the rotation \(\mathbf{R}_{j-i}\) of Equation 10.4.4 — while values pass unrotated: position can influence where a head looks but never what it moves. That separation is one reason we train with pos='rope' below; the other is continuity, since RoPE is Section 10.4’s headline scheme and the default of the models this chapter builds toward.
10.6.1.3 Bigrams, Skip-Trigrams, and the Limits of One Layer
The path expansion tells us what depth buys before we run anything. With zero blocks, \(\mathbf{z}_i = \mathbf{E}\,\mathbf{e}_{x_i}\): the logits are a function of the current token only — a bigram model, and a constrained one. Because the output head is tied to the embedding, the logit of token \(b\) after token \(a\) is \(\mathbf{e}_b^\top \mathbf{e}_a\), the Gram matrix of the embeddings: symmetric and of rank at most \(d\), not an arbitrary \(|\mathcal{V}| \times |\mathcal{V}|\) lookup table. Training on text fills it with as much of the bigram statistics as that form can hold. With one block,
\[ \mathbf{z}_i = \mathbf{E}\,\mathbf{e}_{x_i} + \sum_{h=1}^{H} \sum_{j \leq i} \alpha_{ij}^{h}\; \mathbf{E}\,\mathbf{W}_{\mathrm{OV}}^{h}\,\mathbf{e}_{x_j}. \tag{10.6.3}\]
Each new term reads: if the QK circuit sends attention from the current token to some earlier token \(x_j\), then the OV circuit adds \(\mathbf{E}\,\mathbf{W}_{\mathrm{OV}}\,\mathbf{e}_{x_j}\) to the logits. Statements of this shape are called skip-trigrams — “when \([\mathrm{A}]\) appears somewhere before \([\mathrm{B}]\), boost token \([\mathrm{C}]\)” — and they are all a one-layer model has beyond bigrams. A useful special case is copying: attend to earlier occurrences of the current token and boost that same token’s logit, so text that repeats a word becomes more likely to repeat it again.
Now consider the task that will occupy the rest of this section: the sequence contains \([\mathrm{A}][\mathrm{B}] \ldots [\mathrm{A}]\), and at the second \([\mathrm{A}]\) the model should predict \([\mathrm{B}]\), continuing the pattern by finding what followed last time. No skip-trigram expresses this. The head would need to attend from the second \([\mathrm{A}]\) to the token after the earlier \([\mathrm{A}]\), but in Equation 10.6.3 the score \(\alpha_{ij}\) sees only \(\mathbf{e}_{x_i}\), \(\mathbf{e}_{x_j}\), and the offset: the key at position \(j\) carries no trace of its neighbor at \(j - 1\). Predicting \([\mathrm{B}]\) through content requires a key that announces its predecessor, and that is precisely what a second layer provides. A head in layer 1 attends to the previous token and writes its identity into the stream; a head in layer 2 can then match the query “I am \([\mathrm{A}]\)” against keys enriched with “I follow \([\mathrm{A}]\)”, land on position \(j\), and let its OV circuit copy \(x_j = [\mathrm{B}]\) upward (Figure 10.6.2). This two-hop circuit — a previous-token head composing with a match-and-copy head through the residual stream — is the induction head. Note what the argument establishes: two layers can express it, one layer cannot (except through position alone, a loophole we will meet shortly). Whether gradient descent actually finds the circuit is an empirical question, and the rest of the section answers it by experiment.
10.6.2 Copying and Induction
10.6.2.1 Repetition as a Task
To catch a mechanism in the act, pose a task that only this mechanism solves. We train on sequences of repeated random tokens: sample a pattern of tokens uniformly from a vocabulary of 64, then tile it until the sequence is full. Within one sequence the pattern repeats; across sequences, nothing does. There are no corpus statistics to absorb — every bigram is equally likely on average — so the first pass over a pattern is unpredictable by design, and every later pass is perfectly predictable for a model that can look things up in its own context. The gap between a model’s loss on the first copy and on later copies therefore measures exactly one ability, in-context retrieval, with none of the confounds of real text.
device = d2l.try_gpu()
def repeated_batch(batch_size, num_steps, vocab, Lmin, Lmax):
"""Random patterns of length Lmin..Lmax, tiled to num_steps tokens."""
L = torch.randint(Lmin, Lmax + 1, (batch_size,), device=device)
s = torch.randint(0, vocab, (batch_size, num_steps), device=device)
pos = torch.arange(num_steps, device=device)
return torch.gather(s, 1, pos[None, :] % L[:, None]), L
torch.manual_seed(0)
x, L = repeated_batch(batch_size=2, num_steps=16, vocab=10, Lmin=5, Lmax=8)
print(f'pattern lengths: {L.tolist()}')
print(x.cpu())pattern lengths: [6, 6]
tensor([[4, 0, 3, 4, 0, 0, 4, 0, 3, 4, 0, 0, 4, 0, 3, 4],
[1, 0, 3, 5, 0, 7, 1, 0, 3, 5, 0, 7, 1, 0, 3, 5]])def repeated_batch(key, batch_size, num_steps, vocab, Lmin, Lmax):
"""Random patterns of length Lmin..Lmax, tiled to num_steps tokens."""
key_L, key_s = jax.random.split(key)
L = jax.random.randint(key_L, (batch_size,), Lmin, Lmax + 1)
s = jax.random.randint(key_s, (batch_size, num_steps), 0, vocab)
pos = jnp.arange(num_steps)
return jnp.take_along_axis(s, pos[None, :] % L[:, None], axis=1), L
x, L = repeated_batch(jax.random.key(0), batch_size=2, num_steps=16,
vocab=10, Lmin=5, Lmax=8)
print(f'pattern lengths: {L.tolist()}')
print(x)pattern lengths: [8, 6]
[[2 2 6 5 1 7 2 0 2 2 6 5 1 7 2 0]
[1 5 8 1 4 1 1 5 8 1 4 1 1 5 8 1]]The training loop is the fixed-step loop of Section 9.6 with one change: we log the loss separately on the first copy (target positions before the pattern length) and on the rest of the sequence, because the split is the measurement. Sequences are 64 tokens long and every batch is freshly sampled — the model never sees the same pattern twice.
def split_losses(logits, Y, L):
"""Cross-entropy on the first copy, on later copies, and overall."""
ce = F.cross_entropy(logits.reshape(-1, logits.shape[-1]),
Y.reshape(-1), reduction='none').reshape(Y.shape)
rep = torch.arange(1, Y.shape[1] + 1, device=device) >= L[:, None]
return ce[~rep].mean(), ce[rep].mean(), ce.mean()
def train_repeat(model, num_steps, Lmin, Lmax, batch_size=128, vocab=64):
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3,
weight_decay=0.0)
history = []
for step in range(num_steps):
x, L = repeated_batch(batch_size, 64, vocab, Lmin, Lmax)
first, rep, loss = split_losses(model(x[:, :-1]), x[:, 1:], L)
optimizer.zero_grad()
loss.backward()
optimizer.step()
history.append((first.item(), rep.item()))
return history
@torch.no_grad()
def eval_copy(model, Lmin, Lmax, batch_size=256, vocab=64):
x, L = repeated_batch(batch_size, 64, vocab, Lmin, Lmax)
logits = model(x[:, :-1])
first, rep, _ = split_losses(logits, x[:, 1:], L)
mask = torch.arange(1, 64, device=device) >= L[:, None]
acc = (logits.argmax(-1) == x[:, 1:])[mask].float().mean()
return first.item(), rep.item(), acc.item()def split_losses(logits, Y, L):
"""Cross-entropy on the first copy, on later copies, and overall."""
ce = optax.softmax_cross_entropy_with_integer_labels(
logits.reshape(-1, logits.shape[-1]), Y.reshape(-1)).reshape(Y.shape)
rep = jnp.arange(1, Y.shape[1] + 1)[None, :] >= L[:, None]
first = (ce * ~rep).sum() / (~rep).sum()
return first, (ce * rep).sum() / rep.sum(), ce.mean()
def train_repeat(model, num_steps, Lmin, Lmax, batch_size=128, vocab=64):
optimizer = nnx.Optimizer(model, optax.adamw(1e-3, weight_decay=0.0),
wrt=nnx.Param)
@nnx.jit
def step_fn(model, optimizer, X, Y, L):
def loss_fn(model):
first, rep, loss = split_losses(model(X), Y, L)
return loss, (first, rep)
(_, (first, rep)), grads = nnx.value_and_grad(
loss_fn, has_aux=True)(model)
optimizer.update(model, grads)
return first, rep
key, history = jax.random.key(0), []
for step in range(num_steps):
key, sub = jax.random.split(key)
x, L = repeated_batch(sub, batch_size, 64, vocab, Lmin, Lmax)
first, rep = step_fn(model, optimizer, x[:, :-1], x[:, 1:], L)
history.append((float(first), float(rep)))
return history
def eval_copy(model, Lmin, Lmax, batch_size=256, vocab=64):
x, L = repeated_batch(jax.random.key(7), batch_size, 64, vocab,
Lmin, Lmax)
logits = model(x[:, :-1])
first, rep, _ = split_losses(logits, x[:, 1:], L)
mask = jnp.arange(1, 64)[None, :] >= L[:, None]
acc = ((logits.argmax(-1) == x[:, 1:]) * mask).sum() / mask.sum()
return float(first), float(rep), float(acc)10.6.2.2 The Positional Shortcut
The most natural version of the task concatenates each pattern with itself: pattern length fixed at 32, sequence length 64. Our expressiveness analysis said one layer cannot run the two-hop algorithm — so let us falsify something by training a one-block model on it. (One equipment choice, whose reason will be clear in a moment: this model, alone in the section, gets its projection biases back with bias=True.)
torch.manual_seed(0)
model_shortcut = d2l.TinyCharLM(vocab_size=64, num_blks=1, pos='rope',
bias=True)
train_repeat(model_shortcut, 800, Lmin=32, Lmax=32)
first, rep, acc = eval_copy(model_shortcut, 32, 32)
print(f'first copy {first:.2f} nats, second copy {rep:.2f} nats, '
f'second-copy accuracy {acc:.2f}')first copy 4.42 nats, second copy 0.09 nats, second-copy accuracy 1.00model_shortcut = d2l.TinyCharLM(vocab_size=64, num_blks=1, pos='rope',
bias=True, rngs=nnx.Rngs(0))
train_repeat(model_shortcut, 800, Lmin=32, Lmax=32)
first, rep, acc = eval_copy(model_shortcut, 32, 32)
print(f'first copy {first:.2f} nats, second copy {rep:.2f} nats, '
f'second-copy accuracy {acc:.2f}')first copy 4.41 nats, second copy 0.09 nats, second-copy accuracy 1.00The one-block model matches the second copy essentially perfectly. Did we just refute the composition argument? No. We mis-posed the task. When the pattern length is always 32, the correct source for every prediction sits exactly 31 positions back, and a RoPE head can attend to a fixed relative offset without reading content at all (Section 10.4). No matching, no composition — just “look back 31”. The probe below evaluates the same model on patterns of length 24 and measures where its attention mass actually sits:
first, rep, acc = eval_copy(model_shortcut, 24, 24)
torch.manual_seed(1)
x, L = repeated_batch(256, 64, 64, 24, 24)
m = model_shortcut.attention_weights(x[:, :-1])[0]
print(f'patterns of length 24: second-copy accuracy {acc:.2f}')
print(f'attention mass at offset 31: '
f'{m.diagonal(-31, dim1=2, dim2=3).mean():.2f}')
print(f'attention mass at offset 23: '
f'{m.diagonal(-23, dim1=2, dim2=3).mean():.2f}')patterns of length 24: second-copy accuracy 0.02
attention mass at offset 31: 0.96
attention mass at offset 23: 0.00first, rep, acc = eval_copy(model_shortcut, 24, 24)
x, L = repeated_batch(jax.random.key(1), 256, 64, 64, 24, 24)
m = model_shortcut.attention_weights(x[:, :-1])[0]
print(f'patterns of length 24: second-copy accuracy {acc:.2f}')
print(f'attention mass at offset 31: '
f'{jnp.diagonal(m, -31, axis1=2, axis2=3).mean():.2f}')
print(f'attention mass at offset 23: '
f'{jnp.diagonal(m, -23, axis1=2, axis2=3).mean():.2f}')patterns of length 24: second-copy accuracy 0.02attention mass at offset 31: 0.96attention mass at offset 23: 0.00Accuracy collapses to roughly chance, and the reason is on display: nearly all attention mass still sits at offset 31, almost none at the offset the new patterns require. The model learned look back exactly 31 positions — the cheapest circuit that fits fixed-period data, and a circuit that has nothing to do with induction. This also explains the equipment choice above. A pure-position head needs a query that ignores content, and the qkv bias supplies exactly that: a constant query–key pair whose RoPE-rotated score depends only on the offset, peaking at a fixed distance. With biases disabled the shortcut only half-forms (second-copy accuracy roughly 0.5–0.7 in our probe runs), which is why this demo re-enables them — while every model we analyze keeps biases off, so the QK/OV algebra above stays exact. This is a lesson about synthetic tasks worth stating once and remembering: a benchmark rewards the cheapest shortcut it admits, not the mechanism its designer had in mind. The repair is to make the shortcut impossible: sample the pattern length uniformly from 16 to 32, so no fixed offset works and only content-based matching finds the previous occurrence.
10.6.2.3 Two Blocks Learn to Look Things Up
With variable pattern lengths, we train the two-block model — the smallest one our analysis says can express induction — and the one-block model as its control, both for 2,000 steps:
torch.manual_seed(0)
model_two = d2l.TinyCharLM(vocab_size=64, num_blks=2, pos='rope')
hist_two = train_repeat(model_two, 2000, Lmin=16, Lmax=32)
torch.manual_seed(0)
model_one = d2l.TinyCharLM(vocab_size=64, num_blks=1, pos='rope')
hist_one = train_repeat(model_one, 2000, Lmin=16, Lmax=32)
for name, model in (('2 blocks', model_two), ('1 block', model_one)):
first, rep, acc = eval_copy(model, 16, 32)
print(f'{name}: first copy {first:.2f}, later copies {rep:.2f}, '
f'accuracy {acc:.2f}')2 blocks: first copy 4.50, later copies 0.45, accuracy 0.87
1 block: first copy 4.34, later copies 3.65, accuracy 0.11model_two = d2l.TinyCharLM(vocab_size=64, num_blks=2, pos='rope',
rngs=nnx.Rngs(0))
hist_two = train_repeat(model_two, 2000, Lmin=16, Lmax=32)
model_one = d2l.TinyCharLM(vocab_size=64, num_blks=1, pos='rope',
rngs=nnx.Rngs(0))
hist_one = train_repeat(model_one, 2000, Lmin=16, Lmax=32)
for name, model in (('2 blocks', model_two), ('1 block', model_one)):
first, rep, acc = eval_copy(model, 16, 32)
print(f'{name}: first copy {first:.2f}, later copies {rep:.2f}, '
f'accuracy {acc:.2f}')2 blocks: first copy 4.45, later copies 0.45, accuracy 0.87
1 block: first copy 4.31, later copies 3.67, accuracy 0.11Now the composition argument holds up. The two-block model drives its loss on later copies below half a nat and predicts roughly nine out of ten repeated tokens; the one-block model is stuck above three nats and fewer than one token in five — better than chance, because copying-style heads can at least concentrate probability on tokens present in the context, but nowhere near retrieval. Both models sit at chance on the first copy: at ln 64, around 4.2 nats, or a little above it, the overshoot being the price of betting on repetitions that a fresh random pattern keeps refusing to deliver. How the gap opens during training is at least as informative as the endpoint:
def smooth(vals, k=25):
return [sum(vals[i:i + k]) / k for i in range(0, len(vals) - k + 1, k)]
d2l.plot(list(range(0, 2000, 25)),
[smooth([h[1] for h in hist_two]),
smooth([h[1] for h in hist_one]),
smooth([h[0] for h in hist_two])],
'step', 'loss (nats)',
legend=['later copies, 2 blocks', 'later copies, 1 block',
'first copy, 2 blocks'], figsize=(5, 3))
The two-block curve does not descend gradually. It hugs the one-block curve for hundreds of steps — while the model learns only what one layer can learn — and then drops by a couple of nats within a window of one or two hundred steps: a phase change, after which the model behaves qualitatively differently. The pattern replicates in every run we tried, in both frameworks; when the drop comes does not: it shifts by hundreds of steps between seeds and frameworks, which is why we describe its shape and not its schedule. Olsson et al. (2022) observed the same cliff, at vastly larger scale, as a bump in the training-loss derivative of real language models, and traced it to the same event we can now go and verify directly: the abrupt formation of induction heads.
10.6.2.4 The Heads, Caught in the Act
TinyCharLM.attention_weights returns every head’s attention map. We feed one evaluation sequence with pattern length 24 and look at all eight heads:
torch.manual_seed(2)
x, L = repeated_batch(1, 64, 64, 24, 24)
maps = torch.stack(model_two.attention_weights(x[:, :-1]))
d2l.show_heatmaps(maps[:, 0], xlabel='key position',
ylabel='query position',
titles=[f'head {h}' for h in range(4)],
figsize=(9, 4.5), cmap='Blues')
x, L = repeated_batch(jax.random.key(2), 1, 64, 64, 24, 24)
maps = jnp.stack(model_two.attention_weights(x[:, :-1]))
d2l.show_heatmaps(maps[:, 0], xlabel='key position',
ylabel='query position',
titles=[f'head {h}' for h in range(4)],
figsize=(9, 4.5), cmap='Blues')
The top row is block 1, the bottom row block 2, and the two halves of the predicted circuit are both visible. In block 1, at least one head shows a sharp line one step below the diagonal — a previous-token head. In block 2, heads show a line displaced 23 steps below the diagonal, beginning where the second copy begins: attention from each query to the token after the query’s previous occurrence, which for period-24 patterns is exactly offset 23. That is the induction stripe. To quantify what the eye sees, we measure each head’s average mass on the previous token and on the induction target \(j = i + 1 - L\):
def head_scores(model, x, L):
maps, T = model.attention_weights(x[:, :-1]), x.shape[1] - 1
tgt = torch.arange(T, device=device)[None, :] + 1 - L[:, None]
valid = tgt >= 0
for b, m in enumerate(maps):
B, H = m.shape[:2]
prev = m.diagonal(-1, dim1=2, dim2=3).mean((0, 2))
mass = m.gather(3, tgt.clamp(min=0)[:, None, :, None]
.expand(B, H, T, 1)).squeeze(-1)
ind = (mass * valid[:, None]).sum((0, 2)) / valid.sum()
print(f'block {b + 1}: previous token '
+ ' '.join(f'{v:.2f}' for v in prev) + ' | induction target '
+ ' '.join(f'{v:.2f}' for v in ind))
torch.manual_seed(3)
x, L = repeated_batch(256, 64, 64, 16, 32)
head_scores(model_two, x, L)block 1: previous token 0.00 0.95 0.02 0.01 | induction target 0.00 0.00 0.01 0.00
block 2: previous token 0.04 0.05 0.04 0.04 | induction target 0.67 0.71 0.66 0.65def head_scores(model, x, L):
maps, T = model.attention_weights(x[:, :-1]), x.shape[1] - 1
tgt = jnp.arange(T)[None, :] + 1 - L[:, None]
valid = tgt >= 0
for b, m in enumerate(maps):
B, H = m.shape[:2]
prev = jnp.diagonal(m, -1, axis1=2, axis2=3).mean((0, 2))
mass = jnp.take_along_axis(
m, jnp.broadcast_to(jnp.clip(tgt, 0)[:, None, :, None],
(B, H, T, 1)), axis=3).squeeze(-1)
ind = (mass * valid[:, None]).sum((0, 2)) / valid.sum()
print(f'block {b + 1}: previous token '
+ ' '.join(f'{v:.2f}' for v in prev) + ' | induction target '
+ ' '.join(f'{v:.2f}' for v in ind))
x, L = repeated_batch(jax.random.key(3), 256, 64, 64, 16, 32)
head_scores(model_two, x, L)block 1: previous token 0.02 0.01 0.92 0.03 | induction target 0.01 0.00 0.00 0.01
block 2: previous token 0.04 0.04 0.04 0.04 | induction target 0.62 0.68 0.69 0.49The division of labor across the two blocks is stark. Block 1 contains at least one head that spends well over half of its attention on the previous token and essentially none on the induction target; block 2’s heads do the reverse, with the strongest putting well over half of its mass on the single position the algorithm calls for, out of up to 63 candidates. Which head plays which role varies from seed to seed, and in this small model several block-2 heads usually share the induction work rather than one doing it alone; the structure — previous-token attention below, induction attention above, never the other way around — is what replicates.
10.6.3 In-Context Learning as Pattern Completion
10.6.3.1 Completing Patterns It Has Never Seen
Step back and consider what the trained model does at evaluation time. Every test sequence is freshly sampled: the pattern it completes, and every adjacent pair inside that pattern, has never occurred in training. There is no association between tokens for the weights to have stored — what the weights store is an algorithm, match-and-copy, that binds tokens to their successors at inference time, inside the context window. That is in-context learning, in miniature: the model “learns” each new pattern from a single exposure, without a gradient step. Plotting accuracy per position for period-24 patterns shows the single-exposure character directly:
torch.manual_seed(4)
x, L = repeated_batch(512, 64, 64, 24, 24)
with torch.no_grad():
acc = (model_two(x[:, :-1]).argmax(-1) == x[:, 1:]).float().mean(0)
d2l.plot(torch.arange(1, 64), [acc.cpu()], 'target position', 'accuracy')
x, L = repeated_batch(jax.random.key(4), 512, 64, 64, 24, 24)
acc = (model_two(x[:, :-1]).argmax(-1) == x[:, 1:]).mean(0)
d2l.plot(jnp.arange(1, 64), [acc], 'target position', 'accuracy')
Accuracy is near zero across the whole first copy, stays low at position 24 itself — nothing in the context yet announces that repetition has begun, so the pattern’s restart is unpredictable in principle — and then jumps within a couple of positions to near-perfect for the rest of the sequence. One exposure to a pattern suffices; the second exposure is already being completed from memory. The residual imperfection has an instructive cause: when a token happens to occur twice inside one pattern with different successors, a single-token match is ambiguous, and our two-layer circuit matches on exactly one preceding token. Disambiguating would require matching on a longer prefix: deeper composition, more layers.
10.6.3.2 The Circuit Is in the Weights
The attention maps show where the heads look; the QK/OV factorization lets us check what they do — from the weights alone, no forward pass. If block 2’s heads implement copying, then attending to token \(a\) should raise the logit of token \(a\) itself. The claim is about
\[ C_{ab} = \mathbf{e}_b^\top \mathbf{W}_{\mathrm{OV}}\, \mathbf{e}_a, \tag{10.6.4}\]
the boost that attending to token \(a\) gives to the logit of token \(b\) under a head’s OV circuit: copying means each row of the \(64 \times 64\) matrix \(\mathbf{C}\) is largest on its diagonal. We sum the OV circuits of block 2’s heads (they share the copying work) and test exactly that:
E, D = model_two.token_emb.weight, 128
qkv, proj = model_two.blks[1]['qkv'], model_two.blks[1]['proj']
OV = sum(proj.weight[:, h * 32:(h + 1) * 32]
@ qkv.weight[2 * D + h * 32:2 * D + (h + 1) * 32]
for h in range(4))
C = (E @ OV.T @ E.T).detach()
frac = (C.argmax(-1) == torch.arange(64, device=device)).float().mean()
print(f'rows whose largest entry is the diagonal: {frac:.2f} '
f'(chance 1/64)')
d2l.show_heatmaps(C[None, None].cpu(), xlabel='logit boosted',
ylabel='token attended to', figsize=(3.6, 3.6),
cmap='Blues')rows whose largest entry is the diagonal: 0.94 (chance 1/64)
E, D = model_two.token_emb.embedding[...], 128
qkv = model_two.blks[1]['qkv'].kernel[...]
proj = model_two.blks[1]['proj'].kernel[...]
OV = sum(qkv[:, 2 * D + h * 32:2 * D + (h + 1) * 32]
@ proj[h * 32:(h + 1) * 32] for h in range(4))
C = E @ OV @ E.T
frac = (C.argmax(-1) == jnp.arange(64)).mean()
print(f'rows whose largest entry is the diagonal: {frac:.2f} '
f'(chance 1/64)')
d2l.show_heatmaps(C[None, None], xlabel='logit boosted',
ylabel='token attended to', figsize=(3.6, 3.6),
cmap='Blues')rows whose largest entry is the diagonal: 0.77 (chance 1/64)
A diagonal emerges from weights that were never told about copying: for most tokens — in some runs all of them — the logit most boosted by attending to a token is that token itself, where chance would manage one row in 64. This check is deliberately partial: it feeds the OV circuit raw embeddings, ignoring what block 1 added to the stream, and it says nothing about the QK side of the match (the exercises take the analysis further, composing block 1’s OV circuit into block 2’s QK circuit). Partial as it is, it closes a loop that attention maps alone cannot: the where said match, the weights say copy, and together they are the induction algorithm.
10.6.3.3 Induction Heads in the Wild
None of this would matter much if it stopped at 64 tokens. It does not. Olsson et al. (2022) found induction heads in real transformer language models across sizes, using a probe nearly identical to our lab: score each head by its attention from the current token to the token after that token’s previous occurrence, on repeated random sequences. Three of their observations map directly onto what we just built. First, induction heads form abruptly early in training, and the formation coincides with the visible bump in the loss curve, our phase change at scale. Second, the same window is when models gain most of their in-context learning ability, measured as the gap between loss late and early in the context; ablating induction heads after training removes a large part of that gap. Third, the heads generalize off-distribution: heads identified on repeated random tokens also complete structured patterns, translate copied text fragments, and support few-shot-like completion. Prefix matching is a general algorithm, not a repetition trick. In large models the story cannot be as clean as in ours — feed-forward layers and normalization sit between the reads and the writes, the analysis becomes approximate, and attribution is contested at the margins — but the mechanism you can fully dissect in TinyCharLM is recognizably the one operating when a chatbot picks up a name, a format, or a made-up word from earlier in your conversation and reuses it correctly.
10.6.4 What Attention Weights Do and Do Not Tell You
This section closes the chapter, so it should also close a promissory note the chapter has been carrying: attention weights look interpretable — what do they actually license? Our own lab supplies the calibration.
Weights answer where, never what. An attention map fixes the mixture \(\alpha_{ij}\), but the head’s effect runs through its OV circuit: a head attending sharply to a token might move its identity, some feature of it, or, if the OV circuit annihilates the relevant directions, nothing at all. Conversely a diffuse map can implement precise computation (Section 10.3’s averaging construction did). The observation is not hypothetical: Jain and Wallace (2019) showed that on many tasks one can find quite different attention distributions that leave a model’s predictions unchanged, and the rejoinder of Wiegreffe and Pinter (2019) — that adversarially chosen alternatives do not refute every explanatory use — sharpened rather than settled the debate. “The attention weights mean something” is, as an unqualified claim, false; qualified versions must say what the weights feed into.
What our identification of the induction head actually rested on is worth listing, because the list is the method. A behavioral result: second-copy loss collapses, and only for models deep enough to express the circuit. A causal handle: change the input distribution (the period probe) or the model (the ablation exercise) and performance moves as the mechanism predicts. A weight-level check: the OV circuit is a copying matrix. The attention stripes were the least of it: they told us where to look, not what we had found. In our attention-only model all three checks were cheap because the model is linear once the patterns are fixed. In a real transformer, feed-forward layers and normalization break that linearity, features are packed into shared directions rather than neat subspaces, and each step of the triangulation becomes a research problem — this is the active field of mechanistic interpretability, where a handful of circuits, induction heads first among them, are understood at the level we reached here, while most of what large models do is not. Attention weights are evidence, and this section is a template for what turning evidence into understanding takes.
10.6.5 Summary
An attention-only transformer is best read through its residual stream: each position carries a \(d\)-dimensional vector from embedding to logits, and attention heads are the only writers. Every head factors into a QK circuit \(\mathbf{W}_q^\top \mathbf{W}_k\) that decides where to attend and an OV circuit \(\mathbf{W}_o \mathbf{W}_v\) that decides what the attended stream contributes, two rank-\(d_h\) matrices readable from the checkpoint. Depth buys expressiveness in discrete steps: zero layers store bigram statistics in the low-rank Gram form the tied head allows, one layer adds skip-trigrams (including copying), and two layers can compose a previous-token head with a match-and-copy head into an induction head that finds the previous occurrence of the current token and predicts what followed it. Trained on repeated random tokens with variable period, TinyCharLM discovers exactly this circuit — abruptly, in a phase change visible in the second-copy loss — while a one-block control cannot, and a fixed-period version of the task is solved by a positional shortcut instead (one that a projection bias, restored for that demo alone, makes cheap to express), a reminder that synthetic benchmarks reward the cheapest circuit they admit. The trained model completes patterns it has never seen, which is in-context learning in miniature, and the mechanism matches the induction heads implicated in in-context learning in large language models. Attention maps alone license none of these conclusions: the identification rested on behavior, causal probes, and weight-level checks together.
10.6.6 Exercises
- Build a copying head by hand — the skip-trigram “\([\mathrm{A}] \ldots [\mathrm{A}] \to [\mathrm{A}]\)” — with no training at all. Instantiate
TinyCharLM(64, num_blks=1, pos='none'), pick head 0, and overwrite its slices of the fusedqkvweight: set the query and key maps to the same random \(d_h \times d\) matrix scaled by a large constant (so that the score \(\mathbf{e}_a^\top \mathbf{W}_{\mathrm{QK}} \mathbf{e}_b\) is large precisely when \(a = b\)), set the value map to another random matrix and the head’s output columns to its transpose, and zero the output columns of the remaining heads. Verify on random sequences that the tokens with the most-boosted logits are those present in the context. Why does a random projection preserve the matching structure? (Recall the near-orthogonality of random high-dimensional vectors, Section 23.1.) Explain why this construction, whatever the constants, cannot beat the one-block plateau of the variable-period experiment. - Ablate the circuit. Zero out each block-2 head of
model_twoin turn (its columns ofprojsuffice — why?) and measure the second-copy loss; then ablate the two strongest heads together. How much of the ability survives single-head ablation, and what does that say about redundancy? Compare with ablating block-1 heads. - Retrain the variable-period experiment with
pos='none'andpos='alibi'. Which schemes still develop induction heads, and how does the phase change move? Before running, predict: the circuit needs a previous-token head in block 1 — which positional schemes make attending to a fixed offset of \(-1\) easy, and what does ALiBi’s recency bias do to the long-range match in block 2? - Complete the weight-level analysis with K-composition. The induction head’s match works because block 2’s QK circuit reads what block 1’s OV circuit wrote: for candidate head pairs, compute the \(64 \times 64\) matrix \(\mathbf{E} \mathbf{W}_{q}^\top \mathbf{W}_{k}
\mathbf{W}_{\mathrm{OV}}^{(1)} \mathbf{E}^\top\) (block-2 query/key maps, block-1 OV circuit) and compare its diagonal concentration against the direct path \(\mathbf{E} \mathbf{W}_{q}^\top \mathbf{W}_{k}
\mathbf{E}^\top\). Which pairs of heads carry the match, and do they agree with the attention-based scores of
head_scores? - Evaluate
model_twoon sequences that repeat a length-21 pattern three times. Is the third copy predicted better than the second? Connect your answer to the ambiguity discussed under pattern completion: what does a second previous occurrence do to a single-token match? - Retrain
model_twodrawing training patterns only from tokens 0–47, then evaluate on patterns drawn from tokens 48–63. Completion fails badly. Explain why, tracing the failure through Equation 10.6.2 and Equation 10.6.4: which matrices in the matching and copying paths involve embeddings that never received a gradient? Large models’ induction heads generalize to arbitrary tokens — what do their inputs pass through before reaching the heads that our model lacks?