Cross-Attention

Cross-attention is attention in which the queries come from one sequence and the keys and values come from another. Self-attention lets the positions of a sequence read each other; cross-attention lets the positions of one sequence read a different sequence. It is the standard way to condition one stream on another: a decoder reading an encoder’s output in translation, and a diffusion denoiser reading a text prompt when it generates an image.

The computation

Let \(X\in\mathbb{R}^{N_q\times d_x}\) be the sequence that is being updated (for a diffusion model, one row per spatial position of the latent image) and \(C\in\mathbb{R}^{N_c\times d_c}\) the sequence that is read (one row per text token). Three learned projections give

\[ Q = XW_Q,\qquad K = CW_K,\qquad V = CW_V,\qquad \text{out} = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right)V , \]

with \(Q\in\mathbb{R}^{N_q\times d}\), \(K,V\in\mathbb{R}^{N_c\times d}\) and \(d\) the head dimension. The score matrix \(QK^\top\) has one row per query position and one column per token. The softmax runs along each row, so every position receives a convex combination of the token values, and its weights say which words matter at that position. The input and output lengths match the updated sequence: \(N_q\) rows in, \(N_q\) rows out. The conditioning length \(N_c\) disappears in the product. In self-attention the same formula applies with \(C=X\).

Shapes in a diffusion model

The program below uses a 32×32 latent grid (1,024 positions), a prompt padded to 77 tokens, and made-up feature widths of 320 for the image and 768 for the text. Weights are random, so it checks shapes and cost, not meaning.

import numpy as np
rng = np.random.default_rng(0)
H = W = 32; Nq = H * W; Nt = 77; d_img = 320; d_txt = 768; d = 64
x = rng.normal(size=(Nq, d_img)); c = rng.normal(size=(Nt, d_txt))
Wq = rng.normal(size=(d_img, d)) / np.sqrt(d_img)
Wk = rng.normal(size=(d_txt, d)) / np.sqrt(d_txt)
Wv = rng.normal(size=(d_txt, d)) / np.sqrt(d_txt)
Q, K, V = x @ Wq, c @ Wk, c @ Wv
S = Q @ K.T / np.sqrt(d)
A = np.exp(S - S.max(1, keepdims=True)); A /= A.sum(1, keepdims=True)
out = A @ V
print("Q", Q.shape, "K", K.shape, "V", V.shape, "A", A.shape, "out", out.shape)
print("row sums min/max:", A.sum(1).min().round(6), A.sum(1).max().round(6))
for res in (32, 64, 128):
    n = res * res
    print(f"latent {res}x{res}: self-attn pairs {n*n:>12,}  cross-attn pairs (77 tokens) {n*77:>10,}")

Q (1024, 64) K (77, 64) V (77, 64) A (1024, 77) out (1024, 64)
row sums min/max: 1.0 1.0
latent 32x32: self-attn pairs    1,048,576  cross-attn pairs (77 tokens)     78,848
latent 64x64: self-attn pairs   16,777,216  cross-attn pairs (77 tokens)    315,392
latent 128x128: self-attn pairs  268,435,456  cross-attn pairs (77 tokens) 1,261,568

Three consequences follow from the shapes.

• The attention matrix is 1,024×77, one distribution over the tokens for each position. Reshaping a column back to the 32×32 grid gives a map of where one word acts, which is why cross-attention maps are used to inspect and edit generations.
• The cost grows with \(N_qN_c\), not \(N_q^2\). With 77 tokens, text conditioning at a 32×32 latent needs about 13 times fewer score entries than self-attention over the image, and the gap grows with resolution (about 213 times at 128×128).
• Text enters only through \(K\) and \(V\). Every position reads the same token matrix, with different weights, so the text is a shared memory that the image consults, not a signal added to a fixed place.

Why it conditions well, and where it falls short

Because each position chooses its own weights, a prompt such as “a red cube and a blue sphere” can send the left part of the image toward “red cube” and the right toward “blue sphere”. Nothing forces this: which attribute attaches to which object is learned, and a prompt with several objects often swaps or blends attributes. Keys and values also depend on the text encoder, so a token that the encoder represents poorly is poorly available to every position.

The softmax requires each position to spend its full weight of 1 on some token, even where the prompt says nothing about that region, so models tend to park the weight on a few uninformative tokens such as the start token. Cross-attention weights therefore show where the model reads, not a proof that a word caused the pixels there.

Neighbors

Self-attention uses one sequence for queries, keys and values. Concatenating text tokens with image tokens and applying self-attention to the joint sequence has the same expressive power for reading text, but costs \((N_q+N_c)^2\) and lets text tokens be updated too, which is how some newer diffusion architectures work. Channel concatenation and feature modulation (a scale and shift computed from a pooled vector) condition on a single vector or an aligned map and cannot let different positions read different tokens.