🎓 A1Learner
🔍

PyTorch The Decoder Layer

A decoder layer is a transformer block that builds an output sequence one position at a time. It runs masked self-attention, then cross-attention over the encoder output, then a feed-forward network.

The parts inside one layer

A decoder layer holds two attention parts, not one. The second attention part is the cross-attention over memory.

import torch
from torch import nn
torch.manual_seed(0)
layer = nn.TransformerDecoderLayer(d_model=16, nhead=4, dim_feedforward=32, dropout=0.0, batch_first=True)
print([name for name, _ in layer.named_children()])

The causal mask

The mask adds 0 where reading is allowed and minus infinity where it is blocked. The minus infinity entries sit above the diagonal.

mask = nn.Transformer.generate_square_subsequent_mask(3)
print(mask)

Target and memory shapes

The target and the memory can hold different lengths. The output always takes the target's shape.

layer.eval()
tgt = torch.randn(2, 3, 16)
memory = torch.randn(2, 5, 16)
out = layer(tgt, memory, tgt_mask=nn.Transformer.generate_square_subsequent_mask(3))
print("Target:", tuple(tgt.shape))
print("Memory:", tuple(memory.shape))
print("Output:", tuple(out.shape))

Example

import torch
from torch import nn

torch.manual_seed(0)

layer = nn.TransformerDecoderLayer(
    d_model=16, nhead=4, dim_feedforward=32,
    dropout=0.0, batch_first=True)
layer.eval()

tgt = torch.randn(2, 3, 16)
memory = torch.randn(2, 5, 16)
print("Target:", tuple(tgt.shape))
print("Memory:", tuple(memory.shape))

mask = nn.Transformer.generate_square_subsequent_mask(3)
self_out, _ = layer.self_attn(tgt, tgt, tgt, attn_mask=mask)
print("Masked self-attention output:", tuple(self_out.shape))

after_self = layer.norm1(tgt + layer.dropout1(self_out))
cross_out, _ = layer.multihead_attn(after_self, memory, memory)
print("Cross-attention output:", tuple(cross_out.shape))

full = layer(tgt, memory, tgt_mask=mask)
print("Full layer output:", tuple(full.shape))

Key points

  • Memory: the encoder's output, which the decoder reads but never changes
  • Masked self-attention: attention on the target that cannot see later target positions
  • Causal mask: a mask that blocks every position from reading later positions
  • Cross-attention: attention where queries come from the target and keys and values come from memory
  • A decoder layer reads the target so far and the whole memory, then refines the target

In this example, the target batch holds two sequences of three vectors. The memory batch holds two sequences of five vectors from a pretend encoder. A causal mask is built for length 3. The masked self-attention runs on the target alone, so no position reads a later one. Its output is added back to the target and rescaled by the first norm. The cross-attention then takes its queries from that result and its keys and values from memory. The full layer call runs every part in order, the feed-forward network included. The final shape matches the target shape, because a decoder layer refines the target and leaves the memory alone.

Note

nn.TransformerDecoderLayer names its cross-attention part multihead_attn, not cross_attn

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

In a decoder layer, where do the cross-attention keys and values come from?