Masked and Causal Attention

Masked attention is not a separate scoring rule. It is a visibility constraint inserted between scaled query–key scores and row-wise softmax. In a causal decoder, query row i may read keys 0…i but not keys i+1…N−1.

That lower-triangular boundary lets a decoder learn every next-token target in parallel under teacher forcing without exposing a row to the future tokens it must predict. During generation, the same dependency rule is realized one new token at a time.

Seal future keys before softmax

Play the full mechanism or choose any stage. The matrix keeps query rows vertical and key columns horizontal so the causal triangle does not change orientation between steps.

The exact operation

For one attention head, first compute the usual scaled scores:

Using zero-based row and column indices, define an additive causal bias:

Add the bias before normalizing across the key axis:

Then mix values as usual:

For query row i, every future entry has conceptual logit −∞, so its normalized weight is zero. The remaining allowed entries are positive and sum to one. A valid causal row always retains at least its own key on the diagonal.

The order matters. Masking probabilities after softmax and failing to renormalize leaves an invalid row sum. Masking values instead of logits also changes the operation: a forbidden key would still consume probability mass.

Why the target must be shifted

Suppose a training sequence is

<BOS>  A  small  robot  paints  stars  <EOS>

The decoder receives the left-shifted inputs <BOS> A small robot paints stars, while the loss compares its outputs with A small robot paints stars <EOS>. Row i therefore predicts the token immediately to its right.

Without a causal boundary, row i could directly read future input rows that contain its target or later ground-truth tokens. That shortcut is unavailable during ordinary left-to-right generation. The mask removes it while teacher forcing still supplies the correct prefix to each row.

This is a dependency constraint, not a statement that earlier tokens caused a model’s final decision in the interpretability sense. Attention weights alone are not causal attributions.

Parallel training and sequential generation

The architecture is causal in both regimes, but the execution shapes differ.

Transformer layers still execute in order during training; “parallel” refers to positions within a layer and to their loss terms. During decode, a KV cache avoids recomputing K/V for earlier prefix tokens, but it does not remove the next-token dependency between generation steps.

Conceptual −∞ versus implementation values

The mathematical mask uses −∞ because its exponential is zero. Several implementations are valid:

• add literal −∞ to forbidden logits when every row has at least one finite entry,
• use a sufficiently negative, dtype-aware finite sentinel,
• pass a boolean mask whose polarity is defined by that API,
• tell a fused attention kernel that the operation is causal.

An all-blocked row needs special handling: softmax over only −∞ entries is undefined and can produce NaN. Padding and causal masks must therefore be combined so valid query rows retain at least one key, while outputs or losses for padded query rows are ignored as appropriate.

Boolean polarity is not universal. In PyTorch’s scaled-dot-product attention API, True means a pair participates; other APIs use True to mean “blocked.” Check the exact function rather than assuming masks are interchangeable.

Shape-faithful reference

import math
import torch


def causal_attention(q, k, v):
    # q, k, v: [batch, heads, tokens, head_width]
    n = q.shape[-2]
    scores = q @ k.transpose(-2, -1)
    scores = scores / math.sqrt(q.shape[-1])

    # True means visible in this reference implementation.
    visible = torch.ones(n, n, dtype=torch.bool, device=q.device).tril()
    scores = scores.masked_fill(~visible, float('-inf'))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ v
    return output, weights

For production code, torch.nn.functional.scaled_dot_product_attention can accept causal intent directly. The selected FlashAttention, memory-efficient, or math backend depends on framework version, device, dtype, shape, and other options; causal masking alone does not guarantee a particular kernel or speedup.

Pair counts and runtime boundaries

For sequence length N, a causal pattern contains

The six-token figure therefore has 21 allowed and 15 forbidden positions. These are semantic pair counts. A dense reference implementation can still construct an N × N score tensor and then mask its upper triangle.

A causal-aware fused kernel can avoid materializing an N × N mask and may skip work for forbidden tiles. Whether that changes memory traffic or latency depends on the backend and workload. The causal pattern by itself does not make dense attention linear in sequence length.

Other masks use different boundaries

“Masked attention” is broader than causal attention:

• Padding mask: blocks padded key positions so variable-length examples can share a batch. Padded query outputs or loss terms must also be ignored appropriately.
• Prefix-LM mask: allows bidirectional visibility inside a supplied prefix, then uses causal visibility through the continuation.
• Sliding-window or block-sparse mask: exposes selected local or global key regions. It reduces arithmetic only when the implementation exploits the sparse structure rather than computing a dense matrix and masking afterward.
• Cross-attention mask: often blocks padded source keys. Decoder-to-encoder cross-attention is usually noncausal over the available source, but streaming or task-specific systems can impose additional restrictions.

Masks can be combined by intersecting their allowed pair sets. For example, a batched decoder commonly needs both a causal boundary and a padding boundary.

Common implementation errors

• Applying the mask after softmax: block logits first, then normalize the surviving keys.
• Using the wrong triangle: with query rows and key columns, causal visibility is on and below the main diagonal (j ≤ i). Tensor conventions must be checked explicitly.
• Masking the wrong axis: softmax normalizes keys independently for every query row and head.
• Treating targets as same-row inputs: next-token labels are shifted and belong to the loss, not to that query’s visible context.
• Creating all-masked rows: define padding behavior so valid rows retain an allowed key and ignored rows do not enter the loss.
• Assuming mask booleans have universal polarity: APIs disagree on whether True means allowed or blocked.
• Equating half the semantic pairs with half the runtime: measure the actual kernel, shapes, and hardware.

Primary sources

• Attention Is All You Need
• PyTorch scaled dot-product attention
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Related concepts