Masked attention restricts which sequence positions can contribute to an attention output. In causal self-attention, a token can use itself and earlier tokens but not later ones. This lets a decoder predict the next token without seeing future tokens in the sequence. section 11.1.3 section 11.1.1
Ordinary self-attention can weigh every token in the input sequence. A causal mask changes the attention scores before those scores become weights: entries for future positions are set to negative infinity. Softmax then assigns those positions zero weight. The current position remains available because the mask includes the main diagonal. section 11.1.3
A four-token prompt
For “this is a prompt,” the first position can attend only to “this” when predicting “is.” A lower-triangular 4 × 4 mask allows positions 1 through 4 to attend to 1, 2, 3, and 4 positions respectively. After softmax, every masked entry has weight zero. section 11.1.3
To build causal multi-head attention in PyTorch:
- Project the input embeddings into query, key, and value vectors, then divide each into parallel heads. GPT-2XL uses 25 heads for 1,600-dimensional embeddings. section 11.1.3 section 11.2.3
- Compute scaled query–key scores for each head. Create a lower-triangular mask with torch.tril and use masked_fill to replace scores for future positions with negative infinity. section 11.1.3
- Apply softmax, multiply the resulting weights by the value vectors, and join the heads back into one attention output. Register the reusable mask as a nonlearnable buffer in a CausalSelfAttention module. section 11.1.3 section 11.2.3
Figure 11.3 How to calculate masked attention weights in causal self-attention. A mask is applied to the scaled attention scores so that values corresponding to future tokens (those above the main diagonal in the matrix) become –∞. We then apply the softmax function on the masked scaled attention scores and obtain the masked attention weights. The masking ensures that predictions for a given token can only be influenced by the tokens that precede it in the sequence, not by future tokens. This is essential for generating coherent and contextually relevant text outputs.
The score-matrix diagram shows why masking must happen before softmax: future-position scores become negative infinity, then receive zero attention weight. section 11.1.3
The same mechanism fits different decoder-only models. It forms the attention sublayer of GPT-2XL decoder blocks, a smaller four-head text generator, and an eight-head music Transformer that predicts sequences of musical events. The head count and sequence length change; the restriction on future positions does not. section 11.2.4 section 12.3.2 section 14.1.2
When loading pretrained GPT-2XL weights into an implementation with its own registered mask, omit the checkpoint’s attn.masked_bias entries: they serve future-token masking rather than providing weights the new mask needs. section 11.3.1
