Causal Masking and the Autoregressive Constraint

An upper-triangular mask zeroes out future tokens so every position is trained to predict what comes next.

Why would you train a model to predict the next word if you let it see the next word?

That question sounds almost too obvious to be interesting. But the mechanism that enforces it, the causal mask, sits quietly at the center of every decoder-only transformer and shapes everything about how the model learns and generates. We touched on masking back in Post 3.7, where a block-diagonal mask prevented cross-document contamination in a packed row. Here we're going deeper. The causal mask is the mathematical constraint that makes next-token prediction work at all.

Prerequisites

You'll want to be comfortable with the attention score matrix from Post 4.1 and the softmax function. We'll also reference the packing masks from Post 3.7. If you've followed the arc so far, nothing here requires new math.

The information flow constraint

A language model trained with next-token prediction learns to assign probabilities to what comes next, given everything that came before. At position $i$ the model looks at tokens $0, 1, 2, \ldots, i$ and predicts token $i + 1$. That's the autoregressive assumption from Post 2.1, and it's the same constraint the original GPT paper used to scale next-token prediction to the whole stack.

So when the model is computing the representation for position $i$, it must not have any information about positions $i + 1, i + 2, \ldots, n$. If it did, the training signal would be meaningless. The model would just learn to copy the answer from its own input instead of learning the statistical structure of language.

This seems obvious during generation, when future tokens literally don't exist yet. You're producing one token at a time; there's nothing ahead to peek at. But during training, the entire sequence is sitting right there in the batch. All the tokens exist simultaneously. The model could, in principle, look at position 47 when computing the representation for position 12. We need a mechanism to prevent that.

How the mask works

The mechanism lives inside the attention computation. Recall from Post 4.1 that self-attention computes scores between every pair of positions, and the original transformer paper diagrammed this exact pattern as "Masked Multi-Head Attention":

For a sequence of $n$ tokens, $S$ is an $n \times n$ matrix where entry $S_{ij}$ tells us how much position $i$ wants to attend to position $j$. Without any mask, every position attends to every other position. That's bidirectional attention, and it's what BERT does.

For a decoder, we want to zero out the upper triangle. The causal mask $M$ is defined directly in additive form: zero on the lower triangle, $-\infty$ above the diagonal:

We apply it additively before softmax:

Allowed positions get $S_{ij} + 0 = S_{ij}$, and forbidden positions get $S_{ij} + (-\infty) = -\infty$. After softmax, any position that had $-\infty$ gets exactly zero weight. It contributes nothing to the output. The math is clean: $e^{-\infty} = 0$, so those entries drop out of the softmax normalization entirely.

The result is that position $i$'s output depends only on positions $0$ through $i$. Information flows backward in the sequence, never forward. The entire causal constraint reduces to this one matrix operation, and you can read the residual stream as a strict left-to-right write history because of it.

The growing triangle

Click through the tokens above, starting from "The" and working your way down. You'll see the pattern immediately.

Position 0 ("The") attends only to itself. One token in its context. Position 1 ("cat") attends to positions 0 and 1. Two tokens. Position 2 ("sat") gets three. Position 5 ("mat") attends to all six.

This is the growing triangle of causal attention. Each position has a strictly larger receptive field than the one before it. The first token in any sequence is the most impoverished: it has to predict the second token based on nothing but its own embedding. The last token is the richest, with the entire sequence history available to it.

When people talk about "context length," they're really talking about the size of this triangle for the final position. A 128K context window means the last token can attend to 128,000 previous positions. But the first token in that window always attends to exactly one position, regardless of how long the context is.

Why this makes training efficient

During training, the causal mask lets us compute loss at every position in the sequence simultaneously.

Think about how an RNN trains. It processes token 0, produces a hidden state, processes token 1, produces a new hidden state, and so on. Each step depends on the previous step's output. You can't parallelize it. For a sequence of 2,048 tokens, you need 2,048 sequential steps.

A transformer with a causal mask does something fundamentally different. We feed all 2,048 tokens in at once. The attention computation produces an output for every position in parallel. Position 0's output predicts token 1. Position 1's output predicts token 2. Position 2,047's output predicts token 2,048. We get 2,048 loss terms from a single forward pass.

This works because the causal mask guarantees each position's output only depends on tokens at or before that position. Position 12's representation is computed using only tokens 0 through 12, so it's a valid basis for predicting token 13, even though tokens 13 through 2,047 are sitting right there in the same batch.

This is called teacher forcing. During training, the model always conditions on the ground-truth prefix rather than on its own predictions. It's the natural consequence of feeding a full sequence through a causally masked transformer. The mask is what makes it valid, and the parallelism is what makes it fast.

Bidirectional vs. causal: the fundamental fork

Not every transformer uses causal attention. BERT uses bidirectional attention. Every token attends to every other token. No mask at all (well, a padding mask, but no causal mask). This gives each token access to the full sequence context, which is great for understanding tasks. When BERT computes the representation of the word "bank" in "I walked along the river bank," it can look at "river" to disambiguate, even though "river" comes earlier.

A causal model processing the same sentence would compute the representation of "bank" using only "I walked along the river bank." It can also see "river," because "river" comes before "bank." But in the sentence "the bank of the river was muddy," a causal model computing "bank" would not be able to see "river" at all. A bidirectional model would.

So bidirectional attention gives richer representations. Why doesn't everyone use it?

Because you can't generate text with bidirectional attention. Generation is inherently autoregressive: you produce one token, then the next, then the next. Each new token has to condition on the tokens before it and nothing else. Bidirectional attention would need to condition on tokens that don't exist yet. The causal mask is what makes generation possible.

This is the fundamental fork in transformer architectures. Encoder models (BERT, RoBERTa) use bidirectional attention for understanding tasks. Decoder models (GPT, Llama) use causal attention for generation. We'll get deeper into this in Post 4.7 when we look at encoder-decoder architectures like T5, which use both.

Worked example

Let me trace through a small example to make the numbers concrete. Say we have 5 tokens and our attention scores $S$ (after scaling by $\sqrt{d_k}$) look like this:

The causal mask sends the upper triangle to $-\infty$:

After softmax (applied row-wise), the $-\infty$ entries become exactly 0:

Look at row 0. Position 0 puts all its weight on itself because it has no other option. Row 1 splits its weight between positions 0 and 1, with more on position 1 (the higher score). Row 4 distributes across all five positions. Each row sums to 1, and all the upper-triangle entries are exactly 0.

The bidirectional version on the same scores would spread row 0 across all 5 positions instead of concentrating everything on itself. The representations would be completely different.

Prefix-LM: a middle ground

Some models split the difference. In a prefix-LM architecture, a designated prefix of the sequence uses bidirectional attention (every prefix token can see every other prefix token) while the rest uses causal attention. T5 and PaLM use variants of this.

The intuition is practical. If you're conditioning on a prompt and then generating a response, the prompt is fixed. No generation happens there, so there's no reason to restrict it to causal attention. Let the model see the full prompt bidirectionally for richer representations, then switch to causal for the generation part.

In mask-matrix form it looks like a block of ones in the upper-left (the prefix portion, fully bidirectional) plus a standard causal triangle in the lower-right (the generation portion). The generation tokens can attend to all prefix tokens but not to future generation tokens.

This is an active design space. Different models draw the boundary differently, and some make it dynamic. But the principle is the same: use bidirectional attention where generation isn't happening, and causal attention where it is.

Connection to packing masks

In Post 3.7 we built block-diagonal masks to prevent cross-document contamination during training. Those masks serve a different purpose than the causal mask, but in practice the two get combined.

The causal mask says "don't look at future tokens." The packing mask says "don't look at tokens from other documents." When you pack multiple sequences into a single training row with causal attention, you need both. The result is a block-diagonal matrix where each block is lower-triangular. Each document gets its own causal triangle along the diagonal, and everything outside the blocks is masked.

FlashAttention's flash_attn_varlen_func handles this natively through the cu_seqlens interface, so you never have to materialize the combined mask as a dense matrix. You pass the cumulative sequence lengths and set causal=True, and the fused kernel handles the rest.

Implementation

The causal mask itself is trivial to build. In PyTorch:

import torch
import torch.nn.functional as F
 
def causal_attention(Q, K, V):
    """
    Scaled dot-product attention with causal masking.
    Q, K, V: (batch, heads, seq_len, head_dim)
    """
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5)
 
    # Build the causal mask: lower-triangular matrix of True values
    seq_len = scores.size(-1)
    mask = torch.tril(torch.ones(seq_len, seq_len, device=scores.device))
 
    # Set the upper triangle to -inf so softmax zeros it out
    scores = scores.masked_fill(mask == 0, float('-inf'))
 
    weights = F.softmax(scores, dim=-1)
    return weights @ V
 
# Example: 4 tokens, 1 head, head_dim = 8
Q = torch.randn(1, 1, 4, 8)
K = torch.randn(1, 1, 4, 8)
V = torch.randn(1, 1, 4, 8)
 
output = causal_attention(Q, K, V)
# output[0, 0, 0] depends only on V[0, 0, 0]    (position 0)
# output[0, 0, 1] depends on V[0, 0, 0:2]       (positions 0-1)
# output[0, 0, 3] depends on V[0, 0, 0:4]       (positions 0-3)

That's the explicit version. In practice you almost never write this yourself. FlashAttention's flash_attn_func takes a causal=True flag and handles the masking inside a fused CUDA kernel, never materializing the $n \times n$ mask matrix at all. That matters for long sequences. At 8,192 tokens the mask would be a 67-million-element matrix. At 128K tokens it would be 16 billion. The fused kernel avoids all of that.

Misconceptions

"The mask removes attention to those positions." The mask doesn't delete anything. It's an additive operation that runs before softmax. We add $-\infty$ to the upper-triangle scores; softmax then drives those entries to exactly zero in the output weights because $e^{-\infty} = 0$. The Q, K, and V projections still get computed for every position. The mask only suppresses the contributions of future keys to past queries, which is enough to enforce causality.

"Encoders see the future, decoders never do." True for self-attention, oversimplified for the full encoder-decoder picture. In a model like T5, the encoder is bidirectional self-attention and the decoder uses causal self-attention. But the decoder also has a cross-attention layer that attends to the encoder's outputs, and that cross-attention is unmasked. The decoder is allowed to look at the entire encoded source (every token of the input it's translating from, say) while generating each output token. "Causal" applies to the decoder's view of its own generated sequence, not to the source it's conditioning on.

"The mask is a hack." The implementation (slap −∞ on the upper triangle) can look like engineering papering-over a deficiency in the architecture. It's actually the direct mathematical implementation of the autoregressive factorization. The chain rule says $P(x_1, \ldots, x_n) = \prod P(x_i | x_1, \ldots, x_{i-1})$. The causal mask is what makes each conditional in that product well-defined inside a parallel attention computation. Remove the mask and you don't have an autoregressive model anymore.

What's next

In Post 4.4 we'll look at how the transformer knows the order of tokens. The causal mask controls which tokens are visible, but nothing in the attention computation itself knows which position is which. That's where positional encodings come in.

Additional reading (and watching)

• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. Introduces scaled dot-product attention with the causal mask diagram that became the canonical picture.
• Radford, A., et al. (2018). Improving Language Understanding by Generative Pre-Training. OpenAI. The original GPT paper, where causally-masked next-token prediction scaled to the whole stack.
• Devlin, J., et al. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL 2019. The bidirectional counterexample, and why MLM replaces next-token prediction when you drop the causal mask.
• Raffel, C., et al. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR 2020. Formalizes prefix-LM attention, the middle-ground between bidirectional and causal.
• Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. Fuses the causal mask into the kernel so the upper triangle is never computed, making causal attention strictly cheaper than full attention.
• Jiang, A. Q., et al. (2023). Mistral 7B. arXiv:2310.06825. Introduces sliding-window attention as a band-shaped extension of the causal mask, paired with a rolling KV cache.
• Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic. The residual-stream framework that motivates thinking of the causal mask as the information-flow rule on the stream.