Masked Multi-Head Attention in Decoders

WHY IT EXISTS: Transformers use self-attention to draw global connections between tokens, but in an autoregressive decoder the future does not exist yet. If a model training to predict the next token could see the target sequence ahead, it would cheat and the training objective would not match inference. Masked multi-head attention was invented to let the decoder attend to its own past positions while strictly blocking future positions, enabling parallel matrix operations across the full sequence without leaking information.

THE MENTAL MODEL: Picture a team of specialists reading a novel where the ink ahead is physically blacked out. Each specialist looks for a different cue—one tracks grammar, another tracks character relationships, a third tracks tense. Every token in the sequence gets its own copy of this team, and each copy may only query the words it has already read. The mask is the blackout rule; the multiple heads are the specialists working simultaneously.

HOW IT WORKS: The input is projected into Query, Key, and Value matrices, then split into multiple heads, each with a reduced dimension. The attention scores are computed as the scaled dot product of Query and Key. Before the softmax, a causal mask—typically an upper-triangular matrix of negative infinity—is added to the scores. This forces all future positions to zero attention weight. After softmax, the result is multiplied by Value, the heads are concatenated, and a final linear projection produces the output. During training, this happens for every position in a single matrix multiplication, which is why GPUs can parallelize what conceptually looks like a left-to-right process.

WHEN TO USE IT: Use it in the decoder of any Transformer that generates tokens autoregressively, such as GPT-style language models or the decoder side of machine-translation systems. It is also the right mechanism whenever you need to train a sequence model on complete target sentences while preserving the causal left-to-right dependency structure required at generation time.

WHEN NOT TO USE IT: Do not use it in encoder-only models like BERT that rely on bidirectional context to build deep contextual embeddings. Applying a causal mask there would strip the model of future signal and severely degrade performance. It is also incorrect for encoder-decoder cross-attention, where the decoder should see the entire encoder output; the mask belongs only to decoder self-attention.

ONE CANONICAL EXAMPLE: In a GPT model generating the sentence "The cat sat," the token "sat" attends to "The" and "cat" but not to any later token. During training, the mask ensures that when the model predicts the third word, its context window is exactly the two preceding words, mirroring the incremental unrolling that happens during inference. If the mask were removed, the model would learn to copy the target directly and fail to generate coherent text when asked to produce new sequences.