The mask is added to the scaled score matrix just before softmax; masked positions get -inf, which softmax converts to exactly 0 weight.
Imagine a vote where every candidate's score gets converted into a vote percentage by a rule that exponentiates first. To truly disqualify a candidate, you have to push their score down to 'negative infinity bad', which is the only score the rule converts to zero votes. A score of zero, surprisingly, still gives them about one vote out of N. So the disqualification has to happen at the right step (before counting) using the right value (negative infinity), or the disqualified candidate still ends up with influence.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Masking is one of those attention details that looks trivial on the surface and turns out to be surprisingly easy to get wrong. The mask has to be placed at exactly the right step in the pipeline, with exactly the right value, or it doesn't actually mask anything, the masked positions still leak attention to the output.
The right answer is: additive mask with -inf at forbidden positions, applied to the scaled score matrix immediately before softmax. Every other placement (on the embeddings, on the projections, on the post-softmax weights, on the W_O output) is either too early (no effect) or too late (breaks the math or fails to actually mask).
This deep dive walks through why softmax's normalization semantics force the mask placement, why -inf is the correct value (and what happens if you use 0 or -100 instead), how additive masking differs from multiplicative masking, and how production kernels like FlashAttention fuse the mask application with softmax for efficiency.
Why placement is forced by softmax semantics
Softmax has two properties that constrain mask placement:
- It always normalizes to sum to 1 along the axis it operates on.
- It's strictly positive: no finite input produces zero output.
These two properties mean the mask must act on the INPUTS to softmax, before the normalization. Acting on the outputs breaks the normalization without a clean way to recover it.
Before softmax (correct)
Add a mask tensor with 0 at allowed positions and -inf at forbidden positions to the scaled score matrix. Allowed scores unchanged; forbidden scores become -inf. Softmax then computes e^{-inf} = 0 for forbidden positions and renormalizes the remaining positions to sum to 1. Clean, no extra work.
After softmax (wrong)
If you zero out weights post-softmax, the row no longer sums to 1. The weighted-sum output sum_j a_ij V_j is biased low for the unmasked positions because some of the attention budget went to the now-zeroed forbidden positions. You'd have to renormalize the unmasked weights manually:
a_ij_new = a_ij / sum_{k unmasked} a_ik
This works but is extra work, and fused kernels like FlashAttention don't do it for you because they assume the standard pre-softmax masking pattern.
On embeddings or projections (wrong)
Zeroing the embedding of a masked position doesn't prevent attention from seeing the position. Even with input zero, the projections produce vectors (because the bias terms and the residual stream contribute), and attention still computes scores against those vectors. The mask has no effect.
On W_O output (wrong)
By the time you get to W_O, the attention weights have already been computed and used. The context vectors are already wrong. Zeroing some output positions doesn't undo the bad attention; it just hides the symptom while the cause persists.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
import torch
import math
B, H, T, D = 1, 1, 4, 8
q = torch.randn(B, H, T, D)
k = torch.randn(B, H, T, D)
v = torch.randn(B, H, T, D)
# Build causal mask: 1 at allowed (i >= j), 0 at forbidden
causal_mask = torch.tril(torch.ones(T, T))
# Compute scaled scores
scores = (q @ k.transpose(-2, -1)) / math.sqrt(D)
# Apply mask BEFORE softmax with -inf at forbidden positions
scores = scores.masked_fill(causal_mask == 0, float('-inf'))
# Softmax produces a valid distribution over unmasked positions
weights = torch.softmax(scores, dim=-1)
# Verify: row 0 attends only to col 0
assert weights[..., 0, 1:].sum().item() == 0.0
# Output
out = weights @ vReal products, models, and research that use this idea.
- PyTorch torch.nn.functional.scaled_dot_product_attention with is_causal=True applies the additive -inf mask internally before softmax.
- FlashAttention 2 and 3 fuse the mask application with the softmax computation, never materializing the full mask matrix.
What an interviewer would ask next. Try answering before peeking at the approach.
QHow would you detect a misplaced mask in production?
Add a behavioral assertion: at any query position t in a causal attention layer, the sum of post-softmax weights over positions greater than t should be exactly 0 (modulo float epsilon). Wire it as a unit test for any custom attention kernel. Also watch for training-inference quality mismatches: a leaky mask gives the model future-token information during training that disappears at inference, producing generation worse than perplexity predicts.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Applying the mask after softmax by zeroing weights. This breaks the row sum to 1 invariant and requires extra renormalization that fused kernels don't do for you.
60 second bullets to scan on the way to the call.
Exact placement of the mask in the attention pipeline
Why -inf is the correct mask value (softmax of -inf is 0)
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.