# Causal Attention Masking

Preventing a position from attending to later positions by masking future scores before the softmax.

**Rests on:** Self-Attention from Scratch

*Deep Learning · Transformers*


## Why learn this

Imagine you are teaching a computer to write a story. If the computer is allowed to read the end of the story while it is trying to write the beginning, it will just copy-paste or become confused about what comes next. Causal attention masking is the rule that forces the computer to look only at what has already been written, ensuring it generates text one word at a time in a logical order. Without this rule, the computer cannot learn to predict the future, which means it cannot write coherent sentences, complete your prompts in ChatGPT, or create any useful generative text. This concept is the foundation of how systems like GPT work.


## The intuition

Think about writing an email. When you type the first sentence, you haven't written the second one yet. You can only rely on what you have already put into the draft. You cannot cheat by looking at what you *will* write in five minutes. Causal attention masking does exactly this for a language model. It ensures that when the model is calculating which word to pick next, it is blinded to all words that come after the current position.

Let's look at the sentence: 'The cat sat on the mat.'
1. When the model tries to predict 'cat', it can see 'The'. It cannot see 'sat'.
2. When the model tries to predict 'sat', it can see 'The' and 'cat'. It cannot see 'on'.
3. When the model tries to predict 'on', it can see 'The', 'cat', and 'sat'. It cannot see 'the'.

In technical terms, each word position in the text is assigned a score for how much it should pay attention to every other word. Normally, words can pay attention to everyone. With causal masking, we take those scores for any word that appears *later* in the sequence and block them out completely. We do this by setting their importance to negative infinity before we convert the scores into probabilities. This effectively means 'ignore these future options entirely.' Because we remove the influence of future words, the model must learn to guess what comes next based only on context from the past.


## An analogy

Imagine a magician performing a card trick where they have to predict which card you will pick next. The rule is that the magician can only look at cards that are already face-up on the table; any face-down cards (which represent future words) must be completely ignored. To enforce this, imagine there is a sliding black curtain placed over the deck. As the trick progresses, the curtain moves right, revealing one new card but always keeping all unturned cards hidden behind it. The magician can never peek ahead because the curtain physically blocks their view of any card that hasn't been revealed yet. This ensures their prediction is based solely on the sequence of revealed cards, just as causal masking ensures a model generates text step-by-step without seeing the answer in advance.


## Definition

Causal attention masking is a mechanism that prevents a language model from using information from future positions in a text sequence when predicting the current or next word, ensuring that generation happens sequentially by only allowing attention to past and present tokens.


## How it connects

This concept sits directly beside Self-Attention and the standard Transformer architecture. In general self-attention, every part of a sentence can look at every other part simultaneously (bidirectional). Causal masking modifies this by blocking those backward-looking connections, turning the system into a generator that moves forward one step at a time. This is the key difference between an encoder (a system like BERT that reads and understands a whole block of text at once) and a decoder (a system like GPT that generates text word-by-word). Understanding this mask is essential for grasping how autoregressive training works, where the model predicts the next token to minimize error over the entire sequence.


## The mathematics

Causal masking in a Transformer model is a way of forcing the system to look only at past information, never at future information. Think of it like reading a book: you can only understand sentence B if you have already read sentence A. You cannot use the words in sentence C to help interpret sentence B.

In the mathematics of attention, the model calculates how much each position in a sequence should "pay attention" to every other position. This creates a grid of numbers called an attention score matrix. Let us call this matrix S. The entry S_ij represents the raw score for position i looking at position j.

If we have a sequence of length T, the matrix S has T rows and T columns. Row i corresponds to the current position being processed. Column j corresponds to the position being looked at.

To make this process "causal" (meaning time flows forward only), we must ensure that for any row i, the scores for all columns j where j is greater than i are removed. In other words, position i cannot look at positions i+1, i+2, etc.

We do this by introducing a second matrix, called the mask M. The mask has the same size as S (T rows and T columns). Every entry in M is either 0 or negative infinity (-∞).

The rule for building M is simple:
- If j is less than or equal to i (j <= i), meaning position j is in the past or is the current position, then M_ij = 0.
- If j is greater than i (j > i), meaning position j is in the future, then M_ij = -∞.

Why negative infinity? Because of how the next step works. The model converts these raw scores into probabilities using a function called softmax. Softmax involves taking the exponential of each number. The exponential of 0 is 1. The exponential of any large negative number is very close to 0. The exponential of negative infinity is exactly 0.

So, when we add the mask M to the scores S, the future scores become negative infinity. When softmax processes these, those future scores turn into exactly zero probability. The past scores (which were added to 0) remain unchanged and keep their relative weights.

The full mathematical operation for one step of causal attention is:
1. Calculate raw scores: S = Q dot K^T / sqrt(d_k). Here, Q and K are matrices of queries and keys, and d_k is a scaling factor (a number that keeps the values from getting too large).
2. Add the mask: S_masked = S + M.
3. Apply softmax: P = softmax(S_masked). This turns the scores in each row into probabilities that sum to 1.
4. Weight the values: Output = P dot V. Here, V is a matrix of values. The output for position i is a weighted average of all previous values, where the weights are given by P.

The critical order is: Mask must be added BEFORE softmax. If you masked after softmax, you would be setting probabilities to zero and then trying to normalize them again, which breaks the math because the remaining numbers would no longer sum to 1 in a way that reflects the original scores.

In summary, causal masking uses arithmetic with negative infinity to zero out future attention weights before they are converted into probabilities. This ensures the model only uses information available up to the current time step.


## A worked example

Let us walk through a concrete example with a very short sequence of 3 tokens. Let the positions be 1, 2, and 3.

Step 1: Define the raw attention scores.
Imagine we have already calculated the raw scores (before masking) for how much each position wants to look at every other position. These are stored in matrix S.
For position 1 (row 1), the scores to look at positions 1, 2, and 3 are: [2.0, 1.0, 0.5].
For position 2 (row 2), the scores to look at positions 1, 2, and 3 are: [1.5, 0.0, -1.0].
For position 3 (row 3), the scores to look at positions 1, 2, and 3 are: [0.8, 0.2, 0.6].

So S looks like this:
Row 1: [2.0, 1.0, 0.5]
Row 2: [1.5, 0.0, -1.0]
Row 3: [0.8, 0.2, 0.6]

Step 2: Build the causal mask M.
The rule is: if column index j <= row index i, put 0. If j > i, put negative infinity (-∞).

For Row 1 (i=1):
- Column 1 (j=1): 1 <= 1, so M_11 = 0.
- Column 2 (j=2): 2 > 1, so M_12 = -∞.
- Column 3 (j=3): 3 > 1, so M_13 = -∞.
Row 1 of M: [0, -∞, -∞]

For Row 2 (i=2):
- Column 1 (j=1): 1 <= 2, so M_21 = 0.
- Column 2 (j=2): 2 <= 2, so M_22 = 0.
- Column 3 (j=3): 3 > 2, so M_23 = -∞.
Row 2 of M: [0, 0, -∞]

For Row 3 (i=3):
- Column 1 (j=1): 1 <= 3, so M_31 = 0.
- Column 2 (j=2): 2 <= 3, so M_32 = 0.
- Column 3 (j=3): 3 <= 3, so M_33 = 0.
Row 3 of M: [0, 0, 0]

So M looks like this:
[0, -∞, -∞]
[0, 0, -∞]
[0, 0, 0]

Step 3: Add the mask to the scores.
We calculate S_masked = S + M. We add the numbers element by element.

For Row 1:
- Position 1: 2.0 + 0 = 2.0
- Position 2: 1.0 + (-∞) = -∞
- Position 3: 0.5 + (-∞) = -∞
Row 1 masked: [2.0, -∞, -∞]

For Row 2 (we will focus on this row for the detailed calculation below):
- Position 1: 1.5 + 0 = 1.5
- Position 2: 0.0 + 0 = 0.0
- Position 3: -1.0 + (-∞) = -∞
Row 2 masked: [1.5, 0.0, -∞]

For Row 3:
- Position 1: 0.8 + 0 = 0.8
- Position 2: 0.2 + 0 = 0.2
- Position 3: 0.6 + 0 = 0.6
Row 3 masked: [0.8, 0.2, 0.6]

Step 4: Apply softmax to Row 2.
Softmax converts the numbers in a row into probabilities. The formula for each element is:
P_j = exp(score_j) / sum(exp(all scores in row)).

Let's calculate this for Row 2 masked: [1.5, 0.0, -∞].

First, calculate the exponential (exp) of each number:
- exp(1.5) ≈ 4.4817
- exp(0.0) = 1.0
- exp(-∞) = 0.0

Next, sum these exponentials:
Sum = 4.4817 + 1.0 + 0.0 = 5.4817

Finally, divide each exponential by the sum to get the probabilities:
- Probability for position 1: 4.4817 / 5.4817 ≈ 0.8176
- Probability for position 2: 1.0 / 5.4817 ≈ 0.1824
- Probability for position 3: 0.0 / 5.4817 = 0.0

Check: 0.8176 + 0.1824 + 0.0 = 1.0. The probabilities sum to 1.

Result:
For position 2, the attention weights are approximately [0.82, 0.18, 0].
Notice that the weight for position 3 (the future) is exactly 0. Position 2 only attends to positions 1 and 2. This confirms the causal constraint: the model did not use any information from the future.

If we had NOT masked, the raw scores were [1.5, 0.0, -1.0].
exp(1.5) ≈ 4.48
exp(0.0) = 1.0
exp(-1.0) ≈ 0.37
Sum ≈ 5.85
Probabilities would have been [0.77, 0.17, 0.06].
Without masking, position 2 would have paid a small amount of attention to position 3 (the future). Masking removed this possibility completely.


## From scratch, in Python

```python
import math

# We'll use plain Python lists, no libraries.
# A list is a collection that keeps things in order. We'll have a list of lists.

# Suppose we have 3 positions, and a tiny attention score matrix (before softmax).
# In a real transformer, this comes from dot products of queries and keys.
# But for now, let's make up some numbers.
scores = [
    [2.0, 1.0, 0.5],  # row for position 0 (attending to positions 0,1,2)
    [1.5, 0.0, -1.0], # row for position 1
    [0.8, 0.2, 0.6]   # row for position 2
]

# Causal mask: for each position i, we allow columns 0 to i, and block columns i+1 onwards.
# We block by setting them to a very negative number, so that after exp() they become 0.
# We use -1e9 as a stand-in for negative infinity.

T = 3
masked = []
for i in range(T):
    row = []
    for j in range(T):
        if j > i:
            row.append(-1e9)
        else:
            row.append(scores[i][j])
    masked.append(row)

# Now apply softmax to each row.
# Softmax: for each number, compute exp(number), then divide by the sum of all exp in that row.
def softmax(row):
    # First compute exp of each, but to avoid overflow, we can subtract the max.
    # Since we have -1e9, exp(-1e9) is 0, but exp(2) is fine. We'll do it directly.
    exps = [math.exp(x) for x in row]
    total = sum(exps)
    return [e / total for e in exps]

attention = []
for row in masked:
    attention.append(softmax(row))

print("Masked scores (after adding mask):", masked)
print("Attention weights after softmax:")
for row in attention:
    print([round(w, 4) for w in row])
```


## The same thing, with a framework

In PyTorch (the library that powers most LLMs), this is done with `masked_fill` or directly with a boolean mask. For example:

```python
import torch
# scores: a tensor of shape (batch, heads, seq_len, seq_len)
# mask: a tensor of shape (seq_len, seq_len) with 0 where allowed, -inf where not
# We can create the causal mask as:
seq_len = 3
mask = torch.triu(torch.ones(seq_len, seq_len) * float('-inf'), diagonal=1)
# Then add it to scores before softmax:
masked_scores = scores + mask
attention = torch.softmax(masked_scores, dim=-1)
```

Under the hood, this does exactly what our from-scratch code does: it adds -inf to the future positions, then softmax turns them into zero. The `triu` function creates an upper triangular matrix, which is the same as our loop.


## What the code does

This code starts by defining a 3x3 matrix of attention scores (a list of lists). In real life, these come from multiplying queries and keys, but we just make them up. Then we loop over every position i and j. If j is greater than i (meaning a future position), we replace the score with -1e9 (effectively negative infinity). Otherwise, we keep the original score. This is the mask. Then we define a softmax function: for each number, we compute e to the power of that number, sum them all up, and divide each by the sum. The result is a list of probabilities that add to 1. For position 1 (the second row), after masking, the third column becomes -1e9, so its exp is 0, and the probabilities become [0.82, 0.18, 0]. So the attention weights are exactly as intended. The key takeaway: we mask before softmax, not after.


## Common misconceptions

**Masking is done by zeroing out the scores** — I can just multiply the future scores by zero to mask them.

If you multiply by zero, the score becomes 0. Then after softmax, that position gets some probability (non-zero), because softmax of 0 is exp(0)=1, which is a positive number. So zeroing doesn't zero the attention — it just makes it small. In fact: You must add negative infinity (or a very large negative number) to the score before softmax. Then after softmax, the probability becomes exactly zero because exp(-inf) is 0.

**Masking is applied after softmax** — I can compute attention, then set the probabilities for future positions to zero.

If you set probabilities to zero after softmax, the remaining probabilities no longer sum to 1. This breaks the interpretation of attention as a weighted average and the gradient flow during training. In fact: Always apply the mask to the scores (before softmax). That way, softmax produces a valid probability distribution over allowed positions only.

**Causal masking makes training sequential** — Because each position can only see earlier ones, the model must be trained token by token, one at a time.

The mask is applied to all positions in parallel. In one forward pass, the model computes attention for every position at once, with the mask preventing any info flow from future to past. So training can happen on the whole sequence at the same time, which is why GPUs are so effective. In fact: Causal masking enables parallel training while preserving the autoregressive property. The model is trained on all tokens at once, but each position only uses previous tokens.

**Masking is a feature of the model, not the data** — I can just mask the input data and the model will learn fine.

The mask is part of the attention computation, not data preprocessing. You can't achieve causal masking by trimming the input or shifting it. The model must know which positions are allowed to attend. In fact: The causal mask is a fixed matrix in the attention layer, ensuring the attention pattern is always causal regardless of the input.


## In short

Causal attention masking is the operation of setting future attention scores to negative infinity before the softmax, so each position only attends to itself and earlier positions. It is the key to making a transformer generate text, and it allows parallel training while preserving autoregressive semantics. Remember: mask before softmax, not after, and use negative infinity, not zero.


## Where this is used

Causal masking is the core of autoregressive language models like GPT. In a large language model (LLM), during training, you feed a long sequence of tokens (say a sentence of 1000 tokens). Without masking, each position could look at all 1000, including the ones after it, which would let the model 'cheat' and predict the next token without actually learning the patterns. The causal mask forces each position to only use earlier tokens, so the model learns to predict the next word from the previous context. This is what allows it to generate text during inference: when Sam Altman asks GPT-4 to 'write a poem about a cat', the model starts with a token and then, at each step, only uses the tokens it has already generated to produce the next one. The mask ensures that training and inference are consistent — the model never sees future tokens.

For small language models that run on your phone or a laptop (like TinyLlama or a model in the Ollama app), the causal mask is still there, but because the model is small, the mask computation is cheap. Developers use the same PyTorch code. The mask is applied in each attention layer. When you run a model like Llama 3.2 1B locally, the causal mask is a fixed matrix that is precomputed and reused. The mask also affects memory and speed: for a sequence of length N, the attention matrix is N x N, and the mask zeros out half of it, but you still compute all N^2 scores, so the computation is quadratic. That's why long contexts are expensive. Many optimizations focus on skipping the masked-out positions, but that's advanced.

In retrieval-augmented generation (RAG), a system retrieves relevant documents and then feeds them to a language model. The LM generates the answer autoregressively, so it still uses a causal mask. The retrieved context is part of the input sequence, and the mask ensures that when generating, the model only attends to the retrieved text and what it has already generated, not to future words.

In agentic systems that plan and call tools, an agent might have a loop: it thinks, calls a tool, gets a result, and then decides the next action. Each of these actions is a token sequence generated by an LLM, and the causal mask is applied at every generation step. The mask is what makes the generation possible — without it, the model couldn't produce the next action because it would be looking at a future action that doesn't exist yet.

In everyday applied machine learning, causal masking is less common. It's mostly for sequence generation. If you're building a recommender system, a fraud detector, or a speech model that doesn't generate, you probably don't need it. But if you're building anything that generates text, speech, or code, you'll use it.


---

Read this on the web: https://basavapurushottam.com/learn-ai/causal-attention-masking
