Attention: the one idea the whole architecture is built around
Every position in the sequence needs to gather information from the other positions — but which ones, and how much? Attention answers that with a mechanism that is entirely differentiable, so the model learns the answer instead of being told it. This chapter derives it, reads a real attention map out of a model you trained, and then has you implement it and prove your version correct against a reference.
The problem attention solves
Consider predicting the last word of: "The trophy would not fit in the suitcase because **it** was too large." To resolve it, the model must connect that position to trophy, nine tokens back — and in the near-identical sentence ending "too small", to suitcase instead. The relevant context is not at a fixed offset. It depends on the content.
Earlier architectures handled this badly. A recurrent network compressed everything seen so far into one fixed-size hidden state, so distant information had to survive many lossy rewrites; a convolution could only see a fixed window. Attention takes the opposite approach: keep every position available, and learn a content-based rule for which ones to read.
Queries, keys and values
Q, K and V are three different linear projections of the same input — three learned perspectives on each position, produced by three weight matrices of shape [d, d]:
Why three, rather than dotting positions against each other directly? Because the criterion for "should I read from you" and the content "here is what you get if you do" are different questions, and the model needs to be able to learn them independently. The query/key split further lets relevance be asymmetric: a pronoun can look for a noun without every noun looking for a pronoun.
The full formula
That single line is the entire mechanism. The rest of this chapter is what each piece is for, and what breaks when you get it wrong — which you can then verify by breaking it on purpose.
Every line, and what it is for
splitHeads — the meaning of "multi-head"
The d dimensions are cut into H contiguous slices of size d_head = d / H, and attention runs independently within each. This costs almost nothing — it is the same total arithmetic, just partitioned — and it buys the ability to track several relationships at once. One head can follow syntactic agreement while another tracks quotation marks, because they cannot interfere with each other until the output projection mixes them back at the end.
With one giant head, everything competes for the same subspace and a single softmax must choose one thing to look at. With too many tiny heads, each has too few dimensions to represent anything useful. GPT-2 small uses 12 heads of 64 dimensions; 64 per head has proved a remarkably durable choice across model scales.
bmm(q, k, transposed) — the score matrix
Every query is dotted with every key, producing a [T, T] matrix per head: entry (i, j) is how relevant position j is to position i. This is where the transformer's famous quadratic cost lives. Double the context and this matrix quadruples — which is exactly why long context windows were an unsolved engineering problem for years, and why FlashAttention (which never materialises the full matrix in memory) mattered so much.
scale — divide by √d_head
A dot product of two d_head-dimensional vectors with unit-variance entries has variance d_head, so its magnitude grows like √d_head. Feed numbers that large into softmax and it saturates: one weight goes to ~1, the rest to ~0, and the derivative of softmax at that point is nearly zero. The layer becomes a hard argmax that cannot learn — it looks like it is working while receiving no gradient at all.
| Scaling | Effect on softmax | Effect on learning |
|---|---|---|
| Correct 1/√d_head | Scores stay in a useful range | Gradients flow; attention patterns sharpen as training proceeds |
| None | Saturates to one-hot | Gradient vanishes through the softmax; the layer freezes |
| Too small | Flattens toward uniform | Attention averages everything and distinguishes nothing |
causalMask — the line that makes generation possible
The model is trained to predict token t+1 from tokens ≤ t, at every position simultaneously — that parallelism over positions is what makes training efficient. But it only works if position t genuinely cannot see beyond itself. The mask sets every score above the diagonal to a large negative number before the softmax, so those weights come out at effectively zero.
Mask before softmax, never after. Zeroing weights after the softmax leaves the remaining weights not summing to 1, which silently rescales the output. And use a large negative number rather than literal -Infinity: -Inf times a zero elsewhere in the computation produces NaN.
softmax — scores become a distribution
Softmax exponentiates and normalises, so each row of the attention matrix becomes a set of non-negative weights summing to 1. The implementation subtracts the row max first — mathematically a no-op, numerically essential, since exp(89) overflows float32 to Infinity and every downstream number becomes NaN.
bmm(att, v) and mergeHeads — the actual attending
Everything before this decided the blend; this line performs it. Each position walks away with a weighted average of the value vectors it cared about. The heads are then concatenated back into one [T, d] array and passed through an output projection W_o, which is the first point at which the heads are allowed to interact.
Reading a real attention map
The heatmap in the lab is not an illustration — it is the actual [T, T] softmax output from a forward pass of the model you just trained, on the probe text you typed. Row i is where position i looked; brighter means more weight. Some things to look for:
- The black upper triangle. That is the causal mask, made visible: physical proof the model cannot see ahead.
- A bright diagonal or near-diagonal. Very common: many heads attend mostly to the previous token or two. Local context is genuinely most of the signal in character-level text.
- Vertical stripes. A column that everything attends to. Often a delimiter — a newline, a space, or the first position, which frequently becomes an "attention sink" that heads park on when they have nothing to retrieve.
- Different heads doing visibly different things. Switch heads in the dropdown. If two heads look identical, the model has not needed to specialise them yet — at this size and training length, that is normal.
Implementing it, and proving it
The editor holds the real function the training engine calls. It is written against a whitelisted ops namespace of differentiable primitives, so the autograd system supplies correct gradients for whatever you compose — you can change the mathematics and still get a model that trains, rather than a crash.
Check against reference runs your implementation and the built-in one on identical random input and compares outputs elementwise, reporting the largest disagreement. This is the same technique used to test real numerical code, and it catches the failure mode that matters most here: an implementation that runs fine and is quietly wrong.
The four bugs everyone writes
- Masking after the softmax instead of before. Rows no longer sum to 1.
- Forgetting the transpose in
Q·Kᵀ. Shapes may still line up when T and d_head coincide, which makes it worse — it runs. - Dropping the scale factor. Trains, slowly and badly, with no error anywhere.
- Softmax over the wrong axis. Normalising down columns instead of across rows gives "how much did everyone attend to me", which is a meaningful quantity and not the one you want.
Vocabulary
- Query / Key / Value
- Three learned linear projections of the input: what a position wants, what it advertises, and what it offers.
- Attention weights
- The [T, T] post-softmax matrix; row i sums to 1 over the positions i is allowed to see.
- Causal mask
- Large negative scores above the diagonal, applied before softmax, so no position can read the future.
- Head
- One independent slice of the d dimensions running its own attention. H heads of size d/H.
- Scaling factor
- 1/√d_head, applied to scores to keep softmax out of saturation.
- Attention sink
- A position (often the first) that heads attend to when they have nothing specific to retrieve.
Check yourself
If you can answer these without re-reading, the lab below will make sense. If you cannot, the relevant section is worth a second pass — that is a better use of your time than clicking buttons.
- Why are there three projections instead of dotting the input against itself?
- What exactly goes wrong — mechanically, in the gradient — if you delete the 1/√d_head scaling?
- Why must the causal mask be applied before the softmax rather than after?
- Training loss drops far lower than usual after a change you made. Why is that a reason for suspicion?
- Where does the quadratic cost of a transformer come from, and what happens to it when context doubles?
- What does multi-head attention buy that a single wide head does not?
Where this comes from
- Vaswani et al., "Attention Is All You Need" (2017) — The original. §3.2 is the formula you just implemented.
- Elhage et al., "A Mathematical Framework for Transformer Circuits" (2021) — Where induction heads and the residual-stream view come from.
- Dao et al., "FlashAttention" (2022) — How the quadratic memory cost was made tractable without changing the mathematics.