Causal Self-Attention in Autoregressive Models

Introduction

In our previous lesson, we established that self-attention is a powerful mechanism that allows every token in a sequence to look at and draw context from every other token. We saw how this is achieved by projecting each input into Query, Key, and Value vectors and then computing a weighted sum of the Values, where the weights are determined by Query-Key similarity.

This "all-to-all" communication is perfect for tasks like sentence classification, where understanding the full context is beneficial. However, it presents a fundamental problem for autoregressive models like GPT, which are designed to generate text one token at a time. During training, if a token can "see" the tokens that come after it, it can simply copy the next word instead of learning to predict it. This is like giving a student an exam with the answer key attached.

Today's lesson addresses this exact problem. Our learning objective is to explain causal (masked) self-attention and why it is required for autoregressive models. We will explore the mechanism that prevents the model from "cheating" by looking into the future.

1. The Autoregressive Constraint

Autoregressive models generate sequences step-by-step. When generating a sentence, the model predicts the first word. Then, it uses the first word to predict the second. It uses the first two words to predict the third, and so on.

• Prediction at step t can only depend on the output from steps 1, 2, ..., t-1.

This property must be preserved during training. When we train the model on the sentence "robot must obey orders," the prediction for the word "obey" should only be based on "robot" and "must." It cannot be influenced by "orders." Standard self-attention violates this constraint.

Causal self-attention enforces it by applying a "mask" to the attention scores, ensuring that information from future tokens is blocked.

2. Implementing the Mask: A Tale of Two Methods

So, how do we implement this mask? Our goal is to modify the attention weight matrix so that all entries above the main diagonal are zero, creating a lower-triangular matrix.

A tempting and seemingly straightforward approach would be to:

1. Compute the full attention weight matrix using softmax, just like in standard self-attention.
2. Manually set all the values in the upper triangle of the matrix to zero.
3. Renormalize each row so the weights sum to 1 again.

This approach seems logical, but it contains a critical flaw. Let's explore why this "naive" method is incorrect and what the right way is.

3. The Correct Approach: Masking Before Softmax

As the video explained, the flaw in the naive approach is data leakage. The softmax function normalizes scores by dividing by the sum of the exponentiated scores across the entire row. This means even before we zero out the weights for future tokens, their scores have already contributed to the normalization factor for all other tokens in that row.

The correct solution is to apply the mask before the softmax function. We do this by adding negative infinity ($-\infty$) to the attention scores we want to block.

Why does this work? The softmax function involves an exponential, $e^x$.

• $e^{-\infty }\approx 0$

When we add $-\infty$ to the attention scores for future positions, the softmax function will naturally assign them a weight of 0. The weights for the past and current tokens will then be automatically renormalized to sum to 1, with no leakage.

The process is:

1. Compute attention scores: $Q{K}^T/\sqrt{d_k}$
2. Add the causal mask matrix (0 for allowed positions, $-\infty$ for masked positions).
3. Apply softmax. The $-\infty$ entries will become 0.

This short video provides a clear visual of this exact process.

The Mathematical Formulation

Let's formalize this. The masked self-attention calculation is:

$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}+M\right)V$

Where $M$ is the mask matrix:

$M_{ij}=\{\begin{array}{ll}0& \text{if }i\ge j\text{ (allow attention)}\\ -\infty & \text{if }i<j\text{ (block attention)}\end{array}$

This ensures that for a query at position $i$, it can only attend to keys at positions $j\le i$.

To dive deeper into the mechanics and the rationale, please read the following article.

4. Causal Attention in the Transformer Architecture

Now that we understand the mechanism, let's see where it fits. In the original Transformer architecture, this masked attention is a key feature of the decoder. Decoder-only models like GPT use this mechanism in every one of their blocks.

Efficiency: Training vs. Inference

Causal masking provides a massive efficiency benefit during training. Because the mask prevents information from flowing forward, we can feed the entire sequence to the model at once and compute the predictions for all tokens in parallel. The model can't cheat, so this is valid. This process, where the model is trained on the ground-truth sequence, is often called teacher forcing.

During inference, however, we don't have the future tokens. We must generate them one at a time, in a sequential loop. This is where the KV Cache (which we'll discuss in detail in a later lesson) becomes critical for efficiency.

The "Training vs Inference" section of the "Masked and Causal Attention" article you just read provides a good summary of this key distinction.

Conclusion

In this lesson, we bridged the gap between the all-to-all communication of self-attention and the step-by-step nature of autoregressive language generation.

Key Takeaways:

• Purpose: Causal (or masked) self-attention is required for autoregressive models to prevent them from "seeing the future" during training.
• The Problem with Naive Masking: Masking attention weights after softmax leads to data leakage, as future tokens have already influenced the normalization step.
• The Correct Mechanism: The mask must be applied before softmax by adding $-\infty$ to the attention scores of future positions. The exponential in the softmax function naturally turns these into zero-weight connections.
• Efficiency: This mechanism allows for highly parallelized training (teacher forcing) on entire sequences at once, while still enabling correct, sequential generation at inference time.

Preview of the Next Lesson:
So far, we've treated self-attention as a single computational block. However, what if a model needs to focus on different kinds of relationships simultaneously? For instance, one attention pattern might track syntactic dependencies, while another tracks semantic similarity. Our next lesson will introduce multi-head attention, a technique that allows the model to do just that by running the self-attention mechanism multiple times in parallel.