Cross-Attention in transformers

Module 05 · Transformers · 19 concepts

In depth / 15 lines

Cross-Attention, in more detail

Cross-attention is attention where the questions and the answers come from different places.

The queries come from the sequence being produced.

The keys and values come from a separate sequence being consulted.

That is the only difference from self-attention, and it is enough to change what it is for.

Self-attention is a sequence understanding itself.

Cross-attention is one sequence reading another.

A translation decoder uses it to look back at the source sentence while writing.

Each word being generated attends to whichever source words are relevant to it.

That is what dissolved the fixed-vector bottleneck in sequence-to-sequence models.

An image generator uses it to condition pixels on a text prompt.

The image is being written and the prompt is being read, which is the same pattern.

Speech, video and any encoder-decoder model use it the same way.

It also joins the modalities in multimodal models, letting text attend to image features.

Decoder-only language models have no cross-attention, since everything lives in one stream.

Its weights are unusually readable, since they show which source token each output word drew on.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Cross-Attention

Queries from one sequence, keys and values from another

\mathrm{Attn}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^{\top}}{\sqrt{d_k}} \right) V

The formula is identical to self-attention. The only thing that changes is where the three matrices come from.

Q = X_{\mathrm{dec}} W_Q, \qquad K = X_{\mathrm{enc}} W_K, \qquad V = X_{\mathrm{enc}} W_V

$X_{\mathrm{dec}}$ : the sequence doing the asking.
$X_{\mathrm{enc}}$ : the sequence being read.

In self-attention all three come from the same tensor, so a sequence relates to itself. Here the query sequence interrogates a different one.

That single substitution is what connects two modalities or two languages, and it is the mechanism behind every encoder-decoder model.

The shapes, which is where the asymmetry shows

Q \in \mathbb{R}^{T_{\mathrm{dec}} \times d_k}, \qquad K, V \in \mathbb{R}^{T_{\mathrm{enc}} \times d_k}

QK^{\top} \in \mathbb{R}^{T_{\mathrm{dec}} \times T_{\mathrm{enc}}}

The attention matrix is rectangular, one row per output position and one column per input position.

The two sequences need not be the same length, the same modality, or the same vocabulary. They only have to share a model width.

O(T_{\mathrm{dec}} \, T_{\mathrm{enc}} \, d)

The cost is the product of the two lengths rather than the square of one, which matters when a short query reads a long document.

One decoder block, in order

Self-attention over the tokens generated so far, causally masked.
Add and normalise.
Cross-attention, with queries from that result and keys and values from the encoder output.
Add and normalise.
Feed-forward, then add and normalise again.

The encoder output is computed once for the whole generation, so its keys and values are cached and reused at every decoding step.

That cache is a different object from the KV cache of self-attention, and it never grows, because the source sequence does not change while the output is produced.

No mask, and why that is correct

Self-attention in a decoder is masked, because a token must not see its own future.

Cross-attention is not, because the entire source is available before generation begins and there is nothing to hide.

\alpha_{ij}

Those weights are also the closest thing a transformer has to an alignment.

In translation they line up source and target words, and they are readable in a way self-attention weights usually are not.

Where it appears now

Encoder-decoder models such as T5 and the original transformer use it in every decoder block, which is its classic home.

Multimodal models use it to let text queries read image patches, which is how a vision-language model conditions on what it was shown.

Text-conditioned diffusion is the same trick in reverse, with the image as the query sequence and the prompt as the source, which is exactly where the prompt enters the denoising network.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

What distinguishes cross-attention from self-attention?

Transformers

19 concepts in this module, then a graded test with every answer explained.

Study this module →