Causal Language Modeling › Causal Masking

Michael BrenndoerferUpdated July 7, 20254 min read
Language AI HandbookAll chapters
Part 22 of 67Pre-training Objectives
Chapter 155Causal Language Modeling › Causal Masking
Reading section 5 of 13 · Causal Masking · 4 minSection 5 of 13 · 4 min

Book contents

Section 5 of 13 · 4 min read

Causal Language Modeling › Causal MaskingOne section · Sections in this chapter5 / 13 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

For the autoregressive factorization to hold, the model must not "peek" at future tokens when predicting the current one. In transformer architectures, this is enforced through causal masking in the attention mechanism. Without this constraint, a model could trivially solve the prediction problem by looking ahead at the answer, which would produce a model that learns nothing useful and cannot generate text at test time.

The challenge is a structural one. Standard transformer attention is fully bidirectional: every position attends to every other position. This is powerful for understanding tasks, where you want to incorporate all available context. But for generation, it creates a problem. If position 3 can see positions 4, 5, and 6 during training, the model learns to use that future information. At inference time, positions 4, 5, and 6 have not been generated yet, so the model encounters a distribution shift it was never trained to handle.

Causal masking solves this by zeroing out attention from any position to any future position. The model is structurally prevented from accessing information it would not have access to during generation. In practice, this is implemented by adding a large negative constant (negative infinity) to the attention scores for disallowed positions before applying the softmax. After softmax, those positions receive attention weight of zero, as if they do not exist.

The attention pattern looks like this:

Out[10]:
Visualization
Heatmap showing a lower-triangular causal attention mask where each position can only attend to itself and earlier positions.
Causal attention mask for a sequence of 6 tokens. White cells indicate allowed attention (value 0), dark cells indicate blocked attention (value negative infinity). Position 3 can attend to positions 0, 1, 2, and 3, but not to positions 4 or 5. This lower-triangular pattern ensures information flows only from past to present.

In code, the causal mask is applied to attention scores before the softmax:

In[11]:
Code
def create_causal_mask(seq_len):
    """
    Create a causal attention mask.

    Returns a matrix where position (i, j) is:
    - 0 if i >= j (can attend)
    - -inf if i < j (cannot attend)
    """
    # Create upper triangular matrix of ones (above diagonal)
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)

    # Replace 1s with -inf to zero out attention after softmax
    mask = mask.masked_fill(mask == 1, float("-inf"))

    return mask


# Example usage in attention
def masked_attention(query, key, value, mask):
    """Simplified scaled dot-product attention with causal masking."""
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / (d_k**0.5)

    # Apply causal mask
    scores = scores + mask

    # Softmax: -inf becomes 0 probability
    attn_weights = F.softmax(scores, dim=-1)

    return torch.matmul(attn_weights, value), attn_weights
Out[12]:
Console
Causal mask (4 positions):
[[  0. -inf -inf -inf]
 [  0.   0. -inf -inf]
 [  0.   0.   0. -inf]
 [  0.   0.   0.   0.]]

Attention weights after masking and softmax:
[[1.    0.    0.    0.   ]
 [0.5   0.5   0.    0.   ]
 [0.333 0.333 0.333 0.   ]
 [0.25  0.25  0.25  0.25 ]]

The negative infinity values become zero probability after softmax, effectively blocking information flow from future positions. Each row sums to 1.0, distributing attention only over valid (past and present) positions.