Arc 4: Transformers from First Principles

Causal Masking and the Autoregressive Constraint

12 min readbuilds on Packing, Masking, and Tokenization as a Model Interface, Self-Attention: Q, K, V from First Principles and Language Modeling as Next-Token Prediction

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?

The causal mask answers it by restricting each position to itself and earlier tokens. We used block-diagonal masks in Post 3.7 to keep packed documents separate. Here the mask lets us compute next-token losses for a whole sequence in parallel without leaking the answers.

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.

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.

During left-to-right generation, future tokens do not exist yet. During training, however, the full sequence is present in the batch. Without a mask, position 12 could use the token at position 47 while learning to predict token 13.

How the mask works

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":

S = \frac{QK^T}{\sqrt{d_k}}

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:

M_{ij} = \begin{cases} 0 & \text{if } j \leq i \\ -\infty & \text{if } j > i \end{cases}

We apply it additively before softmax:

S_{\text{masked}} = S + M

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. Since $e^{-\infty} = 0$ , those entries drop out of the softmax normalization entirely.

The result is that position $i$ 's output depends only on positions $0$ through $i$ . Queries look backward in the sequence; information flows from earlier positions into later ones, never from future tokens into past representations. The constraint concerns information access, not execution order: all positions in a layer can still be computed in parallel.

How the mask is applied: additive, before softmax

row 5 of 5

Follow one row from raw scores to probabilities. Masked positions become −∞ before softmax and zero afterward.

Attention Mask Pattern

attendedmasked (-inf)

Click a token to see its attention pattern

Toggle the mask and select a token to compare its accessible positions.

The growing triangle

Select the tokens above, starting from "The" and working your way down.

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 position has only its own token and position information as input, transformed through the network. The last token can use the entire sequence prefix.

When people talk about "context length," they're really talking about the size of this triangle for the final position. A full-attention window of 128K tokens lets the last position attend to itself and the preceding positions within that window. At the start of an independently processed sequence, the first position sees only itself. A rolling window can also contain cached earlier context.

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. The recurrence along this sequence is sequential. For a sequence of 2,048 tokens, you need 2,048 sequential steps.

For a transformer with a causal mask, 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 if the batch supplies the next token as a target for each input, including a target beyond the last input position. Shifting labels within a 2,048-token chunk instead gives 2,047 targets.

Illustrative next-token losses

sequential vs causal-masked

An RNN has to step through the sequence one token at a time, producing one loss per step. A causal transformer produces every position's loss in a single forward pass, because the mask guarantees each position's output depends only on its own past.

This training recipe is called teacher forcing: each position conditions on the ground-truth prefix. The causal mask prevents it from reading its target, while attention computes the positions in parallel.

Bidirectional and causal attention

BERT uses bidirectional attention, with padding and task-specific masks but no causal mask. In “the bank of the river was muddy,” its representation of “bank” can use “river,” even though it comes later in the sequence.

A causal model cannot use “river” at that position. It can use it in “I walked along the river bank,” where “river” comes before “bank.”

For left-to-right next-token training, a bidirectional model could read its target from the input. The causal mask prevents that leakage and lets the same computation support incremental generation. Other generation methods, including iterative masked-token refinement, use different objectives and need not be left-to-right.

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.

Tap or hover a cell to inspect its weight. Darker cells have larger weights on a fixed 0–1 scale.

Tap or hover cells to inspect weights. Select a cell or row to see the weighted sum.

Bidirectional (BERT-style): every row attends to every column. 'bank' can see 'river' in either direction.

Tap or hover a cell to inspect its weight. Darker cells have larger weights on a fixed 0–1 scale.

Tap or hover cells to inspect weights. Select a cell or row to see the weighted sum.

Causal (GPT-style): the upper triangle is zero. Each row sees only itself and tokens to its left.

Attention Matrix — Decode Steps

The grid is a history of attention computations. The cache stores K and V activations, not these attention rows. Prefill logits select the first continuation token; processing that token creates the next row and predicts the following token.

previously computed rowsmasked (future)

4 tokens prefilled — click "Process next token"

The first four positions represent causal prompt prefill. Processing a sampled token computes a new attention row over the prefix and itself, then predicts the following token. The KV cache stores projected keys and values, not these attention rows.

Worked example

Say we have 5 tokens and our attention scores $S$ (after scaling by $\sqrt{d_k}$ ) look like this:

S = \begin{bmatrix} 0.5 & 0.2 & 0.8 & 0.1 & 0.3 \\ 0.4 & 0.9 & 0.1 & 0.6 & 0.2 \\ 0.7 & 0.3 & 0.5 & 0.4 & 0.8 \\ 0.2 & 0.6 & 0.3 & 0.7 & 0.1 \\ 0.3 & 0.5 & 0.2 & 0.4 & 0.6 \end{bmatrix}

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

S_{\text{masked}} = \begin{bmatrix} 0.5 & -\infty & -\infty & -\infty & -\infty \\ 0.4 & 0.9 & -\infty & -\infty & -\infty \\ 0.7 & 0.3 & 0.5 & -\infty & -\infty \\ 0.2 & 0.6 & 0.3 & 0.7 & -\infty \\ 0.3 & 0.5 & 0.2 & 0.4 & 0.6 \end{bmatrix}

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

\text{softmax}(S_{\text{masked}}) \approx \begin{bmatrix} 1.0 & 0 & 0 & 0 & 0 \\ 0.38 & 0.62 & 0 & 0 & 0 \\ 0.40 & 0.27 & 0.33 & 0 & 0 \\ 0.1906 & 0.2844 & 0.2107 & 0.3143 & 0 \\ 0.18 & 0.22 & 0.16 & 0.20 & 0.24 \end{bmatrix}

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 resulting representations can differ because they mix different value vectors.

Tap or hover a cell to inspect its weight. Darker cells have larger weights on a fixed 0–1 scale.

Tap or hover cells to inspect weights. Select a cell or row to see the weighted sum.

Post-softmax attention weights from the worked example. Tap or hover over cells to inspect the rounded values. The zeros in the upper triangle are the causal mask at work.

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. The T5 study compared this prefix-LM variant with causal and encoder-decoder alternatives; its released T5 models use encoder-decoder stacks.

In a prefix-LM, the fixed prompt can attend bidirectionally because the objective only predicts the continuation. That changes the training recipe: ordinary next-token losses inside the bidirectional prefix would leak their targets. A causal prompt remains useful when training on every position or reusing causal-prefix computation.

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.

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 excludes future tokens; the packing mask excludes other documents. Training recipes that require both constraints use a block-diagonal mask with a lower-triangular block for each document. Recipes that permit cross-document attention omit the document-isolation constraint.

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

In PyTorch, we can build the mask with an upper-triangular tensor:

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, dtype=torch.bool))
 
    # Set the upper triangle to -inf so softmax zeros it out
    scores = scores.masked_fill(~mask, 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)

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. 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." Q, K, and V are still projected for every input position. The mask suppresses future keys’ contributions to earlier queries; it does not remove those tokens or their representations from the batch.

"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." Adding −∞ above the diagonal implements the autoregressive factorization inside attention. 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 from this parallel training computation and later inputs can leak into predictions for earlier targets.

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. Uses tiled attention and causal masking without allocating a dense mask; kernel speed depends on shapes and hardware.
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.