Part II — Attention

Causal & multi-head attention

Lesson 05 gave us a trainable attention layer: three projections W_q, W_k, W_v and the derived 1/√d_k scale. But that layer has two fatal flaws for a language model — it lets a token read the future it is supposed to predict, and it computes only one notion of "relevant." This lesson fixes both with a causal mask and multiple heads, and adds dropout to keep the weights honest during training.

What the last step left broken

Look back at lesson 05's heatmap. Token 2 ("journey") attends freely to tokens 3, 4, 5, 6 — every one of which comes after it in the sentence. A model whose job is to predict the next token can simply copy the answer from a future position: it would score a perfect training loss by cheating and learn nothing generalizable, and at inference time — where the future genuinely does not exist yet — it would have nothing to read. On top of that, the layer holds exactly one set of W_q, W_k, W_v, so it can encode only one similarity geometry, when language relevance is many things at once (grammatical agreement, coreference, position).

Linear position

Forced by: trainable self-attention (05) lets every position attend to every other — including future tokens it must predict — and computes just one similarity view.
New idea: a causal mask sets the upper triangle of the score matrix to −∞ before softmax, so each position attends only to itself and the past; dropout randomly zeros attention weights during training to prevent co-adaptation; and multi-head attention splits d into h parallel heads of width d/h, giving h different views at essentially the same cost.
Forces next: attention only mixes information across positions — it applies no per-position nonlinear transform, and a naive deep stack of these layers will not train. Lesson 07 wraps attention in a feed-forward MLP, LayerNorm, and residual connections to make the block deep-stackable.

The plan

Five moves. (1) Pin down exactly why peeking at the future is fatal for a next-token predictor. (2) Build the causal mask and derive why we set scores to −∞ before softmax rather than zeroing weights afterward. (3) Add dropout on the attention weights and see why survivors must be rescaled. (4) Split one head into h heads — one geometry each — and account for the shapes and the cost. (5) Drive the mask, the heads, and the dropout in the widget and watch the future go dark.

1 · The peeking problem

Recall the whole point of the model from lesson 00: it computes p(x_t+1 | x₁ … x_t). The target at position t is the token that actually comes next, x_t+1 (lesson 03 gave us that label for free by shifting the sequence). Training slides this prediction across every position at once: position 1 predicts token 2, position 2 predicts token 3, and so on — all in a single forward pass, which is what makes training efficient.

Here is the danger. In lesson 05, the context vector for position t is a blend over all positions, including t+1, t+2, … Those later positions carry the very tokens we are asking the model to predict. If attention at position t can look at position t+1, the model can route the answer straight through and drive its training loss to zero without learning any real structure. This is label leakage, and it produces a model that is useless the moment you deploy it:

Why leakage is fatal, not just untidy

At inference the model generates one token at a time, left to right — when it is predicting token t+1, positions t+1 onward do not exist yet. A model that learned to lean on future positions during training has learned a feature that is simply absent at generation time, so its predictions collapse. Train-time and inference-time information must match exactly, which means position t may only ever see positions ≤ t.

So we need to enforce a hard rule inside attention: position t attends only to positions 1 … t. Everything strictly to its right must contribute exactly zero weight. This is called causal (or masked) attention — causal because information flows only forward in time, never backward from the future.

2 · The causal mask

Lay out the T×T score matrix from lesson 05, where entry S[i][j] is how much query i attends to key j. The rule "position i may only see j ≤ i" means every entry above the main diagonal (where j > i) must vanish. That upper triangle is exactly the set of "future" cells.

The clean way to kill those cells is to set them to negative infinity before the softmax:

S_masked[i][j] = { S[i][j] if j ≤ i ; −∞ if j > i }

Softmax exponentiates each score, and e^−∞ = 0, so every masked cell receives probability exactly 0. The surviving cells in each row are exponentiated and divided by their own sum, so each row still sums to 1 — but now only over the legal, past-and-present positions. Row 1 attends only to itself; row 2 splits its weight over positions 1 and 2; the last row is the only one that sees the whole sequence.

Why −∞ before softmax

Masking with −∞ and then softmaxing renormalizes the distribution over the surviving cells automatically. The illegal positions never enter the denominator in any meaningful way (e^−∞=0), so each row is a clean probability distribution over the past alone — no extra step required.

Why not zero the weights after

If you softmax first and then zero the future cells, the rows no longer sum to 1: the denominator already included the (now-removed) future weights, so the surviving weights are shrunk by whatever mass the future stole. You would have to renormalize each row by hand. Masking pre-softmax folds that renormalization into the softmax for free.

Concretely, both routes reach the same numbers, but the pre-softmax route is one operation instead of three. In code it is a triangular boolean mask filled with −∞:

# scores S: (T, T) from lesson 05,  S = (Q @ K.T) / sqrt(d_k)
mask = triu(ones(T, T), diagonal=1)      # 1s strictly above the diagonal
S = S.masked_fill(mask == 1, -inf)       # future cells -> -inf
A = softmax(S, axis=-1)                   # rows sum to 1 over the past only
Z = A @ V                                 # causal context vectors

One subtlety often trips people up: is it really "no leakage" if the future values are still sitting in the V matrix? Yes. Because their attention weights are exactly 0, they multiply into the weighted sum Z = A@V as 0·v = 0 — they contribute nothing. The mask acts on the weights; the values are simply never selected.

3 · Dropout on the attention weights

The mask makes the layer correct. Dropout makes it generalize. On a small corpus — and this book's example is a single short story — a model with 124 million parameters will happily memorize its training text, driving train loss down while validation loss climbs. Dropout is the standard antidote.

Dropout picks a fraction p of the entries in a tensor and sets them to 0, at random, on every training step. Applied to the attention weights, it forces each query to spread its bets: it cannot rely on one favorite key always being available, because that key might be dropped this step. This breaks up brittle co-adaptations where a few weights conspire to overfit, and it acts like training an ensemble of slightly different attention patterns.

There is one required correction. If you zero out fraction p of the weights, the surviving weights sum to less than 1, so the context vector shrinks on average — the layer's output would be systematically smaller during training than at inference (when dropout is off). To keep the expected magnitude constant, the survivors are scaled up by 1/(1−p):

a_kept ← a_kept / (1 − p) (so the expected sum stays ≈ 1)

For p = 0.5 the survivors double; for GPT-2's typical p = 0.1 they are scaled by 1/0.9 ≈ 1.11. Two things to remember: dropout is a training-only tool — it is switched off entirely at inference, so generation is deterministic given the weights — and it is applied after the softmax, to the normalized weights (the more common variant in GPT-style models).

The trap

Forgetting the 1/(1−p) rescale creates a silent train/inference mismatch: the model trains on attention outputs that are ~(1−p) too small, then sees full-magnitude outputs at test time and underperforms for no obvious reason. The rescale is what lets you drop weights during training yet leave the inference path untouched.

4 · Multi-head attention

Now the second flaw. A single head holds one W_q, W_k, W_v, so it can learn exactly one way to decide "what is relevant." But relevance in language is plural: one relationship is subject–verb agreement, another is a pronoun reaching back to its antecedent, another is "the token immediately before me." One geometry cannot be all of these at once.

The fix is embarrassingly cheap. Instead of one attention over the full width d, run h attentions in parallel, each over a slice of width d_head = d/h. Each of these heads has its own small projections and computes its own causal attention over its own slice; the h resulting context vectors are concatenated back to width d and passed through one final output projection W_o that lets the heads mix:

d_head = d / h                     # each head gets a 1/h slice of the width
for head i in 1..h:
    Q_i, K_i, V_i = project x into d_head             # per-head Q/K/V
    Z_i = softmax(causal(Q_i @ K_i.T / sqrt(d_head))) @ V_i   # (T, d_head)
Z = concat(Z_1, …, Z_h)            # (T, h * d_head) = (T, d)
out = Z @ W_o                      # final mix across heads -> (T, d)

The striking part is the cost. Because each head works on width d/h, the total arithmetic is essentially the same as one full-width head — you have rearranged the same FLOPs into h independent geometries rather than added new ones. For GPT-2 "124M", d = 768 and h = 12, so d_head = 768/12 = 64 — twelve parallel 64-dimensional views for the price of one 768-dimensional one. In practice, heads specialize: interpretability work finds "previous-token" heads and "induction" heads that implement copy-and-continue patterns.

Shapes, and the quadratic wall

In real code the heads are not a Python loop — they are a reshape. You project once to full width, then split the width axis into (h, d_head) and move the head axis next to the batch so the matmuls batch over it:

x (B, T, d) → Q,K,V (B, T, d) → reshape (B, T, h, d_head) → transpose (B, h, T, d_head)

scores = Q @ Kᵀ → (B, h, T, T) → causal-mask + softmax + dropout → @ V → (B, h, T, d_head) → concat → (B, T, d) → W_o

Notice the score tensor (B, h, T, T). It is quadratic in the sequence length T: doubling the context quadruples this tensor's memory. For short toy sequences this is nothing, but at long context it becomes the dominant memory cost — and the reason FlashAttention computes attention tile-by-tile without ever materializing the full T×T matrix (see the cross-reference below).

5 · Watch the future go dark

The widget shows the causal multi-head attention over the six-token sentence "Your journey starts with one step," with each head a different seeded geometry (so the patterns genuinely differ). Toggle causal mask and watch the upper triangle grey out to exactly 0 while every row renormalizes over the past — the KPI "sees future?" flips from yes to no and the masked-cell count jumps to the triangular number. Slide heads from 1 to 4 to switch which head's map you see (note d_head = d/h shrinking) and slide dropout to zero a seeded fraction of the legal weights, rescaling the survivors so the rows still sum to ~1.

Failure modes & checklist

Failure modes

No causal mask (or masking the wrong triangle). The model reads future tokens and reports a suspiciously low training loss that collapses at inference. Signal: great train perplexity, useless generation.
Zeroing weights after softmax without renormalizing. Rows no longer sum to 1; each query's attention is silently scaled down by the mass the future would have taken.
Dropout left on at inference, or missing the 1/(1−p) rescale. Nondeterministic generation, or a train/inference magnitude mismatch that quietly hurts accuracy.
Choosing h that does not divide d. d_head = d/h must be an integer; otherwise the reshape into heads is undefined.
Forgetting the output projection W_o. The concatenated heads never get to mix, wasting much of the multi-head benefit.

Checklist

Upper triangle of the score matrix set to −∞ before softmax.
Each row sums to 1 over positions ≤ i only.
Dropout applied to the post-softmax weights, training only, survivors scaled by 1/(1−p).
h divides d; each head runs over width d_head = d/h.
Heads concatenated back to width d, then one output projection W_o.
Score tensor is (B, h, T, T) — remember it is quadratic in T.

Checkpoint

Try it

Turn the causal mask OFF and read the KPIs: "sees future?" says yes and "masked cells" is 0 — this is the leaky lesson-05 layer. Now turn it ON: the upper-triangle cells go grey (weight exactly 0), "masked cells" jumps to 15 (that is 1+2+3+4+5, the future cells of a 6-token sequence), and "sees future?" flips to no. Note that the top row never changes — position 1 only ever saw itself — while the bottom row is unchanged too, because the last token is allowed to see everything before it. Then push heads to 4 and step viewing head across 1–4: the maps differ, yet d_head has dropped to d/4. Which of the three properties from lesson 00 — accurate, trainable, steerable — does the causal mask protect?

Where this points next

We now have the real attention layer used inside every GPT: causal, multi-headed, regularized with dropout. But step back and notice what attention actually does — it takes each token's vector and replaces it with a weighted blend of other tokens' vectors. That is a mixing operation across positions; it moves information sideways but performs no heavy nonlinear computation within a position. And if you simply stack many of these layers to get depth, the activations drift in scale and the gradients vanish, so the deep stack refuses to train. Lesson 07, The transformer block — norm, GELU, residuals, adds the three pieces that fix both: a feed-forward MLP for per-position processing, LayerNorm to hold activations in range, and residual connections to keep a clean gradient path through depth.

Takeaway

Causal attention forbids a token from seeing the future by setting the upper triangle of the score matrix to −∞ before softmax, so future cells get exactly 0 weight and each row renormalizes over the past alone — masking pre-softmax folds that renormalization in for free, which is why we do not zero weights afterward. Dropout on the post-softmax weights (training only, survivors scaled by 1/(1−p)) breaks up overfitting co-adaptations. Multi-head attention splits the width into h heads of size d/h, runs h causal attentions in parallel for h different "views" at essentially the same FLOPs, then concatenates and projects with W_o. The score tensor is (B, h, T, T) — quadratic in context length, the wall FlashAttention was built to climb.

Interview prompts

Why must a next-token predictor use causal masking? (§1 — position t predicts x_t+1; if it can attend to positions > t it copies the answer during training and has nothing to read at inference, so it must only see positions ≤ t.)
Why set masked scores to −∞ before softmax instead of zeroing weights after? (§2 — e^−∞=0 and softmax renormalizes over the survivors automatically, giving rows that sum to 1 over the past in one step; zeroing after leaves rows summing to less than 1 and needs a manual renormalize.)
If the future values are still in V, how is there no leakage? (§2 — their attention weights are exactly 0, so 0·v=0 in Z=A@V; the mask acts on the weights and the future values are never selected.)
Why must dropout rescale the surviving weights by 1/(1−p)? (§3 — dropping fraction p shrinks the weighted sum; rescaling keeps the expected magnitude ≈ 1 so the training-time output matches the dropout-free inference path.)
What does multi-head attention buy, and at what cost? (§4 — h independent similarity geometries (views) instead of one; because each head is width d/h the total FLOPs are essentially unchanged — the same compute rearranged.)
Why must h divide d, and what is d_head for GPT-2 124M? (§4 — heads split the width evenly, so d_head=d/h must be an integer; for d=768, h=12 it is 64.)
What is the shape of the score tensor and why does it matter at scale? (§4 — (B, h, T, T), quadratic in T; it dominates memory at long context and is what FlashAttention avoids materializing.)

Companion reads: cs336 · 07 Kernels shows how FlashAttention computes the (B, h, T, T) scores tile-by-tile without ever storing the full matrix — the answer to the quadratic wall; cs336 · 02 The Transformer counts the params and FLOPs of multi-head attention at scale.