Causal Masking for Autoregressive LMs

Introduction

Welcome to the third module of our course, where we'll explore variants of the attention mechanism and the diverse architectures they enable. In the previous lesson, you mastered the skill of calculating a Transformer's parameter count, giving you a solid grasp of the model's static composition and scaling properties.

Today, we transition from the static architecture to the dynamic control of information flow. Our learning outcome is to implement causal masking to enable autoregressive language modeling. This mechanism is the cornerstone of generative models like GPT. It's what allows a model to generate text one token at a time, just as a human would, by ensuring it only has access to the past, not the future.

By the end of this lesson, you will understand why causal masking is critical, how it's implemented mathematically, and how to build it in PyTorch.

1. The "Cheating" Problem in Language Generation

In Module 2, we learned that in self-attention, every token calculates attention scores with every other token in the sequence. This is called bidirectional attention and is perfect for tasks where the goal is to understand the entire context of a sentence, such as sentiment classification (as in BERT-style models).

However, for language generation, this poses a fundamental problem. Our goal is to train a model to predict the next token in a sequence. If, during training, the token at position t is allowed to "see" the token at position t+1, it doesn't need to learn any patterns in the language. It can simply "cheat" by copying the answer.

This leads to perfect performance during training but complete failure during inference, where the future tokens don't exist yet. To make this concrete, let's watch a short, clear explanation of this issue.

Decoder-Only Transformer for Next Token Prediction: PyTorch Deep Learning Tutorial

The video 'Decoder-Only Transformer for Next Token Prediction' by Luke Ditria provides a concise overview of why standard self-attention is unsuitable for autoregressive tasks and introduces the high-level solution.

Watch the segment from 00:19 to 06:33. Pay close attention to: The explanation of how a token can 'cheat' by querying future tokens. The concept of the attention map and how modifying it can solve the problem. The visual representation of the triangular mask that emerges.

As the video explains, the solution is to prevent a token at a given position from attending to any subsequent tokens. This is called causal masking or sometimes look-ahead masking.

2. The Core Mechanism: Masking Attention Scores

Causal masking ensures that information flows only from the past to the present. For a query at position i, it can only attend to keys at positions j where j <= i.

Causal Masking for Autoregressive Language Modeling — This image illustrates the core principle of causal masking. For any token (query), attention is allowed for itself and all preceding tokens (keys), but blocked for all future tokens.

The implementation of this concept is an elegant mathematical trick applied directly to the attention score matrix, $W_{scor es} = Q K^{T}$ , before the softmax function is applied.

We create a "mask" matrix. This is typically an upper-triangular matrix where the values corresponding to illegal, future positions are 1 (or True) and the allowed past positions are 0 (or False).
We use this mask to replace the attention scores at the illegal positions with a very large negative number (effectively $- \infty$ ).
When we apply the softmax function, softmax(z), the components with $- \infty$ will evaluate to e-inf, which is zero.

This ensures that future tokens have zero weight in the final weighted sum of the value vectors, effectively blocking any information flow from the future.

3. Implementing Causal Self-Attention From Scratch

There is no better way to understand this than to build it from the ground up. The following resource from Andrej Karpathy does exactly that, evolving from a simple average to a full, causally-masked self-attention head. Your background in PyTorch and deep learning will allow you to follow the code and logic closely.

Let's build GPT: from scratch, in code, spelled out.

In his renowned 'Let's build GPT: from scratch' tutorial, Andrej Karpathy walks through the implementation of a decoder-only Transformer. We will focus on the exact moment he introduces and implements the causal attention mechanism.

Watch the segment from 42:18 to 1:20:49. This is the most critical part of the lesson. Focus on: The Mathematical Trick (42:18 - 54:41): Understand how a matrix multiplication with a lower triangular matrix can be used to perform a weighted aggregation of past elements efficiently. This is the foundation of the technique. Building the Self-Attention Head (56:17 - 1:20:49): Observe how this trick is integrated into a self-attention head. Pay close attention to how the tril (lower triangular) matrix is used to create a mask, which is applied to the raw attention scores (wei) before the softmax. This directly implements causal attention.

Let's summarize the key PyTorch components from the video that bring this to life:

torch.tril(torch.ones(T, T)): This creates a $T \times T$ lower-triangular matrix of ones. This forms the basis of our mask. When used to mask the attention scores, it allows a token to attend to itself and previous tokens.
wei.masked_fill(mask == 0, float('-inf')): This is the core operation. It finds all positions in the attention scores wei where our mask is 0 (i.e., the future positions in an upper-triangular section) and fills them with negative infinity.
register_buffer: Karpathy registers the tril matrix as a buffer. This is a best practice in PyTorch for tensors that are part of the model's state but are not learnable parameters (i.e., they should not be updated by the optimizer but should be moved to the GPU along with the model).

For a complementary text-based guide that shows a clean, encapsulated CausalAttention module, you can review the following resource.

Creating a Transformer From Scratch - Part One: The Attention ...

The article 'Creating a Transformer From Scratch' by Benjamin Warner provides a well-commented PyTorch implementation of a Causal Attention layer, reinforcing the concepts from the video.

Read the section 'Causal Self-Attention'. Note how a permanent causal_mask is created using torch.triu and registered as a buffer. Observe its application in the forward pass using masked_fill before the softmax.

The use of torch.triu (upper triangular) with diagonal=1 in the article achieves the same outcome as Karpathy's torch.tril approach. It creates a boolean mask where True indicates a position to be masked. The choice between triu and tril is a matter of implementation preference, as long as the logic correctly masks future tokens.

Test your understanding!

Consider a sequence of length 4.

Write down the 4x4 boolean mask matrix that would be used for causal attention. Use True for positions that should be masked (not attended to) and False for positions that are allowed.
Why is it crucial to apply this mask before the softmax function instead of after?

Show answer

The 4x4 boolean mask would be an upper triangular matrix (excluding the diagonal):
```
[[False, True,  True,  True],
 [False, False, True,  True],
 [False, False, False, True],
 [False, False, False, False]]
```
- The first token (row 0) can only attend to itself (column 0).
- The second token (row 1) can attend to the first and itself (columns 0, 1).
- And so on.
The mask must be applied before the softmax. The softmax function normalizes a set of scores into a probability distribution that sums to 1.
- If you apply the mask before, you set the scores of future tokens to $- \infty$ . The softmax of $- \infty$ is 0, so these tokens receive zero attention probability.
- If you were to apply the mask after the softmax, you would be zeroing out probabilities in an already-normalized distribution. The remaining probabilities would no longer sum to 1, breaking the mathematical foundation of a weighted average. The attention weights would be invalid.

Conclusion

Today, you've implemented one of the most important mechanisms in modern LLMs. Causal masking is what transforms the general-purpose attention mechanism into a powerful tool for autoregressive sequence generation.

Key Takeaways:

Autoregressive models require causal masking to prevent the model from accessing future tokens during training, which would be a form of data leakage or "cheating."
This is implemented by modifying the attention score matrix (QK^T) before the softmax normalization step.
The standard technique involves creating a triangular mask and using it to set the scores for future positions to negative infinity (-inf).
In PyTorch, this is achieved using functions like torch.tril or torch.triu to create the mask, register_buffer to store it within the nn.Module, and masked_fill to apply it.
An architecture that exclusively uses causally masked self-attention is known as a decoder-only Transformer, the foundation for models like GPT.

What's Next?

Now that you have a deep understanding of causally masked self-attention, you are prepared to situate it within the broader family of attention mechanisms. In our next lesson, we will compare self-attention, cross-attention, and masked self-attention, identifying their specific use cases. This will clarify when and why each type is used, from encoder-only models (BERT) to decoder-only models (GPT) and encoder-decoder models (T5).