Cross-Attention

Cross-attention is the Q/K/V generalization of the original Bahdanau attention idea from Why Attention — the query comes from one sequence, but the keys and values come from a different sequence, letting one sequence selectively pull information from another.

The Defining Feature

\(\mathbf{X}_1\) and \(\mathbf{X}_2\) are two different sequences. In the machine translation setting this category started with: \(\mathbf{X}_1\) is the decoder's current representations (what's being generated), and \(\mathbf{X}_2\) is the encoder's output (the full, uncompressed representation of the input sentence) — every decoding step's query attends over every position of the source sentence's keys and values.

Cross-Attention IS the Formal Solution to Seq2Seq's Bottleneck

This is worth stating explicitly: cross-attention is the precise mathematical mechanism that resolves the exact limitation identified all the way back in Context Vector and Seq2Seq Limitations. Instead of the decoder relying on one fixed-size context vector, it now has a query at every decoding step that directly, selectively attends over every encoder position's key and value — nothing has to survive being compressed into a single vector anymore.

Self-Attention vs Cross-Attention, Side by Side

Diagram

Code

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CrossAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k

    def forward(self, decoder_seq, encoder_seq):   # TWO DIFFERENT sequences
        Q = self.W_Q(decoder_seq)     # query from the decoder
        K = self.W_K(encoder_seq)      # key from the encoder
        V = self.W_V(encoder_seq)      # value from the encoder
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
        weights = F.softmax(scores, dim=-1)
        return weights @ V

cross_attn = CrossAttention(d_model=16, d_k=8)
decoder_tokens = torch.randn(1, 3, 16)    # 3 tokens generated so far
encoder_tokens = torch.randn(1, 7, 16)     # 7 tokens in the source sentence
output = cross_attn(decoder_tokens, encoder_tokens)
print(output.shape)   # (1, 3, 8) -- one output per DECODER token, each informed by ALL encoder tokens

Common Mistakes

• Assuming cross-attention requires the two sequences to have the same length — they don't; the query sequence's length determines the number of outputs, while the key/value sequence's length determines how many positions each query can attend over, and these can differ freely.
• Forgetting that cross-attention's keys and values both come from the same (second) sequence — only the query comes from the first sequence; it's not a three-way split across three different sources.

Interview Relevance

Q: "How does cross-attention relate to the original attention mechanism introduced for machine translation?" Cross-attention is the formal Q/K/V generalization of exactly that original idea — the decoder's query attends over the encoder's keys and values, letting the decoder selectively draw on any part of the source sentence at every generation step. This directly replaces the single fixed-size context vector from basic Seq2Seq with a mechanism that preserves and makes accessible every encoder position's information.

Practice Question

Why must cross-attention's keys and values come from the same sequence as each other, even though the query comes from a different sequence?

Related DL Notes

• Self-Attention
• Why Attention
• Transformer Decoder