Predict the softmax output for a causal masked attention row.
Position 2 (0 indexed) in a length 4 sequence has raw attention scores [1.0, 2.0, 3.0, 4.0] before masking. Apply causal masking (position 2 can attend to positions 0, 1, 2 but NOT position 3), then softmax. What are the 4 output weights, in order?
Mask sets position 3's score to negative infinity before softmax, giving the four weights approximately 0.0900, 0.2447, 0.6652, 0.0000.
Imagine a teacher converting test scores into slices of a pie chart where every slice has to add up to one whole pie. To remove a student entirely, you can't just give them a score of zero, a zero score still claims a sliver of the pie, because the conversion rule turns zero into a real positive number first. You have to give them an impossibly bad score, think 'negative infinity bad', which then converts to no pie at all. The remaining students share the whole pie among themselves, like causal masking sweeps a future position out of the attention distribution.
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.
Attention masking looks trivial, change one number and run softmax, but the choice of mask value is one of the most common implementation bugs in custom attention kernels. Getting it wrong silently leaks attention to the masked position, breaks the autoregressive property, and creates a hard to localize training inference mismatch.
This question is a small numeric exercise, but it tests whether you understand the underlying mechanic: masking acts on the pre-softmax score, and only negative infinity drives the post-softmax weight to exactly zero.
The walkthrough below covers the correct procedure, the arithmetic step by step, the two common bug modes, and the numerical details that come up at the kernel level.
The correct procedure
Causal masking happens before softmax. For position t in a length-N sequence, every position j greater than t is set to negative infinity in the score row, while positions less than or equal to t are left alone.
Why before, not after
The softmax normalizes whatever it receives into a probability distribution. If you apply softmax first and then zero out masked weights, the remaining weights no longer sum to one and you have to manually renormalize. That is wasteful and a known source of subtle bugs.
Why negative infinity, not zero
The exponential function never returns zero for any finite input. Setting a score to zero gives e to the zero, which is one, meaning the masked position still claims a real share of the probability mass. Only negative infinity sends the exponential to exactly zero.
The mask value has to act on the exponent, not on the post-softmax weight. Negative infinity is the only value that makes the masked position vanish from both numerator and denominator simultaneously.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
import torch
scores = torch.tensor([1.0, 2.0, 3.0, 4.0])
mask = torch.tensor([0.0, 0.0, 0.0, float('-inf')])
masked_scores = scores + mask # [1, 2, 3, -inf]
weights = torch.softmax(masked_scores, dim=-1)
print(weights) # tensor([0.0900, 0.2447, 0.6652, 0.0000])
assert weights[3].item() == 0.0
assert abs(weights[:3].sum().item() - 1.0) < 1e-6Real products, models, and research that use this idea.
- PyTorch's torch.nn.functional.scaled_dot_product_attention with is_causal=True applies the canonical mask internally and selects an efficient backend.
- HuggingFace transformers uses torch.finfo(dtype).min as the mask value, which is roughly -3.4e38 for fp32.
What an interviewer would ask next. Try answering before peeking at the approach.
QWhy use a large finite negative instead of literal negative infinity in production kernels?
Numerical safety inside fused kernels. Some accumulator paths can propagate -inf through intermediate reductions and produce NaN. A value like -1e9 in fp32 or torch.finfo(dtype).min in fp16 makes e to the value underflow to zero far below precision while avoiding -inf propagation through fused ops.
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.
Zero as the mask value leaks ~3% weight to the masked slot. e^0 is 1, not 0. Only negative infinity vanishes.
60 second bullets to scan on the way to the call.
Where the mask applies relative to softmax
Why negative infinity and not zero as the mask value
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.