Transformer Attention: Why can not GPT look ahead?

Published

This is a series where I answer the questions my tutoring students really ask – the ones textbooks tend to skip. I teach machine learning, statistics, and math to learners of all levels.

The question. A student who had just worked through the first three posts of this series came back with a worried look: “During training, the model sees the whole sentence at once, and every token can attend to every token. So when the model is supposed to predict the word after ‘journey,’ what stops it from just looking at the answer? Isn’t that cheating?”

It is cheating, and without one small fix it works perfectly – the model copies the future instead of predicting it, learns nothing, and then fails the moment it has to generate text on its own. The fix is called a causal mask, and it is two lines of code.

Where the problem comes from

Recall the attention matrix from the first post in this series. Every row is one token’s distribution over where to look, and every row uses the entire sentence. The row for “with” spread its weight across all six tokens, including “one” and “step” – words that have not been said yet if the model is generating left to right.

At generation time this is impossible by construction: the future tokens literally do not exist yet. But at training time the whole sentence is right there in the input matrix, and nothing in the score-normalize-mix routine respects the arrow of time. The model will happily attend to the very word it is supposed to predict, because copying is easier than learning language.

The fix: hide the future before the softmax

Here is the whole trick, applied to the same six-token example we have used all along. First compute the scores exactly as before, then set every score that points to a future token to minus infinity, before the softmax runs:

scores = inputs @ inputs.T

mask = torch.tril(torch.ones(6, 6))
masked_scores = scores.masked_fill(mask == 0, float("-inf"))
masked_weights = torch.softmax(masked_scores, dim=-1)

print(masked_weights[3])  # the row for "with"
# tensor([0.2046, 0.2956, 0.2915, 0.2084, 0.0000, 0.0000])

Two things to notice in that one line of output. First, the last two weights are exactly zero – “with” can no longer see “one” or “step,” because softmax of minus infinity is zero, so the future is not just discounted but erased. Second, the remaining weights were renormalized: before masking, this row was [0.14, 0.21, 0.20, 0.15, 0.13, 0.17], and now the same four allowed entries have grown to [0.20, 0.30, 0.29, 0.21] so the row still sums to one. The token’s attention budget did not shrink; it was redistributed over the past.

Here is what the whole matrix looks like after masking:

The attention matrix after applying a causal mask. Every entry above the diagonal is exactly zero, so each token can only attend to itself and the tokens before it.

The upper triangle is all zeros, and the diagonal is the last survivor in each row. Look at the first row: “Your” attends to itself with weight 1.00, because at position one there is nothing else to look at. Every row is still a proper probability distribution – the softmax sees to that – but each one is confined to its own past.

Student follow-up: “Doesn’t throwing away half the matrix waste training signal?”

The opposite, and this is the part my student found genuinely elegant. After masking, every row is a complete, independent prediction problem: the row for “Your” is the task of continuing a one-word sentence, the row for “journey” is the task of continuing “Your journey,” and so on down to “step,” which has seen everything. One sentence of six tokens becomes six training examples of increasing difficulty, computed in a single pass. The mask does not reduce what the model learns per sentence – it multiplies it.

One thing to try

Predict the masked row for “step” before running the code: how many nonzero entries should it have? Then flip the mask with torch.triu instead of torch.tril, so every token can only see the future, and predict what the last row becomes. If your answer was “all zeros except a 1.00 on itself,” run it and check.

The code in this article is adapted from the teaching notebooks I use with my students, and it follows the progression in Sebastian Raschka’s Build a Large Language Model (From Scratch), which I highly recommend if you want to continue from here to multi-head attention and beyond.

References

  1. Vaswani, A., et al. (2017). “Attention Is All You Need.” Advances in Neural Information Processing Systems 30 (NeurIPS 2017).
  2. Radford, A., et al. (2018). “Improving Language Understanding by Generative Pre-Training.” OpenAI.
  3. Raschka, S. (2024). Build a Large Language Model (From Scratch). Manning.

I tutor machine learning, deep learning, statistics, and math. If attention – or anything else – is not clicking for you, you can find me on Wyzant.