Cross-Attention in Multimodal AI: When One Stream Needs to Read Another

Cross-attention is not limited to the original encoder–decoder Transformer. It is a general mechanism for connecting two different streams of tokens or features. One stream asks a question through its queries; another stream provides the keys and values from which relevant information is retrieved. That simple pattern explains its use in translation, vision-language systems, retrieval-augmented generation, and text-to-image diffusion models.

The practical rule

A useful rule is: use cross-attention when a representation in stream A should selectively read information from a separate stream B. Stream A produces Query (Q); stream B produces Key (K) and Value (V). Self-attention, in contrast, uses a single stream for all three.

In the original Transformer, a decoder state queries the encoder’s output. This is why the mechanism is often called encoder–decoder attention. But the name describes one important application, not a requirement. The underlying operation applies whenever two token streams should be fused asymmetrically. Vaswani et al. specify this original case clearly: queries come from the previous decoder layer, while memory keys and values come from the encoder output.

Self-attention and cross-attention are one formula with different inputs

Scaled dot-product attention is written as:

Attention(Q, K, V) = softmax(QKᵀ / √dₖ) V

The formula does not change between self-attention and cross-attention. The source of its inputs does.

• Self-attention: Q, K, and V are all projected from the same tensor X.
• Cross-attention: Q is projected from X_A; K and V are projected from a different tensor X_B.

For example, a visual-question-answering model may let text tokens query image tokens. A text-to-image model reverses the direction: image-side features query text tokens. The direction is set by the task: the component that needs additional context produces Q; the component acting as reference produces K and V.

Stable Diffusion: image positions query text tokens

Stable Diffusion gives a concrete and intuitive use case. It includes a text encoder, a latent-space diffusion U-Net, and a VAE. A prompt such as a small red fox wearing a blue scarf is tokenized and encoded into a sequence of contextual text embeddings. The diffusion U-Net then denoises a noisy latent representation over many timesteps. It is inside the U-Net that cross-attention repeatedly makes the prompt available to image generation.

Latent Diffusion Models introduced cross-attention layers to make diffusion models flexible conditional generators for inputs such as text and bounding boxes. They perform the diffusion process in the latent space of a pretrained autoencoder rather than directly in pixel space. Rombach et al. describe this conditional design in the original LDM paper.

At one U-Net layer, suppose the latent-image feature map has shape [B, C_img, H, W]. It is flattened into N = H × W spatial tokens:

X_img ∈ ℝ^(B × N × C_img)

The prompt becomes a text-token matrix:

C_text ∈ ℝ^(B × M × C_text)

Here, M is the number of text tokens, including special or padding tokens. The layer learns separate projections:

Q = X_img  W_Q     # image features become queries
K = C_text W_K     # text features become keys
V = C_text W_V     # text features become values

After projection, each spatial image token compares itself with every text token. The resulting attention map has shape [B, heads, N, M]. Its value at a particular image location says how strongly that location reads each word token at that layer and denoising step. The weighted text information is then written back into the image feature stream through an output projection and residual connection.

This does not mean that an individual attention map is a complete causal explanation of an image. It does, however, capture the operational role of cross-attention: a region developing the fox can draw more information from the token for fox; a region developing the scarf can draw more information from scarf and blue. The Hugging Face diffusion course summarizes the intuition well: cross-attention layers are distributed through the U-Net, and each spatial location can attend to different tokens in the text conditioning.

Why different modalities do not need the same width

Text features and image features frequently have different channel dimensions. This is expected, not an error. Let:

X_img  ∈ ℝ^(B × N × C_img)
X_cond ∈ ℝ^(B × M × C_cond)

Cross-attention uses independently learned projections:

W_Q ∈ ℝ^(C_img  × h·d_h)
W_K ∈ ℝ^(C_cond × h·d_h)
W_V ∈ ℝ^(C_cond × h·d_v)

After splitting into h heads, Q and K both have final per-head width d_h, so their dot product is valid. V may theoretically use another width d_v, although standard implementations often take d_v = d_h for simplicity and efficiency. A final output projection maps the concatenated result back to C_img, allowing it to be added to the image feature stream.

This is the key engineering point: the input widths of different modalities may differ; only the projected Q and K widths must agree per attention head. Modern implementations may first map the conditioning encoder output to a shared cross_attention_dim. The Diffusers conditional U-Net documentation describes this option for projecting encoder hidden states before cross-attention.

Minimal PyTorch shape sketch

q = q_proj(image_tokens)   # [B, N, image_dim] → [B, heads, N, d_head]
k = k_proj(text_tokens)    # [B, M, text_dim]  → [B, heads, M, d_head]
v = v_proj(text_tokens)    # [B, M, text_dim]  → [B, heads, M, d_head]

scores  = q @ k.transpose(-2, -1) / sqrt(d_head)  # [B, heads, N, M]
weights = scores.softmax(dim=-1)
context = weights @ v                               # [B, heads, N, d_head]

# Concatenate heads and project back to the image-side width.
image_update = out_proj(context)  # [B, N, image_dim]

The code above highlights the only essential difference from self-attention: K and V originate in text_tokens, not in image_tokens. A practical implementation also handles head reshaping, layer normalization, residual paths, dropout, and padding masks for invalid condition tokens.

When should you expect cross-attention?

Expect it when a model maintains two streams that should stay distinct while still exchanging information selectively. Typical examples include text-conditioned image, audio, or video generation; image captioning; visual question answering; multimodal encoders; and models that read retrieved documents. However, not every multimodal model uses an explicit cross-attention module. Some project image tokens into an LLM’s embedding space, concatenate them with text tokens, and use ordinary causal self-attention over the combined sequence. Spatially aligned controls such as depth maps, edge maps, or segmentation masks may also be injected through concatenation, additive modulation, or a dedicated condition branch rather than cross-attention.

Takeaway

Cross-attention should be remembered as selective, directional reading across streams, rather than as a component tied to a particular architecture label. In Stable Diffusion, image-space latent features ask questions and text-token embeddings supply answers. Separate learned projections make that possible even though language and vision begin with different dimensions. This same pattern is one of the central building blocks of modern multimodal AI.

References

1. Vaswani et al., “Attention Is All You Need”.
2. Rombach et al., “High-Resolution Image Synthesis with Latent Diffusion Models”.
3. Hugging Face Diffusion Course: Stable Diffusion.
4. Hugging Face Diffusers: UNet2DConditionModel.