Implementing Causal Masking in Autoregressive Models

Introduction

In our previous lesson, we established the statistical rationale behind scaling dot-product attention, culminating in the formula that governs how tokens weigh each other's importance:

Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V

As it stands, this equation allows every token in a sequence to attend to every other token, both past and future. While powerful, this bidirectionality presents a fundamental conflict with the nature of generative language models. Such models, often called autoregressive, build text sequentially, one token at a time. When predicting the next word, they can only rely on the words they have already generated; the future is, by definition, unknown.

This lesson addresses this conflict by introducing a crucial modification to the attention mechanism: causal masking. You will learn why this is necessary for any generative model and, most importantly, how it is implemented as an elegant mathematical trick that prevents the model from "seeing" into the future, ensuring it learns meaningful patterns from the past.

The Principle of Causality in Autoregressive Models

The term "autoregressive" means that the model's prediction at a given step is a function of its own previous outputs. For a language model, this translates to predicting the next token, $x_{t + 1}$ , based solely on the sequence of preceding tokens, $x_{1}, x_{2}, \dots, x_{t}$ .

This creates two scenarios where unrestricted attention would fail:

Inference Time: When generating text, the model simply does not have access to future tokens. If it were trained to rely on them, it would be unable to function.
Training Time: During training, we typically provide the model with a complete text sequence all at once for computational efficiency—a technique related to "teacher forcing." If attention were unrestricted, the model could trivially predict the next token by simply "peeking" at the answer already present in the sequence. For instance, in predicting the word after "The cat," it could look ahead to see "sat" and achieve perfect accuracy without learning anything about grammar or semantics.

To solve this, we must enforce a strict causal or temporal ordering: information must only flow from past tokens to the current token being processed.

Implementing the Causal Mask

The solution is to modify the attention scores before they are passed to the softmax function. We want to programmatically force the attention weights for all future tokens to be zero. We can't simply multiply the weights by zero after the softmax, as this would violate the property that the weights for a given query must sum to one.

The correct approach is to add a mask matrix to the scaled dot-product scores.

This diagram visualizes the core principle of causal masking. A query from a token at a given position (e.g., T2) is only allowed to attend to keys from itself and previous tokens (T0, T1, T2). Attention to future tokens (T3, T4) is blocked.

The video below provides a clear, animated overview of this process and its motivation.

Attention in transformers, step-by-step | Deep Learning Chapter 6

This segment from 3Blue1Brown's "Attention in transformers, step-by-step" explains why masking is needed during training and how it's implemented.

Please watch the section from this point. The key insight is how we can manipulate the pre-softmax scores to get a zero-weight after softmax, without disrupting the normalization of the other weights.

The Mathematical Mechanism

As explained in the video, the trick is to add negative infinity ( $- \infty$ ) to the attention scores corresponding to future positions. When the softmax function is applied, an input of $- \infty$ results in an output of 0.

Let's formalize this. The equation for masked attention is:

Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}} + M) V

The only new component is the mask matrix, $M$ . For a sequence of length $n$ , $M$ is an $n \times n$ matrix defined as:

M_{ij} = {0 - \infty if j \leq i if j > i

Here, $i$ is the row index representing the query's position, and $j$ is the column index representing the key's position.

When $j \leq i$ , the key is at or before the query's position (a "past" or "present" token). We add 0, leaving the original score unchanged.
When $j > i$ , the key is in a future position. We add $- \infty$ .

The effect on the softmax is profound. For any score $s_{ij}$ where the mask is $- \infty$ :

softmax (..., s_{ij} - \infty, ...)_{i} = \frac{e ^{s_{ij} - \infty}}{\sum _{k} e ^{s_{ik} + M_{ik}}} \to \frac{0}{\sum _{k} e ^{s_{ik} + M_{ik}}} = 0

The numerator becomes zero, forcing the final attention weight to be zero. The attention from token $i$ to a future token $j$ is effectively blocked.

The following article provides a superb, detailed explanation of this concept, including the mathematical formulation and a clear visualization of the mask matrix.

Decoder Architecture: Causal Masking & Autoregressive ...

This article by Michael Brenndoerfer offers a comprehensive look at the why and how of causal masking.

Please read the first three sections, starting from The Causal Masking Requirement up to the end of the section Why −∞ Creates Zero Attention. Focus on how the text connects the autoregressive principle to the design of the mask matrix M and the resulting effect on the softmax calculation.

A Concrete Example

Let's trace the data flow with a simple 3-token sequence. Assume we have computed our scaled attention scores $S = \frac{Q K ^{T}}{d _{k}}$ :

S = s_{11} s_{21} s_{31} s_{12} s_{22} s_{32} s_{13} s_{23} s_{33}

Our causal mask matrix $M$ will be:

M = 000 - \infty 00 - \infty - \infty 0

We add this to our scores matrix:

S + M = s_{11} s_{21} s_{31} - \infty s_{22} s_{32} - \infty - \infty s_{33}

Now, we apply the softmax function row by row.

For token 1 (row 1): softmax([s_11, -inf, -inf]) will result in [1, 0, 0]. Token 1 can only attend to itself.
For token 2 (row 2): softmax([s_21, s_22, -inf]) will distribute probability between tokens 1 and 2, with a weight of 0 for token 3.
For token 3 (row 3): softmax([s_31, s_32, s_33]) will distribute probability across all three tokens.

This process ensures that for each token, the attention weights are properly normalized (sum to 1) across only the allowed past and present tokens.

This diagram provides a complete overview of the masked self-attention process. The panel "Computing the attention matrix (with masking)" explicitly shows the unnormalized scores being combined with a mask before the softmax step, illustrating where our causal mask fits into the larger computation.

To reinforce this, the following video provides another excellent walkthrough of applying the mask.

The matrix math behind transformer neural networks, one step at a time!!!

This clip from StatQuest's video on the Transformer decoder focuses specifically on the application of the mask matrix.

Watch the segment from this timestamp. It shows a concrete example of adding the mask matrix (containing zeros and negative infinities) to the scaled similarity matrix and how this affects the softmax output.

Conclusion

In this lesson, we have augmented the self-attention mechanism with a crucial feature for generative models. Causal masking is not merely an add-on but a fundamental requirement that aligns the Transformer's architecture with the principles of autoregressive sequence generation.

Here are the essential takeaways:

The "Why": Autoregressive models generate text sequentially, so they must not have access to future information. Causal masking prevents the model from "cheating" during training and ensures it functions correctly during inference.
The "How": It is implemented by creating a mask matrix $M$ where entries corresponding to future positions are $- \infty$ and all other entries are 0.
The Mechanism: This mask matrix is added to the scaled dot-product scores before the softmax function. The properties of the exponential function ensure that any score modified by $- \infty$ will result in a final attention weight of 0.

You now understand the complete architecture of a single, causal self-attention mechanism. However, a single attention "view" of the sequence can be limiting. In our next lesson, we will explore multi-head attention, a technique that allows the model to simultaneously attend to different aspects of the sequence, dramatically increasing its expressive power.