Masked self-attention is regular self-attention with one specific, deliberate restriction: each position is prevented from attending to any position that comes after it in the sequence — essential for the decoder, where "future" tokens genuinely don't exist yet at generation time.
The Problem It Solves
Recall from Transformer Decoder: during training, the entire target sequence is available at once (via teacher forcing, see Teacher Forcing). Without any restriction, self-attention would happily let position 3 attend to positions 4, 5, 6 — information that, at real inference time, simply doesn't exist yet, since the model hasn't generated them. A model trained this way would learn to "cheat" by peeking ahead, and would then fail badly when deployed, since that future information is never actually available during real generation.
Formula
\(\mathbf{M}\) is a mask matrix: 0 for allowed (position \(j \le i\), i.e. current or earlier) positions, and \(-\infty\) for disallowed (future, \(j>i\)) positions. Adding \(-\infty\) to a score before softmax forces that position's post-softmax weight to become exactly 0 — the term \(e^{-\infty}=0\) — completely eliminating any influence from future positions, without needing any special-case logic in the softmax computation itself.
Numerical Example
For a 3-token sequence, the mask matrix looks like:
Row 1 (position 1) can only attend to position 1 itself. Row 2 (position 2) can attend to positions 1 and 2. Row 3 (position 3) can attend to all three positions. This exact upper-triangular pattern (with \(-\infty\) above the diagonal) is called a causal mask.
Diagram
The classic causal (upper-triangular) mask — each position can attend to itself and everything before it, never anything after.
Code
import torch
import torch.nn.functional as F
import math
def masked_self_attention(Q, K, V):
seq_len = Q.shape[-2]
d_k = Q.shape[-1]
scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() # True above the diagonal
scores = scores.masked_fill(mask, float('-inf')) # set future positions to -infinity
weights = F.softmax(scores, dim=-1) # softmax turns -inf into exactly 0 weight
return weights @ V
Q = K = V = torch.randn(1, 4, 8) # a 4-token sequence, self-attention
output = masked_self_attention(Q, K, V)
print(output.shape) # (1, 4, 8)Where This Is Used
Masked self-attention is used in the decoder of encoder-decoder Transformers, and it's the only kind of self-attention used throughout decoder-only architectures like GPT (covered fully in the NLP and LLM Fundamentals categories) — since generating text left-to-right always means position \(t\) has genuinely no access to tokens \(t+1\) and beyond.
Common Mistakes
- Using a mask of exactly 0 (rather than \(-\infty\)) for disallowed positions — this wouldn't remove those positions' influence at all; softmax would still assign them some (incorrectly non-zero) weight based on their actual score.
- Applying masking to the encoder's self-attention — the encoder sees the full input at once and has no "future" restriction; masking only applies to the decoder's self-attention (and equivalent decoder-only architectures).
Interview Relevance
Q: "Why does the decoder's self-attention need to be masked, and how is the masking actually implemented?" At real inference time, a decoder generating text left-to-right has genuinely not produced future tokens yet — allowing self-attention to see them during training would let the model cheat by using information it will never have available at inference, producing a model that fails at real generation. The mask is implemented by adding \(-\infty\) to the attention scores for disallowed (future) positions before the softmax step, which forces softmax to assign those positions exactly zero weight.
Practice Question
For a 5-token sequence with a causal mask, how many key positions can position 3 attend to?
