Cross-attention lets a decoder use a separate encoded input while generating an output sequence. In an encoder–decoder transformer, it sits between masked self-attention and a feed-forward network in each decoder layer. Masked self-attention uses earlier output tokens; cross-attention connects the resulting decoder state to the encoder output. section 2.1.2 section 2.3.2
Trace one decoder layer:
- Masked self-attention updates the decoder state without exposing future output positions. section 2.3.2
- Cross-attention uses that state as its query and the encoder output, called memory, as both key and value. In image captioning, the queries come from caption representations and the keys and values come from image representations, each passed through linear layers. section 2.3.2 section 4.1.2
- A feed-forward sublayer processes the cross-attention output. Layer normalization and residual connections accompany the decoder sublayers. section 2.1.2
Listing 2.6 Creating a decoder layer
class DecoderLayer(nn.Module):
def __init__(self, size, self_attn, src_attn,
feed_forward, dropout):
super().__init__()
self.size = size
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
self.sublayer = nn.ModuleList([deepcopy(
SublayerConnection(size, dropout)) for i in range(3)])
def forward(self, x, memory, src_mask, tgt_mask):
x = self.sublayer[0](x, lambda x:
self.self_attn(x, x, x, tgt_mask)) #1
x = self.sublayer[1](x, lambda x:
self.src_attn(x, memory, memory, src_mask)) #2
output = self.sublayer[2](x, self.feed_forward) #3
return output
The decoder-layer implementation makes the data flow explicit: self_attn(x, x, x, tgt_mask) runs first, src_attn(x, memory, memory, src_mask) performs cross-attention second, and the feed-forward sublayer runs last. section 2.3.2
Conditioning a caption on an image
An image encoder produces image embeddings. Starting from a <start> token, the text decoder generates a caption one token at a time. At each decoder layer, cross-attention connects the current caption state to those image embeddings so the caption can align with image content. The same encoder–decoder pattern also supports German-to-English translation, with encoded German tokens conditioning English output. section 4.1.2 section 4.3.2 section 2.1.2
