Cross-Attention & Conditioning

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Cross-Attention Formulation
5. Encoder-Decoder Cross-Attention
6. Text-to-Image Conditioning
7. Retrieval-Augmented Generation (RAG)
8. Cross-Attention vs Self-Attention
9. Computational Properties
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive cross-attention and distinguish it from self-attention.
2. Explain the asymmetry: queries from one sequence, keys/values from another.
3. Show how cross-attention enables text conditioning in diffusion models.
4. Analyze the computational complexity of cross-attention.
5. Explain why cross-attention KV-cache doesn't grow during generation.

Notation

• $\mathbf{X} \in \mathbb{R}^{T_x \times d}$ — primary sequence (e.g., decoder states)
• $\mathbf{C} \in \mathbb{R}^{T_c \times d}$ — context sequence (e.g., encoder output)
• $T_x$ — length of primary sequence
• $T_c$ — length of context sequence

Core Intuition

Self-attention asks: "what parts of my own sequence are relevant?" Cross-attention asks: "what parts of another sequence are relevant to me?" The query comes from one sequence, while keys and values come from a different sequence — enabling information flow between two modalities or representations.

Cross-Attention Formulation

Decomposition:

• Queries from primary: $\mathbf{Q} = \mathbf{X}\mathbf{W}^Q \in \mathbb{R}^{T_x \times d_k}$.
• Keys from context: $\mathbf{K} = \mathbf{C}\mathbf{W}^K \in \mathbb{R}^{T_c \times d_k}$.
• Values from context: $\mathbf{V} = \mathbf{C}\mathbf{W}^V \in \mathbb{R}^{T_c \times d_v}$.
• Attention matrix: $\mathbf{A} \in \mathbb{R}^{T_x \times T_c}$ (rectangular, not square).
• Output: $\mathbb{R}^{T_x \times d_v}$ (same length as primary sequence).

Key difference from self-attention: K, V come from a different sequence than Q. No causal masking (context is fully available).

Encoder-Decoder Cross-Attention

In T5/BART-style models:

• Encoder processes input (bidirectional self-attention) → $\mathbf{C}$.
• Decoder generates output one token at a time:
  • Causal self-attention over decoder tokens.
  • Cross-attention: decoder queries, encoder keys/values.

Layer structure of decoder block:

KV-cache for cross-attention: Since the encoder output $\mathbf{C}$ is fixed during generation, the K, V projections are computed once (during prefill) and reused for all decoder steps. Cross-attention KV-cache doesn't grow during generation.

Text-to-Image Conditioning

In Stable Diffusion / DALL-E, cross-attention injects text information into the image generation process:

• Context $\mathbf{C}$: Text encoder output (e.g., CLIP embeddings of the prompt), $T_c \approx 77$ tokens.
• Primary $\mathbf{X}$: Spatial features of the noisy image at each diffusion step.

Interpretation: Each spatial position in the image "looks at" the text and decides which words are relevant for generating that patch. The word "dog" will receive high attention from spatial positions where the dog should appear.

Classifier-free guidance works by comparing cross-attended vs unconditioned outputs:

where $s > 1$ amplifies the cross-attention signal.

Retrieval-Augmented Generation (RAG)

Cross-attention can integrate retrieved documents:

• Primary $\mathbf{X}$: Current generation state.
• Context $\mathbf{C}$: Concatenated retrieved passages.

RETRO architecture: Interleave self-attention layers with cross-attention layers that attend to retrieved neighbors:

Advantage: The model can access a vast knowledge base without storing everything in parameters.

Cross-Attention vs Self-Attention

Mathematical properties:

• Self-attention ($\mathbf{X} = \mathbf{C}$): $\mathbf{A} \in \mathbb{R}^{T \times T}$, symmetric input, causal masking possible.
• Cross-attention ($\mathbf{X} \neq \mathbf{C}$): $\mathbf{A} \in \mathbb{R}^{T_x \times T_c}$, asymmetric, no causal constraint on context.

Gradient flow:

• Self-attention: Gradients flow between all positions bidirectionally (or causally).
• Cross-attention: Gradients flow from decoder → encoder (during training), enabling end-to-end learning.

Permutation properties:

• Self-attention: equivariant to permutation of input.
• Cross-attention: equivariant to permutation of primary sequence, but NOT to permutation of context (unless we want it to be).

Computational Properties

Complexity: $O(T_x \cdot T_c \cdot d)$ — scales with product of sequence lengths.

Typically: $T_c$ is much smaller than $T_x$ (e.g., text prompt = 77 tokens, image = 4096 spatial positions). So cross-attention is cheap relative to self-attention on the primary sequence.

Memory: Attention matrix is $T_x \times T_c$. For $T_x = 4096, T_c = 77$: only 315K entries (vs 16.7M for self-attention).

Common Pitfalls

Pitfall 1. Applying causal masking to cross-attention. The context (encoder output) is fully available to all decoder positions — masking it defeats the purpose.

Pitfall 2. Forgetting that cross-attention K/V can be pre-computed. Since $\mathbf{C}$ doesn't change during generation, computing K/V for cross-attention at every step wastes compute.

Pitfall 3. Assuming cross-attention order doesn't matter. The order of self-attention → cross-attention → FFN in a decoder block matters: self-attention first allows the decoder to build context, then cross-attention injects external information.

Summary

• Cross-attention: Q from one sequence, K/V from another → information bridge between sequences.
• Encoder-decoder: enables sequence-to-sequence tasks (translation, summarization).
• Text-to-image: spatial features query text embeddings for conditioning.
• RAG: generation state queries retrieved passages for knowledge.
• KV-cache for cross-attention is fixed (doesn't grow during generation).
• Complexity: $O(T_x \cdot T_c \cdot d)$ — typically cheap since $T_c \ll T_x$.

Exercises

Exercise 1. For an encoder-decoder model with encoder length $T_c=512$ and decoder length $T_x=128$: compute the cross-attention matrix size and compare to decoder self-attention.

Exercise 2. Explain why cross-attention KV-cache doesn't grow during autoregressive generation, while self-attention KV-cache does.

Exercise 3. In Stable Diffusion with image resolution 64×64 and text length 77: compute the cross-attention matrix size and FLOPs per layer.

Exercise 4. Derive the gradient $\partial\mathcal{L}/\partial\mathbf{C}$ through the cross-attention operation.

Exercise 5. Design an architecture that uses cross-attention to combine visual and audio features for a multimodal model.