Attention, Fully Understood

After Module 1.1 you have a sequence of vectors: one per token, each holding the identity of that token and where it sits. Nothing else. The vector for cat is byte-for-byte the same in the cat sat on the mat and in the cat 5 file was corrupted.

That cannot be enough. To predict what follows …the store, John gave a drink to, the vector sitting at to has to somehow know that a name appeared earlier and which one it was. Information has to move between positions.

The mechanism is a soft dictionary lookup. In an ordinary hash map you hold a key, compare it against the stored keys, find the one that matches, and take that entry's value. Attention does the same thing with one change: instead of matching exactly one key, it scores every key for similarity and returns a weighted blend of all the values.

Three roles, three names, and each is just a linear projection of the residual-stream vector at that position:

The queries and keys live in a space of size $d_{\text{head}}$, typically much smaller than the residual stream's $d_{\text{model}}$ — 64 vs 768 in GPT-2 small. A head is a narrow window onto a wide stream, and that narrowness is a big part of why heads end up specialised.

Here is the whole of attention. Read it as: score every pair, shrink the scores, forbid looking ahead, normalise into a distribution, take the weighted average of the values.

Term by term. $QK^\top$ is an $n \times n$ matrix of scores: entry $(i,j)$ is $q_i \cdot k_j$, how well token $j$'s advertisement answers token $i$'s question. It is a plain dot product, so it is large when the two vectors point the same way and are long.

$M$ is the causal mask: 0 where $j \le i$ and $-\infty$ where $j > i$. Adding $-\infty$ before the softmax sends those entries to exactly zero probability. Then softmax runs along each row, so every token's weights over its visible predecessors sum to 1. Finally multiplying by $V$ takes the weighted average.

Why divide by $\sqrt{d_{\text{head}}}$

Suppose the entries of $q$ and $k$ are roughly independent with mean 0 and variance 1. Then $q \cdot k = \sum_{t=1}^{d} q_t k_t$ is a sum of $d$ independent terms each with variance 1, so

With $d_{\text{head}} = 64$ the scores would routinely sit around $\pm 8$, and gaps of 8 in logit space are enormous: softmax would be nearly one-hot at initialisation, gradients through it would be nearly zero, and training would stall before it started. Dividing by $\sqrt{d_{\text{head}}}$ puts the scores back at scale $\approx 1$, where softmax is soft and its gradient is healthy.

Why the mask

A language model must predict token $i+1$ from tokens $1..i$. Without the mask, position 3 could read position 7 and the training objective would be trivially cheatable. But there is a second, practical reason the mask is triangular rather than a loop over prefixes: one forward pass over a length-$n$ sequence produces $n$ training examples at once, because each row of the masked matrix already sees exactly the right prefix. That is a large part of why transformers train so much faster than RNNs.

Small enough to check every number with a pen. Sequence: the cat sat, residual stream width $d_{\text{model}} = 4$, one head with $d_{\text{head}} = 2$. The three input vectors, stacked as rows of $X$:

The head's three projection matrices, each $4 \times 2$:

Step 1 — project. Multiply each row of $X$ by each matrix. For sat, $q_3 = (1,1,0,1) W_Q = (1{\cdot}1 + 1{\cdot}0 + 0{\cdot}0 + 1{\cdot}1,\; 1{\cdot}0 + 1{\cdot}1 + 0{\cdot}1 + 1{\cdot}0) = (2,1)$. All nine vectors:

Step 2 — score. $S = QK^\top$, so $S_{ij} = q_i \cdot k_j$. For instance $S_{32} = (2,1)\cdot(1,2) = 4$:

Step 3 — scale and mask. Divide by $\sqrt{d_{\text{head}}} = \sqrt 2 \approx 1.414$ and set the upper triangle to $-\infty$:

Step 4 — softmax each row. Row 3, worked out: subtract the row max 3.54 to get $(-1.41, -0.71, 0)$, exponentiate to $(0.243, 0.493, 1)$, divide by the sum 1.736:

Step 5 — mix the values. $\mathrm{out} = A V$. Row 3: $0.140\,(1,2) + 0.284\,(2,1) + 0.576\,(1,0) = (1.284, 0.564)$.

Row 1 is exactly $v_1$: the first token can only attend to itself, so a head can never do anything useful at position 1. In real models this shows up as attention sinks — heads dump unwanted probability mass onto the first token when they have nothing to say. In production a $\langle$BOS$\rangle$ token is usually prepended precisely to give them somewhere harmless to point.

One head can express one routing rule at a time — the softmax forces its weights to compete for a single unit of probability mass. Real layers run several heads side by side. GPT-2 small has 12 heads per layer, each with $d_{\text{head}} = 64$, so the 12 heads together use exactly the same parameter budget as one head of width 768 would.

The textbook presentation concatenates the head outputs and multiplies by one big $W_O$. Split $W_O$ into per-head blocks and the same equation reads much more usefully:

A sum, not a tangle. Each head reads the residual stream through its own narrow projection, does its own routing, and adds its own contribution back. Heads in the same layer never see each other. That is what makes them tractable to study one at a time — and why the field says “head 5.1” the way a biologist says “this neuron.”

Some head types recur across models trained by different labs on different data — previous-token heads, duplicate-token heads, induction heads (which complete $[A][B]\ldots[A] \to [B]$, the star of Module 3.2), and name-mover heads (Module 3.5). This partial universality is one of the field's most encouraging findings: it suggests there is a shared set of algorithms to discover, not one private mess per model.

Two toys, top-down. The first shows you what attention patterns look like: four hand-designed heads, each imitating a head type that really occurs in trained models, over three sentences. The second shows you where a pattern comes from: drag vectors around a plane and watch geometry turn into probabilities.

These heads are hand-written scoring rules, not weights lifted from a model — the point is to make the shape of a pattern legible before you go stare at real ones (the problem set sends you to do exactly that).

Things to try: (1) On the Mary/John sentence, select the duplicate-token head and click the second John — that single arrow is the raw signal the IOI circuit uses to work out which name is not the answer. (2) Click a token near the start of any sentence and turn the causal mask off: watch mass flood into the future and the entropy jump. That is the model cheating, and it is why the mask exists. (3) In the QK playground, set the divisor to $\sqrt{64} = 8$ and watch the pattern collapse to nearly uniform — then crank query length up to 8× and watch it sharpen back. That is precisely the trade the $\sqrt d$ is managing: score magnitude grows with dimension, so the divisor grows with it too.