Cross-Attention

Cross-attention is attention between two different sets of tokens: one set supplies the queries, the other supplies the keys and values. Each query reads what it needs from the other set, and the result has one output per query. It is how a transformer decoder reads its encoder, how a detector’s learned object slots read an image, how a text-to-image generator’s image features read the prompt, and how a vision-language model brings image features into a language model. It is also called encoder–decoder attention, the name Vaswani et al. (2017) used. Scoring functions, scaling and multi-head attention are explained in the attention mechanism and self-attention articles.

Definition

Take two sets: $N_q$ query tokens $X_q$ and $N_k$ context tokens $X_k$. Cross-attention projects the queries from $X_q$ and the keys and values from $X_k$, then runs ordinary attention: each query scores every key, the scores are turned into weights with a softmax over the keys, and the query’s output is the weighted average of the values. The two sets can differ in size, in feature dimension and in modality. What makes it cross is only where the three inputs come from; the computation is the one in self-attention, where all three come from the same set.

Intuition

Think of a set of questions and a library. Each question (query) is compared with the catalogue entry (key) of every book, and the answer it takes away is a blend of the books’ contents (values), weighted by how well each entry matched. Asking more questions gives more answers; adding books makes each question take longer but does not change how many answers come back.

The query set decides the shape of the output; the context set decides what information is available. A diffusion model’s image features ask which words describe them and get one answer per image position; DETR’s 100 object queries get 100 answers however large the image is.

Formal Definition

With projection matrices $W_Q \in \mathbb{R}^{d_q \times d}$ and $W_K, W_V \in \mathbb{R}^{d_k \times d}$, and scaled dot-product scores as in the transformer,

$$Q = X_q W_Q \in \mathbb{R}^{N_q \times d}, \qquad K = X_k W_K, \;\; V = X_k W_V \in \mathbb{R}^{N_k \times d},$$ $$A = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d}}\right) \in \mathbb{R}^{N_q \times N_k}, \qquad \operatorname{CrossAttn}(X_q, X_k) = A V \in \mathbb{R}^{N_q \times d},$$

with the softmax taken along each row. Row $i$ of $A$ is query $i$’s attention map over the context tokens. Multi-head cross-attention splits $d$ into heads exactly as in self-attention. In a transformer block the output is added back to $X_q$ through a residual connection; $X_k$ is not changed by the layer.

Cost. The scores and the weighted sum each take $N_q N_k d$ multiply-adds, and the map $A$ holds $N_q N_k$ numbers per head. The projections add $N_q d_q d$ and $2 N_k d_k d$. The cost is therefore linear in each set’s size, and quadratic in neither. Carion et al. give the same count for DETR’s decoder: cross-attention between $N$ object queries and an $H \times W$ feature map costs $O(d^2(N + HW) + dNHW)$, much less than the encoder’s self-attention, $O(d^2 HW + d(HW)^2)$, because $N \ll HW$ (arXiv v3, Appendix A.3).

Masks. An additive mask on the scores, $-\infty$ for forbidden pairs, restricts which context tokens each query may read. Flamingo uses one to let each text token see only the image that preceded it (Alayrac et al., 2022, §2.3), and Mask2Former uses one to confine each query to its predicted mask (see below).

Properties

• Output length follows the queries. A fixed set of learned queries therefore turns an input of any size into a fixed number of tokens, which Perceiver, Flamingo and BLIP-2 use as a bottleneck.
• Information flows one way. Only the query tokens are updated. For the context set to read the query set, a second layer with the roles swapped is needed, as in dual-stream vision-language models and in SAM’s mask decoder, which runs cross-attention in both directions (see multimodal fusion).
• Order of the context does not matter. Permuting the key–value pairs together leaves every output unchanged, so where a context token sits is known only if its key or value carries position, typically through a positional encoding. DETR adds fixed spatial encodings to the image features and its learned object queries to the decoder’s inputs at every attention layer (arXiv v3, §3.2 and Appendix A.3).
• Different spaces, one inner product. The projections map both sets into a shared $d$-dimensional space, so text embeddings can be compared with U-Net features of a different width.

Cross-attention against concatenation. An alternative is to concatenate the two sets and run self-attention over all $N_q + N_k$ tokens. That lets both sets update each other and lets the context tokens attend among themselves, but it costs $(N_q + N_k)^2$ scores per head instead of $N_q N_k$, and it returns $N_q + N_k$ outputs, so the size of the output again depends on the input. Hybrids sit between the two. Imagen, as Hertz et al. (2023) describe it, appends the text tokens to the keys and values of image self-attention, so image tokens attend to image and text in one softmax while the output stays image-sized. Flamingo’s resampler likewise appends its latent queries to the visual features that form its keys and values (arXiv v2, Appendix A.1.1).

Examples

Shape templates attending over a rendered scene. The code draws a scene with one square, disk, triangle and cross, each $12 \times 12$ pixels at a random place, with Gaussian noise. Every $12 \times 12$ window of the image, at every pixel offset, is one context token: its key is the window’s pixels, made zero-mean and unit-norm, and its value is the pixels together with the window’s centre coordinates. The four queries are clean renderings of the four shapes, normalized the same way, so a score is a correlation coefficient divided by a temperature $\tau$. Nothing is learned: the projections are the identity, and the temperature plays the role of the $1/\sqrt{d}$ scaling. Because the values contain the window centres, the last two numbers of each output row are an attention-weighted position, which the code compares with the true centre of that query’s shape. Part 1 grows the image from $36 \times 36$ to $288 \times 288$ pixels; parts 2 and 3 compare two temperatures (NumPy 2.3.5, well under a second; two runs gave identical output).

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

P = 12                                   # patch size; patches are taken at every pixel
SHAPES = ["square", "disk", "triangle", "cross"]

def glyph(name):
    """A 12x12 binary drawing of one shape."""
    y, x = np.mgrid[0:P, 0:P] + 0.5
    if name == "square":   m = (abs(x - 6) < 4.5) & (abs(y - 6) < 4.5)
    if name == "disk":     m = (x - 6) ** 2 + (y - 6) ** 2 < 5.0 ** 2
    if name == "triangle": m = (y > 1.5) & (y < 10.5) & (abs(x - 6) < (y - 1.5) * 0.55)
    if name == "cross":    m = ((abs(x - 6) < 1.6) & (abs(y - 6) < 5.5)) | ((abs(y - 6) < 1.6) & (abs(x - 6) < 5.5))
    return m.astype(float)

def scene(rng, H):
    """An HxH noisy image with one of each shape at random, non-overlapping places."""
    img = 0.1 * rng.standard_normal((H, H))
    centres = {}
    for name in SHAPES:
        while True:
            r, c = rng.integers(0, H - P + 1, size=2)
            if all(abs(r + P / 2 - a) >= P or abs(c + P / 2 - b) >= P for a, b in centres.values()):
                break
        img[r:r + P, c:c + P] += glyph(name)
        centres[name] = (r + P / 2, c + P / 2)
    return img, centres

def normalize(X):
    """Zero-mean, unit-norm rows, so that a dot product is a correlation coefficient."""
    X = X - X.mean(axis=1, keepdims=True)
    return X / np.linalg.norm(X, axis=1, keepdims=True)

def keys_values(img):
    """One token per 12x12 window: key = normalized pixels, value = pixels and window centre."""
    win = sliding_window_view(img, (P, P))
    g = win.shape[0]
    feats = win.reshape(g * g, P * P)
    r, c = np.divmod(np.arange(g * g), g)
    centre = np.stack([r + P / 2, c + P / 2], axis=1)
    return normalize(feats), np.hstack([feats, centre]), centre

def cross_attention(Q, K, V, tau):
    S = Q @ K.T / tau                                  # (Nq, Nk) scores
    A = np.exp(S - S.max(axis=1, keepdims=True))
    A /= A.sum(axis=1, keepdims=True)                  # softmax over the keys
    return A @ V, A                                    # (Nq, dv) outputs, (Nq, Nk) maps

Q = normalize(np.array([glyph(s).ravel() for s in SHAPES]))   # 4 queries of d = 144

# 1. The output has one row per query, whatever the number of keys; cost grows with Nq*Nk.
rng = np.random.default_rng(0)
print("image     Nk  output    FLOPs (QK^T + AV)  FLOPs/(Nq*Nk)  self-attn/cross  max error (px)")
for H in [36, 72, 144, 288]:
    img, true = scene(rng, H)
    K, V, centre = keys_values(img)
    out, A = cross_attention(Q, K, V, tau=0.03)
    Nq, Nk, d, dv = Q.shape[0], K.shape[0], K.shape[1], V.shape[1]
    flops = 2 * Nq * Nk * d + 2 * Nq * Nk * dv        # one multiply-add = 2 FLOPs
    self_flops = 2 * (Nq + Nk) ** 2 * (d + dv)         # same layer over all Nq + Nk tokens
    err = max(np.hypot(*(out[i, -2:] - true[s])) for i, s in enumerate(SHAPES))
    print(f"{H:3d}^2 {Nk:6d}  {str(out.shape):8s} {flops:18,d} {flops // (Nq * Nk):14d} "
          f"{self_flops / flops:15.0f}x {err:15.2f}")

# 2. Where each query's attention goes: weight within 2 px of each shape's centre.
img, true = scene(np.random.default_rng(1), 72)
K, V, centre = keys_values(img)
for tau in [0.03, 0.1]:
    out, A = cross_attention(Q, K, V, tau)
    print(f"\ntau = {tau}: attention mass near each shape (rows: queries)")
    print("            " + "".join(f"{s:>10s}" for s in SHAPES) + "   elsewhere   error (px)")
    for i, s in enumerate(SHAPES):
        near = [A[i, np.hypot(*(centre - true[t]).T) <= 2].sum() for t in SHAPES]
        print(f"{s:>10s}  " + "".join(f"{m:10.2f}" for m in near)
              + f"{1 - sum(near):12.2f}{np.hypot(*(out[i, -2:] - true[s])):13.2f}")

# 3. Spread over 20 scenes at 72x72: localization error of the attention-weighted centre.
for tau in [0.03, 0.1]:
    errs = []
    for seed in range(20):
        img, true = scene(np.random.default_rng(100 + seed), 72)
        K, V, centre = keys_values(img)
        out, _ = cross_attention(Q, K, V, tau)
        errs += [np.hypot(*(out[i, -2:] - true[s])) for i, s in enumerate(SHAPES)]
    print(f"\n20 scenes, tau = {tau}: error median {np.median(errs):.2f} px, max {max(errs):.2f} px")

image     Nk  output    FLOPs (QK^T + AV)  FLOPs/(Nq*Nk)  self-attn/cross  max error (px)
 36^2    625  (4, 146)          1,450,000            580             158x            0.05
 72^2   3721  (4, 146)          8,632,720            580             932x            0.30
144^2  17689  (4, 146)         41,038,480            580            4424x            0.27
288^2  76729  (4, 146)        178,011,280            580           19184x            0.81

tau = 0.03: attention mass near each shape (rows: queries)
                square      disk  triangle     cross   elsewhere   error (px)
    square        1.00      0.00      0.00      0.00        0.00         0.06
      disk        0.00      1.00      0.00      0.00        0.00         0.08
  triangle        0.00      0.00      1.00      0.00        0.00         0.00
     cross        0.00      0.00      0.00      1.00        0.00         0.07

tau = 0.1: attention mass near each shape (rows: queries)
                square      disk  triangle     cross   elsewhere   error (px)
    square        0.53      0.18      0.07      0.05        0.18         8.15
      disk        0.14      0.52      0.03      0.13        0.18        10.71
  triangle        0.09      0.04      0.56      0.04        0.26         1.72
     cross        0.05      0.15      0.04      0.56        0.20        10.00

20 scenes, tau = 0.03: error median 0.06 px, max 0.21 px

20 scenes, tau = 0.1: error median 11.99 px, max 24.23 px

In the first table the number of context tokens grows 123-fold, but the output stays four rows, one per template. The FLOP count grows in step with $N_k$: divided by $N_q N_k$ it is the same 580 at every size, which is $2(d + d_v)$ with $d = 144$ and $d_v = 146$. Running the same layer as self-attention over the concatenated $N_q + N_k$ tokens would cost more by a factor that itself grows with the image, from about 160 to about 19,000, because $(N_q + N_k)^2 / (N_q N_k) \approx N_k / N_q$ when $N_k \gg N_q$.

At $\tau = 0.03$ each template’s attention map puts essentially all its weight within 2 pixels of its own shape, and the attention-weighted centre lands within a fraction of a pixel of the true one, across all 20 test scenes as well. In the first table, where each size is a different scene, the error is largest in the largest image; in our reading, the tiny weight that leaks to other shapes is multiplied by their distance, which grows with the image. At $\tau = 0.1$ the softmax is flatter: each query keeps only about half its weight on its own shape and spreads the rest over similar shapes and background. The disk template leaks to the square and the cross, which overlap a disk in many pixels, and the output, an average, lands between the shapes: over the 20 scenes the median error is 12 pixels. The output of soft attention is a blend, so a query that matches several keys returns something in between; in a trained model, the learned projections and scale set this sharpness, as $\tau$ does here.

Common Misconceptions

• “Cross-attention is a different kind of attention.” Libraries implement both with one function: PyTorch’s nn.MultiheadAttention takes the query, key and value as separate arguments.
• “The context set gets updated too.” In a text-to-image denoiser the text tokens leave every cross-attention layer unchanged; only the image features are rewritten.
• “Cross-attention is quadratic like self-attention.” Its cost $N_q N_k$ is large only when both sets are large, which is why models with many context tokens keep the query set small.
• “An attention map shows what the model used.” Cross-attention maps are informative: Prompt-to-Prompt edits images by manipulating them, and DETR builds masks from them. Whether a map explains a prediction is a separate question, discussed under interpretability of neural networks.

Where It Is Used

Encoder–decoder attention in the Transformer. In the original Transformer of Vaswani et al. (2017), each decoder layer has a cross-attention sub-layer whose queries come from the previous decoder layer and whose keys and values come from the encoder’s output, so every target position can attend over the whole source sentence (§3.2.3). The authors note that it mimics the encoder–decoder attention of earlier sequence-to-sequence models, citing among them Bahdanau, Cho and Bengio, whose attention was, in our reading, cross-attention in this sense. The encoder–decoder and transformer articles describe the architecture.

Object queries in detection and segmentation. DETR (Carion et al., 2020) feeds its decoder $N = 100$ learned embeddings, the object queries, which cross-attend to the encoded image feature map and are decoded in parallel into a class and a box each. The authors observed that decoder attention concentrates on object extremities such as heads and legs (§4.2, Fig. 6), and their panoptic segmentation head turns each object’s cross-attention scores over the encoded image into low-resolution heatmaps, the starting point of its mask (§4.4). Mask2Former (Cheng et al., 2022) replaces standard cross-attention in its decoder with masked attention: each query attends only to the foreground of the mask that the previous layer predicted for it, by adding $-\infty$ to the scores elsewhere (their Eqs. 2–3). The authors hypothesized that local features are enough to update a query and that context can come from self-attention among the queries. In their ResNet-50 ablation, masked attention gave 43.7 AP on COCO val2017 instance segmentation, 51.9 PQ on COCO panoptic val2017 and 47.2 mIoU on ADE20K val, against 37.8, 47.1 and 45.5 for standard cross-attention (arXiv v3, Table 4c). Instance segmentation and panoptic segmentation cover the model’s results.

Text conditioning in latent diffusion. Rombach et al. (2022, §3.3) condition the denoising U-Net of a latent diffusion model by adding cross-attention layers in which the flattened U-Net features are the queries and the output of a condition encoder, such as a text transformer, gives the keys and values. Each image position thus reads the prompt tokens relevant to it, and the same mechanism takes class labels or layouts instead of text. Stable Diffusion uses this design, and text-to-image generation surveys the alternatives. Because each cross-attention map ties image positions to individual words, it can be edited. Prompt-to-Prompt (Hertz et al., 2023) observed that the layout of a generated image depends on these maps, and edited images by re-using the maps from the original prompt while generating from an edited one: swapping a word, adding a phrase, or scaling one word’s attention. Their arXiv version used Imagen as the generator.

Latent bottlenecks: Perceiver. Self-attention scales quadratically with the number of inputs, which Jaegle et al. (2021) identify as the obstacle to applying it to the 50,176 pixels of a $224 \times 224$ image. Their Perceiver lets a small learned latent array (512 latents on ImageNet) cross-attend to the pixels, which costs $O(MN)$ for $M$ inputs and $N$ latents, then processes the latents with a deep self-attention stack costing $O(N^2)$ per layer, for $O(MN + LN^2)$ over $L$ layers (§3.1). The authors stress that the output of cross-attention takes the shape of the query input, so network depth is decoupled from input size; cross-attention can be repeated to re-read the input. With Fourier-feature positions and no convolutions, Perceiver reached 78.0% top-1 on ImageNet validation, against 77.6% for ResNet-50 and 77.9% for ViT-B/16 as reported in the literature (Table 1). They name the Set Transformer as the closest earlier work using cross-attention to project a large input onto a smaller array (§2).

Vision-language fusion. Vision-language models use cross-attention in two places. A resampler of learned queries cross-attends to the image features and hands a fixed number of tokens to the language model: 64 from Flamingo’s Perceiver Resampler (Alayrac et al., 2022, §2.1) and 32 from BLIP-2’s Q-Former, whose cross-attention layers sit in every other transformer block (Li et al., 2023, §3.1). Flamingo also inserts new cross-attention layers between the frozen layers of its language model, with text tokens as queries over the visual tokens, gated so that training starts from the unchanged language model (§2.2). The vision-language model article gives the equations, and multimodal fusion compares cross-attention with other ways of combining modalities.

Matching two images. In feature matching, SuperGlue (Sarlin et al., 2020) alternates self-attention among the keypoints of one image with cross-attention to the keypoints of the other, and LoFTR (Sun et al., 2021) does the same with dense features, so that each image’s descriptors are conditioned on both images.