Module 6: The transformer

Module 04 ended with attention as a fix for the encoder–decoder bottleneck. At each output step the decoder took a weighted sum of the encoder states, with the weights computed from a score between its current state and each of them. In Bahdanau’s form the score is a small network:

where \mathbf{s}_{t-1} is the decoder state, \mathbf{h}_j the encoder state at position j and \mathbf{a}_t the context vector the decoder reads at step t. Luong and colleagues replaced the small network by multiplicative scores, the dot product \mathbf{s}^\top\mathbf{h}_j or the bilinear form \mathbf{s}^\top\mathbf{W}\mathbf{h}_j, which turn the scoring of every position into one matrix product. In both, attention was an addition to two recurrent networks. The transformer (Vaswani et al. 2017) removes the recurrence altogether and makes attention the only way positions communicate.

What recurrence costs

Recurrence has two costs that no amount of engineering removes.

It is sequential in time. Step t needs \mathbf{h}_{t-1}, which needs \mathbf{h}_{t-2}, and so on back to the first token. A training sequence of 10,000 tokens is therefore 10,000 dependent steps in every layer, whatever the hardware: a GPU can process many sequences side by side, but inside one sequence each step waits for the one before. Teacher forcing makes every input known in advance and still does not help, because the nonlinearity sits inside the loop.

Its paths are long. Information from position j reaches position t through t - j applications of the recurrence, each of which squeezes it through the same fixed-width state and can lose some of it. The vanishing gradients of Module 04 are the same long path, travelled backwards.

Self-attention removes both. Every position computes a weighted sum over every position it is allowed to see, and all positions do so at once, in one matrix multiply: a 10,000-token sequence becomes one large matrix operation instead of 10,000 small dependent ones. The path between any two positions is one layer long, whatever their distance.

The trade

Nothing is free. Vaswani et al. (2017, Table 1) compare the two kinds of layer on a sequence of T vectors of width d:

A recurrent layer multiplies a d-vector by a d \times d matrix at each of T steps; self-attention scores T^2 pairs of positions with dot products of length d. The ratio of the two costs is T^2 d / (T d^2) = T/d, so self-attention is the cheaper layer when T < d, as it was for the sentences of a few dozen tokens and the width d = 512 of the original translation model (a ratio of about 0.1 at T = 50). Beyond T = d it pays a quadratic price: at T = 32{,}768 and d = 4{,}096 the ratio is 8, and it grows linearly with the context. The table counts only the operation that mixes positions; a self-attention layer also pays O(Td^2) for its projections (Section 4), as a recurrent layer does for its weights. Module 04, Section 12 adds the convolutional alternative and the one thing recurrence keeps, a constant memory per generated token.

The T^2 is a memory cost as well: the T \times T matrix of weights. Section 10 shows how FlashAttention computes attention without ever storing that matrix, which removes the T^2 memory but not the T^2 arithmetic. The state-space models of Module 04, Section 13 take the other road back: a recurrence that is linear, and can therefore be computed in parallel.

Attention as a soft dictionary lookup

A Python dictionary is a hard lookup. Given {"pump": v1, "valve": v2, "tank": v3} and the query "tank", it compares the query with the keys, finds the one that is equal and returns that key’s value, v3, exactly. Any other query finds nothing.

Attention relaxes each of the three steps. It scores the query against every key with a dot product, so a key can match partially. It turns the scores into weights with a softmax, so every key receives a positive weight and the weights sum to 1. It returns the weighted average of the values, so the answer is a blend in which the best matches count most (Figure 6.1).

Two limits bracket this behaviour. Multiply all the scores by a factor c. As c grows the largest score dominates the softmax: at c = 10 the weights are (0.001, 0.001, 0.998) and the output is (2.996, 2.997), almost exactly \mathbf{v}_3, and in the limit the softmax becomes an argmax and the lookup is hard. As c \to 0, or whenever the scores are equal, every weight is 1/3 and the output is the plain average, (1.333, 1.667). Attention lives between the two, and where it sits is set by the size of the scores, which is why Section 2 cares about their scale.

import numpy as np

table = {"pump": (1, 0), "valve": (0, 2), "tank": (3, 3)}
print(table["tank"])                                 # hard: one exact match

q = np.array([1.0, 1.0])
K = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])   # keys: pump, valve, tank
V = np.array([[1.0, 0.0], [0.0, 2.0], [3.0, 3.0]])   # values
s = K @ q / np.sqrt(2)                               # scaled scores
w = np.exp(s - s.max())                              # subtracting the max avoids overflow
w = w / w.sum()                                      # softmax weights
print(np.round(w, 3), np.round(w @ V, 3))            # soft: a weighted average

(3, 3)
[0.248 0.248 0.503] [1.759 2.007]

The softness is what makes the lookup learnable. A dictionary offers no useful gradient: a small change to the query either changes nothing or jumps to another key. The soft lookup’s output is a smooth function of the query, of every key and of every value, so the gradient of a loss tells each of them which way to move. That is what allows the queries, keys and values to be produced by learned projections of the tokens (Section 2), and lets training decide what each position looks for and what it offers.

What must be added back

Removing the recurrence also removes two things it provided for free. The first is order. A set of weighted sums has no notion of position: shuffle the input tokens and every output is the same vector as before, moved to its token’s new place, so “the pressure exceeds the limit” and “the limit exceeds the pressure” produce the same set of outputs. Section 6 injects position, and Exercise 4 proves the symmetry. The second, for generation, is the arrow of time. A recurrent network cannot see its future, because its state at step t is built from steps 1 to t only; an attention layer sees every position at once, so a model trained to predict the next token needs a mask that stops each position from seeing its future (Section 2).

Take a sequence of T token vectors stacked as the rows of \mathbf{X} \in \R^{T\times d}. Project each into three roles with learned matrices:

with \mathbf{W}_Q, \mathbf{W}_K \in \R^{d\times d_k} and \mathbf{W}_V \in \R^{d\times d_v}. Row i of each is a projection of token i: \mathbf{q}_i = \mathbf{x}_i\mathbf{W}_Q, and likewise \mathbf{k}_i and \mathbf{v}_i. A query is what a position is looking for; a key is what a position offers to be matched against; a value is what it hands over if matched. Three matrices rather than one, because the roles differ: what a word looks for (a verb looking for its subject) need not resemble what it offers, and what it hands over need not be what made it match. Queries and keys share the width d_k so that they can be compared; values may have any width d_v. Then

\mathbf{S} = \mathbf{Q}\mathbf{K}^\top/\sqrt{d_k} is a T \times T matrix with one score per (query position, key position), S_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j/\sqrt{d_k}. The softmax is applied along each row, P_{ij} = \exp(S_{ij})/\sum_k \exp(S_{ik}), so \mathbf{P} = \softmax(\mathbf{S}) has non-negative rows that sum to 1: each query gets a probability distribution over the keys. Output row i is \mathbf{o}_i = \sum_j P_{ij}\mathbf{v}_j, a convex combination of the value rows, so it lies in their convex hull (Section 3 draws it). Figure 6.2 follows the shapes through.

The softmax Jacobian, applied to a row

Attention learns through its softmax, so the softmax’s derivative decides how fast it can learn. Module 02, Section 3 derived the Jacobian of \mathbf{p} = \softmax(\mathbf{s}):

Three of its properties do the work in this module. Every row sums to zero, \sum_j p_i(\delta_{ij} - p_j) = p_i - p_i = 0: adding the same constant to every score leaves \mathbf{p} unchanged. \mathbf{J} tends to zero as \mathbf{p} approaches one-hot, because every entry then contains a factor (p_i, p_j or 1 - p_i) that tends to zero. And \mathbf{J} is largest when \mathbf{p} is spread out: its trace, 1 - \sum_i p_i^2, peaks at the uniform distribution.

The gradient through one row. Let one row’s output be \mathbf{o} = \sum_j p_j\mathbf{v}_j and \mathcal{L} any scalar computed from it. Since \partial\mathbf{o}/\partial p_j = \mathbf{v}_j, the gradient with respect to each weight is g_j = \partial\mathcal{L}/\partial p_j = (\partial\mathcal{L}/\partial\mathbf{o})\cdot\mathbf{v}_j. The chain rule through \mathbf{J} then gives

g_j is the rate at which the loss changes as weight moves onto value j, and \bar g is the same rate averaged over the current mixture. Gradient descent therefore raises a key’s score when its value improves the loss more than the current weighted average does (g_j < \bar g), lowers it otherwise, and moves it in proportion to p_j, so a key that already has almost no weight barely moves. The scores pass the gradient on to the queries and keys through S_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j/\sqrt{d_k}; Section 3 does this with numbers.

Why the scores are scaled

Suppose the entries of \mathbf{q} and \mathbf{k} are independent, with mean 0 and variance 1. The score \mathbf{q}\cdot\mathbf{k} = \sum_{i=1}^{d_k} q_i k_i is a sum of d_k terms. Because q_i and k_i are independent, its mean is

The terms q_ik_i are independent of one another, so their variances add, and each has mean 0, so its variance is its second moment, which factorises by independence:

The standard deviation is \sqrt{d_k}: 11.3 at d_k = 128, a common head width. Dividing by \sqrt{d_k} restores unit variance, \operatorname{Var}(\mathbf{q}\cdot\mathbf{k}/\sqrt{d_k}) = d_k/d_k = 1, whatever the head width.

The assumption describes initialisation: a normalised input passed through projections initialised to preserve variance (Module 02, Section 6) has entries of roughly unit variance. Training is free to change the scale afterwards (larger \mathbf{W}_Q and \mathbf{W}_K sharpen a head that benefits from being sharp), so the factor sets the starting temperature of the softmax, not a limit on it.

What goes wrong without it. Unscaled scores at d_k = 128 have a standard deviation of 11.3, so gaps of ten or more between the best key and the rest are typical. The softmax is then saturated, nearly one-hot, where \mathbf{J} is nearly zero. Every gradient that reaches \mathbf{W}_Q and \mathbf{W}_K passes through \partial\mathcal{L}/\partial\mathbf{s} = \mathbf{J}\mathbf{g} (\mathbf{J} is symmetric, so it is its own transpose), so at initialisation those gradients are tiny whatever \mathbf{g} is: the model starts with an arbitrary, almost one-hot attention pattern and learns only slowly to change it.

Lab 1 measures the effect on random rows. At d_k = 128, with 16 keys and 5,000 draws, the unscaled softmax has a median largest weight of 0.978 and a mean entropy of 0.28 nats; scaled, the figures are 0.224 and 2.36 nats, against \ln 16 = 2.77 for a uniform row (Figure 6.3; the last digit depends on the random draw). Scaling is not the only defence: some large training runs also normalise \mathbf{q} and \mathbf{k} before the dot product (QK-normalisation), which bounds the scores outright; Module 08, Section 7 treats it as a stability device.

The causal mask

A model that generates left to right must not see what it is about to predict, so position i may attend only to positions j \le i. Set S_{ij} = -\infty for j > i before the softmax. Since \exp(-\infty) = 0, the weights on the future are exactly zero, not small, and no gradient flows through them either. The mask removes the T(T-1)/2 entries above the diagonal; row i keeps i scores.

The mask is what makes training efficient. Because the output at position i depends only on tokens 1, \dots, i, it can be trained to predict token i + 1, and every position does so at once: one forward pass over a T-token sequence yields T next-token predictions, each with its own loss. The decoder of Module 04, Section 10 ran step by step even under teacher forcing, because its recurrence did.

In code the mask is additive or boolean. An additive mask is a T \times T matrix holding 0 where attention is allowed and -\infty (or the most negative finite value of the dtype) where it is not, added to \mathbf{S}; a boolean mask holds True where attention is allowed. PyTorch’s F.scaled_dot_product_attention accepts either as attn_mask, or builds the causal mask itself with is_causal=True, which also lets a fused kernel skip the masked half (Section 10).

Padding masks. A batch of sequences of unequal length is padded to a common T. Padded key positions must be invisible to every query, so the padding mask is a boolean tensor of shape (B, 1, 1, T) that broadcasts over the heads and the query positions; combined with the causal mask it becomes (B, 1, T, T). The outputs at padded query positions are still computed, because tensors are rectangular, but they mean nothing and must be excluded from the loss, typically by giving their targets the loss function’s ignore_index.

import math, torch, torch.nn.functional as F

B, h, T, dk = 2, 8, 5, 32
q, k, v = (torch.randn(B, h, T, dk) for _ in range(3))
lengths = torch.tensor([5, 3])                         # sequence 2 ends in two pads
key_real = torch.arange(T) < lengths[:, None]          # (B, T)
pad = key_real[:, None, None, :]                       # (B, 1, 1, T)
causal = torch.ones(T, T, dtype=torch.bool).tril()     # (T, T), True = may attend
mask = causal & pad                                    # broadcasts to (B, 1, T, T)

y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)        # fused
s = q @ k.transpose(-2, -1) / math.sqrt(dk)                        # the same, by hand
y_hand = s.masked_fill(~mask, float("-inf")).softmax(dim=-1) @ v

Numerical safety. Compute the softmax as \exp(s_j - m)/\sum_k \exp(s_k - m) with m = \max_k s_k. Subtracting the same constant changes nothing and keeps every exponent at or below zero, so nothing overflows; Section 10 shows what happens otherwise. A row whose every key is masked is a different trap: every exponential is 0, and the softmax divides 0 by 0. A hand-written softmax returns NaN, which then spreads through every later layer; some fused kernels return zeros instead (PyTorch’s did in the version used for this module); and an additive mask of the most negative finite value gives every key the same score, so the row silently becomes the plain average of all the values, padding included. Rely on none of the three. The case arises naturally: in a left-padded batch under a causal mask, the padded positions at the start can see only padding. Make sure every query has at least one valid key, or zero those rows explicitly.

Three tokens, d_k = d_v = 2, and, to keep the arithmetic visible, the projected vectors taken directly instead of computed from embeddings:

The scores, and the scores divided by \sqrt 2 = 1.414:

Only three distinct scaled scores occur, so three exponentials serve for everything that follows: e^{0} = 1, e^{0.707} = 2.028 and e^{1.414} = 4.113.

The causal case

With a causal mask, row 1 sees only column 1, row 2 sees columns 1–2, and row 3 sees all three.

Token 3’s query (1, 1) matched its own key best, and the output is pulled toward its own value. Token 1 can only copy \mathbf{v}_1, whatever its query and key are.

Without the mask, and without the scale

The weights of rows 1 and 2 mirror each other, the first two entries swapped, because \mathbf{Q} = \mathbf{K} and swapping tokens 1 and 2 swaps both their queries and their keys: \mathbf{q}_1\cdot\mathbf{k}_1 = \mathbf{q}_2\cdot\mathbf{k}_2, \mathbf{q}_1\cdot\mathbf{k}_2 = \mathbf{q}_2\cdot\mathbf{k}_1 and \mathbf{q}_1\cdot\mathbf{k}_3 = \mathbf{q}_2\cdot\mathbf{k}_3. The outputs differ because \mathbf{v}_1 and \mathbf{v}_2 differ. Row 3 does not change at all: the mask removes nothing from the last row, which already sees every key.

Without the scale every score gap is \sqrt 2 times larger (in row 3, key 3 now leads the others by 1 instead of 0.707), so the weights sharpen and the outputs move further toward the best-matching value: row 3’s weight on \mathbf{v}_3 rises from 0.503 to 0.576 and its output from (1.759, 2.007) to (1.940, 2.152). Even at d_k = 2 the effect is visible. At d_k = 128, where unscaled scores spread eight times more than at d_k = 2 (11.3 against 1.4), it is the saturation of Section 2.

The geometry

Each output is a convex combination of the value rows its query may see: non-negative weights that sum to 1. So \mathbf{o}_1 is \mathbf{v}_1 itself; \mathbf{o}_2 = 0.330\,\mathbf{v}_1 + 0.670\,\mathbf{v}_2 lies on the segment from \mathbf{v}_1 to \mathbf{v}_2, two thirds of the way along; and \mathbf{o}_3 lies inside the triangle \mathbf{v}_1\mathbf{v}_2\mathbf{v}_3, nearest \mathbf{v}_3 (Figure 6.4). The scores only choose where in that region the output lands. Attention can mix values but never leave their convex hull: it cannot extrapolate, scale a value up, or produce a direction that no visible value contains. That is one reason every transformer block pairs it with a feed-forward network that transforms each position on its own, and with a residual path that keeps each token’s own vector (Section 5).

The backward pass, by hand

Take the scalar \mathcal{L} = o_{3,2}, the second component of token 3’s output (2.007), and push its gradient back through the causal, scaled computation. The general rules follow from \mathbf{O} = \mathbf{P}\mathbf{V}, the row rule of Section 2 and S_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j/\sqrt{d_k}:

with g_{ij} = (\partial\mathcal{L}/\partial\mathbf{o}_i)\cdot\mathbf{v}_j and \bar g_i = \sum_k P_{ik}g_{ik}. Masked entries have P_{ij} = 0 and receive no gradient. Here \partial\mathcal{L}/\partial\mathbf{O} has a single non-zero entry, a 1 in row 3, column 2, so only row 3 contributes.

The numbers say what the formula promised. Raising s_{33} moves weight toward \mathbf{v}_3, whose second component (3) is above the current 2.007; raising s_{31} moves it toward \mathbf{v}_1, whose second component is 0; key 2’s value, 2, sits almost exactly at the current average, so its score hardly matters. In the query, raising the second component increases the match with keys 2 and 3, whose values have large second components, at the expense of key 1: gradient 0.352. Raising the first component increases the match with keys 1 and 3 together, whose pulls (-0.498 and +0.500) cancel almost exactly: gradient 0.001. Lab 1 confirms every number with finite differences and with PyTorch’s autograd.

That is all attention does: a soft, learned lookup. Everything else is arranging many of them: in parallel heads (Section 4), interleaved with feed-forward networks in a block (Section 5), told about position (Section 6) and computed in tiles (Section 10).

One attention gives each query one probability distribution over the keys, and so captures one kind of relationship. Multi-head attention runs h attentions in parallel, each with its own projections into a smaller space of width d_k = d_v = d/h, concatenates their outputs and mixes them with one more matrix:

with \mathbf{W}_Q^{(i)}, \mathbf{W}_K^{(i)}, \mathbf{W}_V^{(i)} \in \R^{d\times d_k}, each \text{head}_i \in \R^{T\times d_k}, the semicolon denoting concatenation along the feature dimension, and \mathbf{W}_O \in \R^{d\times d}. With d = 4{,}096 and h = 32, each head works in 128 dimensions.

Why several heads. A softmax has one unit of weight to share out. A position that needs two pieces of information from two places (the previous word for its grammar, an earlier mention for a name) can get both from one head only as a weighted average, and Section 3 showed that an average is a point between the values, neither of them. With h heads a position can attend to h places for h different reasons. The cost does not grow: because hd_k = d, the projections are d \times d in total whatever h is, and the scores cost h \cdot T^2 \cdot d_k = T^2 d multiply-adds, the same as one head of full width. What changes is that each head compares queries and keys in a space of d_k dimensions instead of d.

The shapes, step by step

In code the heads are never separate objects. One d \times d matrix computes every head’s queries at once (the blocks \mathbf{W}_Q^{(i)} are its column blocks), a reshape splits the result by head, and the head index becomes a batch dimension. The explicit version below follows the attention class of Section 12 line by line, except that it forms the scores instead of calling the fused F.scaled_dot_product_attention:

import math, torch, torch.nn as nn

B, T, d, h = 2, 16, 256, 8
dk = d // h                                                    # 32
x = torch.randn(B, T, d)
wq, wk, wv, wo = (nn.Linear(d, d, bias=False) for _ in range(4))

q = wq(x).view(B, T, h, dk).transpose(1, 2)                    # (B, h, T, dk)
k = wk(x).view(B, T, h, dk).transpose(1, 2)
v = wv(x).view(B, T, h, dk).transpose(1, 2)
s = q @ k.transpose(-2, -1) / math.sqrt(dk)                    # (B, h, T, T)
future = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
p = s.masked_fill(future, float("-inf")).softmax(dim=-1)       # (B, h, T, T)
y = (p @ v).transpose(1, 2)                                    # (B, T, h, dk), not contiguous
out = wo(y.reshape(B, T, d))                                   # (B, T, d)

Step by step, for a batch of B sequences:

• wq(x) maps (B, T, d) to (B, T, d): all the heads’ queries side by side.
• .view(B, T, h, dk) splits the last dimension into h blocks of d_k; no data moves.
• .transpose(1, 2) gives (B, h, T, d_k): each head is now a separate batch entry.
• q @ k.transpose(-2, -1) gives the scores, (B, h, T, T). Matrix multiplication treats every leading dimension as a batch, so all B \times h score matrices come from one batched multiply.
• p @ v gives the per-head outputs, (B, h, T, d_k).
• .transpose(1, 2) returns to (B, T, h, d_k), and .reshape(B, T, d) concatenates the heads. It must be reshape, not view: after the transpose the memory is still laid out head by head, and merging the head dimensions violates view’s stride compatibility condition (here it raises a RuntimeError). Non-contiguous tensors can still support other compatible views. reshape returns a view when possible and copies when it has to.
• wo maps (B, T, d) to (B, T, d).

Concatenate, then project: a sum of per-head writes

Split \mathbf{W}_O into h blocks of d_k consecutive rows, \mathbf{W}_O^{(1)}, \dots, \mathbf{W}_O^{(h)}, each d_k \times d. Block matrix multiplication gives

because the entries of \text{head}_i multiply only the rows of \mathbf{W}_O that sit opposite them. Concatenation is bookkeeping: each head writes its d_k-dimensional result into the d-dimensional output through its own slice \mathbf{W}_O^{(i)}, and the writes add.

Two low-rank circuits per head

Write head i’s score between a query at position t and a key at position s in terms of the layer’s inputs:

This is a bilinear form in the two inputs with the d \times d matrix \mathbf{W}_Q^{(i)}\mathbf{W}_K^{(i)\top}, which passes through d_k dimensions and so has rank at most d_k. Elhage et al. (2021) call it the head’s QK circuit: it decides where the head looks. Likewise, the head writes into position t

the weighted inputs pushed through the d \times d matrix \mathbf{W}_V^{(i)}\mathbf{W}_O^{(i)}, again of rank at most d_k: the OV circuit, which decides what the head writes, given where it looks. Each head therefore reads from a d_k-dimensional subspace of its inputs and writes into a d_k-dimensional subspace of the output: 128 dimensions of 4,096 at d = 4{,}096 and h = 32. The two circuits are independent, so a head can choose where to look by one property of the tokens (position, say) and copy another (identity), which is how the heads described below work.

Counting

Each of the four projections (\mathbf{W}_Q, \mathbf{W}_K and \mathbf{W}_V for all heads together, and \mathbf{W}_O) is d \times d, so an attention layer has 4d^2 parameters, whatever h is. Per token, multiplying a d-vector by a d \times d matrix takes d^2 multiply-adds, or 2d^2 FLOPs (a multiply and an add each), so the four projections cost 8d^2 FLOPs. A token at context t, attending to t positions, pays 2td more for its row of \mathbf{Q}\mathbf{K}^\top (t dot products of length d_k in each of h heads, thd_k = td multiply-adds) and 2td for mixing the values: 4td in all. Summed over a sequence of T tokens that is O(T^2 d), quadratic in the length, the transformer’s best-known cost. Section 11 turns these counts into the FLOP convention the series uses.

The residual stream

Look at a whole model from the point of view of one position. Its vector starts as the token embedding, \mathbf{x}_0. Every attention layer and every feed-forward network reads from this vector through a normalisation and adds its output back:

where F_l is the l-th sublayer, attention or feed-forward (for attention, the read also includes the vectors of the positions it attends to). After the last layer, the logits are a linear read-out of the normalised final vector. Elhage et al. (2021) call this running sum the residual stream. Unrolled, \mathbf{x}_L = \mathbf{x}_0 + \sum_l F_l(\cdot): the final state is the embedding plus everything that every layer wrote.

Two consequences shape how transformers are understood. Layers communicate only through the stream: an attention head reads from the streams of the positions it attends to and writes into its own position’s stream, a feed-forward network reads and writes at one position, and nothing else passes between layers. And the stream’s d dimensions are a shared resource: every layer’s writes must fit into the same d numbers per position, alongside everything that earlier layers wrote and later layers will need. Section 5 writes the residual connection as the equation of a block and explains why the normalisation sits on the branch rather than on the stream.

What heads learn

Trained models contain heads with recognisable jobs. A previous-token head attends from each position t to t - 1. Analyses of trained models report heads that attend to the matching bracket, or from a word to the subject of its sentence; working out what heads do is a research field of its own.

The best-understood case is the induction head (Olsson et al. 2022), which completes a pattern “… A B … A” with B: having read “P-104 pump … P-104”, it predicts “pump”. It takes two heads in two layers. A previous-token head in an earlier layer writes into each position’s stream which token came before it; at the position of B, it writes “the token before me was A”. At the later A, the induction head’s query asks for a position whose previous token was A, and through its QK circuit it matches the key built from what the first head wrote at B. It attends there, and its OV circuit copies B’s identity into the stream, raising B’s logit. The second head’s key depends on information the first head moved. This particular circuit needs two sequential attention stages; it does not establish that every copying task is impossible for one layer.

Induction heads are a general mechanism for copying from context: names, identifiers, repeated phrases. Olsson et al. report that they form fairly abruptly, in a narrow window early in training that shows up as a bump in the loss curve, and that in-context learning improves at the same point. Lab 6 trains a two-layer model on sequences that repeat and finds both kinds of head (Figure 6.7), and Section 12 meets the same abrupt step in a character-level model that learns to copy an identifier.

Reading attention maps with care. An attention map shows where a head read information from, not why, and not what the layers above did with it. Jain and Wallace (2019) found that attention weights often disagree with other measures of which inputs mattered to a prediction, and that quite different attention patterns can give the same prediction. A claim about what a head does needs an intervention: zero the head’s output, measure the loss again, and see whether the behaviour attributed to the head goes away. Lab 6 does exactly this.