Causal masks and multiple heads

Why is ten out of ten worth nothing?

Give someone a ten-question quiz with the answer printed in small grey type under every question. They finish in forty seconds and score ten out of ten. They have learned nothing, and on a real exam, with no answers printed, they would fail.

The attention code from the last lesson trains a model under exactly those conditions.

Pause and predict When a language model is trained, what plays the part of the answer printed under the question?

Two problems with plain attention

Problem A: the answer is visible

Take the training text “the cat sat on a mat”. The model reads it once and must make a prediction at every position:

With plain attention, the row for “cat” can put weight on “sat”. But “sat” is what “cat” is being asked to predict.

Gradient descent is relentless. If copying the answer lowers the loss, the model learns to copy.

There is a second reason. When the model generates text, the next token does not exist yet. There is nothing to peek at. A model trained with peeking would be practising for an exam it never gets to sit.

The problem

During training, position t must predict token t+1, but attention lets position t look straight at token t+1.

The naive fix

Feed the model one prefix at a time: first “the”, then “the cat”, then “the cat sat”… Each run only contains the past, so nothing can leak.

Why it fails

A sequence of 1024 tokens would need 1024 separate forward passes, over prefixes that average half the full length. Training would be hundreds of times slower.

The idea

Feed the whole sequence once, but inside attention forbid every token from looking to its right. Each row then behaves exactly as if it had been given only its prefix.

The trade-off

A token can never use later context, even when that would help understanding. For a model whose job is to continue text, that is precisely the rule we want.

Problem B: one blend per token is not enough

Softmax gives each token one row of weights that sums to 1. One row means one blend of the other tokens.

But look at “it” in “The animal didn’t cross the road because it was tired”. To be useful, “it” would like to know several things at once:

• What do I refer to? (animal)
• What word came just before me? (because)
• What is this sentence about in general? (a bit of everything)

If “it” spends 80% of its weight on “animal”, only 20% is left for everything else. Hiring has the same problem: one interviewer cannot judge coding, communication and design in the same hour, which is why companies use a panel.

A mental model

A. The mask

Take the score table from the last lesson: one row and one column for each of the T tokens in the sequence. The row is who is looking. The column is who is looked at.

Every cell to the right of the diagonal is a token looking at something that comes after it. Cross out that whole triangle.

What remains is a staircase. The first token sees only itself, the second sees two tokens, and so on.

Causal mask

A rule inside attention: a token may look at itself and at earlier tokens, never at later ones. “Causal” because information only flows from past to future.

For example: In “the cat sat”, the row for “cat” may use “the” and “cat”. Its weight on “sat” is forced to exactly 0.

B. The heads

If one search is not enough, run several. Give each token a few small attention mechanisms that work side by side, like the interview panel: each takes its own notes on the same tokens, and the notes are joined at the end. Each one is called a head. Unlike a panel, nobody tells a head what to listen for: its focus is learned because it lowered the loss.

Multi-head attention

Several small attention operations run in parallel on the same tokens. Each head has its own query, key and value projections, so each can look for something different. Their results are joined together at the end.

For example: With D = 8 numbers per token and 2 heads, each head works with 4 numbers per token. Each head also builds its own table of weights, one row and one column per token.

Let's see the numbers

One row, by hand. The sentence is “the cat sat”. We are the row for “cat”, which must predict “sat”. Suppose its three scores (already divided by √d) are [1, 2, 3].

Pause and predict Switch the mask on. “sat” gets no weight. What happens to the weights on “the” and “cat”?

Heads: division and nothing else

Our tiny GPT has D = 128 numbers per token and 4 heads, so each head works with 128 / 4 = 32 numbers. Four heads of 32 hold exactly as many numbers as one head of 128.

The math

The mask is a single addition inside the formula you already know:

And multi-head attention is “do that h times on thinner slices, then join”:

What you should remember

1. Without a mask, position t can look straight at token t+1 and copy it. The causal mask sets every score for a later token to −∞ before softmax. Weight exactly 0, and the remaining weights still sum to 1.
2. The mask also matches generation, where the future does not exist, and it makes training efficient: one pass over T tokens = T training examples.
3. One softmax row is one blend. Multi-head attention runs h thinner attentions in parallel, each with its own Q/K/V slice and its own weight table, then concatenates and mixes with W_O.
4. Heads split the D numbers (width D/h each, in GPT-2-style models). Same parameters as one big head. Shapes: (T, D) → (h, T, D/h) → (T, D).

Where this appears in a real LLM

In this lesson

• 6 tokens, random scores, a hand-built “cheating” head
• 4 heads with hand-written patterns and friendly names
• D = 4, heads of width 1 to 4
• The mask is rebuilt on every call

In a real system

• The identical mask, on sequences of thousands of tokens or more
• GPT-2 small: 12 heads × 12 layers = 144 learned heads, none with a name
• Head width is typically 64 to 128; a 70B-class model has 64 query heads per layer
• The mask is precomputed once (tiny_gpt.py) or never materialised at all (fused kernels)

One modern change: many current models let several query heads share one set of keys and values (grouped-query attention) to save memory during generation. The idea of several parallel heads is unchanged. More in Modern LLM architecture.

If you train the model from the last lesson again with the mask switched on, its training loss ends higher than before. That is expected. The model can no longer copy the answer, so the loss now measures how well it really predicts the next token.