Transformer Decoder

The Transformer decoder layer builds on the encoder layer's structure but adds one crucial sublayer: cross-attention over the encoder's output — plus one crucial modification: its self-attention must be masked.

The Structure of One Decoder Layer

1. Masked multi-head self-attention over the decoder's own input so far
2. Add & norm
3. Multi-head cross-attention, querying the decoder's current state against the encoder's output (keys and values)
4. Add & norm
5. Position-wise feed-forward network
6. Add & norm

\(\mathbf{y}\) here is the decoder's own sequence-so-far (the tokens generated, or during training, the target sequence shifted right — connecting directly to Teacher Forcing).

Why Masking Is Essential Here

The encoder's self-attention can freely let every position see every other position, since the entire input is available upfront. The decoder is different: at generation time, position \(t\) has genuinely not yet seen positions \(t+1, t+2, \ldots\) — they don't exist yet. If the decoder's self-attention were allowed to attend to future positions during training, it would be trivially "cheating" (using information it won't actually have at inference time), and the model learned would be useless for real generation. Masked Self-Attention, the next note, covers exactly how this is prevented.

Why Cross-Attention Here, Specifically

This is exactly the cross-attention mechanism from Cross-Attention: the decoder's query (from its own, masked self-attention output) attends over the encoder's keys and values — letting every decoding step directly access the full, uncompressed representation of the input sequence, resolving the exact context-vector bottleneck that motivated attention in the first place, back in Context Vector.

Diagram — One Decoder Layer

Code

import torch
import torch.nn as nn

decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8, batch_first=True)
decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)

encoder_output = torch.randn(1, 10, 512)    # from the encoder
decoder_input = torch.randn(1, 5, 512)       # target tokens generated/available so far

# PyTorch requires an explicit causal mask for the decoder's self-attention
tgt_mask = nn.Transformer.generate_square_subsequent_mask(5)

output = decoder(decoder_input, encoder_output, tgt_mask=tgt_mask)
print(output.shape)   # (1, 5, 512)

Common Mistakes

• Forgetting to pass a causal mask to the decoder's self-attention — without it, the decoder would be trained with access to future tokens it won't have at inference time, producing a model that appears to perform excellently during training but fails badly at real generation.
• Mixing up which attention sublayer uses which K/V source — the decoder's self-attention uses the decoder's own sequence for Q, K, and V; only the cross-attention sublayer pulls K and V from the encoder's output.

Interview Relevance

Q: "What are the three sublayers in a Transformer decoder layer, in order, and why does the order matter?" Masked self-attention (builds context from the decoder's own sequence so far, without seeing future tokens), then cross-attention (queries the encoder's output for relevant source information), then a feed-forward network. The order matters because cross-attention needs the decoder's own self-attention output as its query — it can't meaningfully query the encoder's output before first building some representation of what's been generated so far.

Practice Question

Why does the decoder's self-attention need to be masked, while the encoder's self-attention doesn't need any masking at all?

Related DL Notes

• Transformer Encoder
• Masked Self-Attention
• Cross-Attention