Cross-Attention: Shapes and What Training Actually Updates

1. Tracing shapes through cross-attention

Setup: d_model=4, num_heads=2 (so d_k=2), decoder has 3 tokens, encoder memory has 5 tokens.

query_input: (batch=1, tgt_len=3, d_model=4)   ← decoder's 3 tokens
kv_input:    (batch=1, src_len=5, d_model=4)   ← encoder's 5 tokens (memory)

Step 1 — split into per-head chunks (slice last dim into num_heads pieces of width d_k):

query_chunks = ( (1,3,2), (1,3,2) )   ← still 3 tokens long
kv_chunks    = ( (1,5,2), (1,5,2) )   ← still 5 tokens long

Sequence lengths don’t change here — only the feature width shrinks, split by head.

Step 2 — per-head projections (head i=0):

Q = W_Q[0](query_chunks[0])  →  (1, 3, 2)   ← still 3 decoder tokens
K = W_K[0](kv_chunks[0])     →  (1, 5, 2)   ← still 5 encoder tokens
V = W_V[0](kv_chunks[0])     →  (1, 5, 2)

Q and K/V now come from different sources — this is the actual cross-attention moment.

Step 3 — scores = Q @ K.transpose(-2, -1):

(1, 3, 2) @ (1, 2, 5)  →  (1, 3, 5)

Not square — a 3×5 grid: for each of the 3 decoder tokens (rows), how relevant is each of the 5 encoder tokens (columns).

Step 4 — mask + softmax. A source padding mask (1, 5) → .unsqueeze(1) → (1, 1, 5) broadcasts across the 3 query rows. softmax(scores, dim=-1) normalizes each of the 3 rows across the 5 source positions; shape stays (1, 3, 5).

Step 5 — output = attention_weights @ V:

(1, 3, 5) @ (1, 5, 2)  →  (1, 3, 2)

Key rule: the output’s sequence length always matches the query’s length (3), never the kv’s (5) — the kv dimension gets contracted away. For each query position you get one blended vector pulled from across all kv positions, weighted by relevance.

After the loop, both heads’ (1, 3, 2) outputs concatenate (dim=-1) into (1, 3, 4) — back to (batch, tgt_len, d_model). The decoder ends with as many tokens as it started with (3), each now carrying a blend of information pulled from the encoder’s 5 tokens.

General rule for any attention (self or cross): output shape mirrors the query’s shape, regardless of how long the thing being attended to is.

2. What training actually updates

Precision worth keeping straight: training updates W_Q, W_K, W_V, W_O — the learned parameters — not “the attention weights” as a stored, persistent thing. Attention weights are recomputed fresh on every forward pass from whatever the current weights are; training nudges those weights so the pattern that emerges gets more useful for predicting the correct target word.

For cross-attention specifically, this is the same “soft alignment” idea from Luong/Bahdanau-style RNN attention, just implemented as scaled dot-product attention instead of additive/multiplicative scoring:

  • Decoder’s W_Q learns what to ask for from the source, conditioned on what’s already been generated in the target so far — “I’m about to generate the next word, what source information do I need right now?”
  • Encoder’s W_K/W_V (reused as cross-attention’s keys/values) learn to produce representations that are useful to be retrieved — not just a good standalone encoding of the source language, but one specifically shaped by what the decoder needs to translate correctly. The encoder and decoder train end-to-end together, so gradients from the translation loss flow all the way back into the encoder’s parameters too.

The other two attention blocks learn different things in parallel: - Encoder self-attention — builds context-aware representations of the source sentence (e.g. disambiguating a word using its neighbors). - Decoder self-attention (causal) — learns to look back at previously generated target tokens for target-language fluency/grammar.

After enough training on parallel sentence pairs, inspecting cross-attention weights for a generated word should show it concentrating on whichever source token(s) that word actually corresponds to — the same kind of alignment pattern visible in classic attention heatmaps from earlier RNN-based translation models.

Comments