Masking and Causality

The score table is the answer key

Which positions is a head allowed to read? So far, all of them, and the question never came up because of an assumption nobody wrote down: that the sentence is a finished object, handed over complete, free to be studied from any angle. Plenty of tasks do look like that. One important task does not. A model shown a sentence one word at a time and asked what comes next is in a different position, and for it the score table built over the last five chapters is the answer key.

That task is next-token prediction. The sentence still goes in as , all five words at once. What changes is what is asked of the output: row is no longer read as a better representation of word , it is read as a guess at word .

So (1) The must guess cat, (2) cat must guess chased, (3) chased must guess the, and (4) the must guess mouse. Five words, four guesses. The loss measures how wrong those four are, and that number is the entire training signal.

Nothing about the mechanism has changed. What changed is what is wanted from it. Earlier, chased pulling 0.60 out of mouse was the mechanism working as intended: a verb is incomplete without its object, and the blend turned chased into something closer to chased a mouse. Reaching forward was fair, because the whole sentence was there to read. The goal is no longer that. chased is not being asked what it means, it is being asked what comes next.

How to delete an entry

Take row (3) again: 0.04, 0.18, 0.12, 0.06 and 0.60. Under the new job, chased has to produce the the at (4). Its row hands it 0.06 of that exact word, and another 0.60 of mouse one step further on. Two thirds of what chased walks away with is built out of the words it is meant to be predicting. The answer is inside the vector it has to answer from.

If the trouble is that a word can see what comes after it, stop letting it. chased may look at (1), (2) and itself, and nothing further. (2) gets one fewer, (4) one more, and (1) is left with only itself.

The cells are not removed, they are made to count for nothing. Before the softmax, every forbidden score is set to . The softmax raises to each score and , so those cells come out with weight zero, and the row still sums to one because the division is by the total of what is left.

All of it is one addition. Build a matrix of the same shape as the score table, holding everywhere a word is allowed to read and everywhere it is not:

Here is the row, the word doing the reading, and the column, the word being read, so says this word came no later than mine. Adding leaves a score as it was and adding takes it out, so the whole rule is one matrix added to the scores:

is , built from the length of the sentence alone. It never looks at the words, so every sentence of five tokens gets the same , reused by every head and every sentence in the batch.

One pass, four guesses

The mask is not the only way to keep the model honest. It could be run once per prefix instead: show it The and ask for cat, then The cat and ask for chased, and so on. Nothing leaks that way either, because the forbidden words are not in the input at all. It also costs four forward passes for one five-word sentence, each recomputing the prefix the pass before it already worked through.

One masked pass already contains those four. Row (3) may read (1), (2) and itself, which is exactly what a run on The cat chased would have had in front of it, and the same holds for every other row. A single pass over the whole sentence computes all four prefixes at once, so the four guesses come out together and the loss collects all four at the same time.

Sentences of different lengths

Training runs a batch of sentences at once, and a batch is a single tensor. Tensors are rectangular, so every sentence in the batch has to be the same length, and sentences are not. The usual answer is to take the longest and fill the others out with a filler token: padding.

The model cannot tell the difference. A pad is a token like any other: it has an embedding, produces a key and a value, scores against every query, and the softmax gives it a share of the weight. Nothing in the mechanism marks it as filler, so real words end up with part of their answer built out of positions that stand for nothing.

Train on sentences at once, pad each to the length of the longest, and the batch arrives as of shape . Sentence has true length , and everything past is filler. Masking it is one more matrix of the same kind:

The index is absent from the right-hand side. The rule is the same for every query, so blanks whole columns, while depends on both indices and blanks a triangle. Different shapes, same mechanism, both added to the scores before the softmax:

is shared across the whole batch, since it depends only on the length. is constant down each column, so in practice it is stored as one boolean per position per sentence — flags — and broadcast across the rows and the heads. Applying both is enough to mask whatever either one targets, because adding to any finite score gives . The loss is computed only over positions holding real words, so whatever the padded rows compute is never used.

def masked_attention(raw_scores, true_lengths):
    """raw_scores (B, n, n) before softmax; true_lengths (B,) words per sentence."""
    B, n, _ = raw_scores.shape
    causal = torch.triu(torch.full((n, n), float("-inf")), diagonal=1)   # M
    is_pad = torch.arange(n).unsqueeze(0) >= true_lengths.unsqueeze(1)
    pad = torch.where(is_pad, float("-inf"), 0.0).unsqueeze(1)           # P
    return F.softmax(raw_scores + causal + pad, dim=-1)

The two masks solve independent problems. enforces the left-to-right reading order, so the model cannot peek at what it is meant to predict. cleans up the artefact left behind by forcing different-length sentences into one rectangle.

What the mask still cannot fix

The mask enforces a strict reading order, so causality is preserved. But look closely at the formula. Scramble the order of the input words and the dot products between their queries and keys are unchanged. Attention treats the sentence as a bag of words, computed all at once, with no inherent sense of sequence.

The mask says which words a position may look at. It says nothing about where those words sit. Order has to be injected before attention begins, and that is the next chapter.