← All tutorials

Cross-Attention

On this page

Cross-attention lets one sequence read another: each token of the first gets a new vector computed from the second sequence's content, so a translation can read its source and a fixed set of learned vectors can read an image. It runs on every forward pass in each layer that has it; the other sequence's keys and values can be computed once and reused while a decoder generates. As in Attention, each token (a word piece or an image patch, as a vector) makes a query, a key and a value, the rows of Q, K and V.

1 Queries come from one sequence, keys and values from another

In self-attention the queries, keys and values all come from one sequence; in cross-attention the queries come from stream X_A and the keys and values from stream X_B.

Each token is a vector of d numbers; here the queries and keys have d numbers too, so d_k = d.

In this cross-attention, 4 French tokens make the queries and 5 English tokens the keys and values, and one row per French token comes out.

import torch

def attention(Q, K, V):
    S = Q @ K.transpose(-2, -1) / Q.size(-1) ** 0.5
    return S.softmax(dim=-1) @ V

d = 64
X_A = torch.randn(1, 4, d)  # stream A, 4 French tokens
X_B = torch.randn(1, 5, d)  # stream B, 5 English tokens
W_Q, W_K, W_V = torch.randn(3, d, d) / d ** 0.5

Q = X_A @ W_Q               # queries from stream A
K, V = X_B @ W_K, X_B @ W_V  # keys, values from B
out = attention(Q, K, V)
print(out.shape)  # torch.Size([1, 4, 64])

1. 1Plain attention, the recap card above Section 1, line 2.
2. 2Streams of 4 and 5 tokens and three weight matrices, the n = 4 and m = 5 rows of Figure 1.
3. 3Queries from stream A, the cross card in Section 1, line 1.
4. 4Keys and values from stream B, the cross card, line 2.
5. 5One row per stream A token, the cross card, line 3.

If some rows of stream B are padding, the key padding mask from Positional Encoding and Masking still applies, on B's keys.

2 In translation, the decoder reads the encoder

The original Transformer translated with an encoder that reads the whole source sentence and a decoder that writes the translation one token at a time; cross-attention is how the decoder looks at the source, in the words of Vaswani et al.:

In "encoder-decoder attention" layers, the queries come from the previous decoder layer, and the memory keys and values come from the output of the encoder. This allows every position in the decoder to attend over all positions in the input sequence.

• Each decoder layer has three parts in order: masked self-attention over the translation so far, cross-attention to the encoder, then a feed-forward network (FFN) on each token.
• Cross-attention needs no causal mask (the mask that hides future tokens), because the whole source sentence is already known.
• The encoder's keys and values do not change while the decoder writes, so they are computed once and reused, the caching idea of Transformer Cost and Efficiency.
• Attention-based translators read the source in whatever order the target needs: in Bahdanau et al. (2015), European Economic Area became zone économique européenne, each French word attending to the English words in reverse order.
• When keys have the same length, a query scores highest on the key that points most nearly its way, the dot product drawn in Attention, Figure 2; a longer query multiplies all of its scores by the same factor, so its weights get sharper, which means more of the weight goes to one word.

In Figure 3 a toy decoder translates the black cat is sleeping into le chat noir dort. Its keys and queries have 2 numbers so they can be drawn, and its values have 4; the code example uses d = 64.

With noir most of the weight goes to black, which comes before cat in English, so the decoder reads the source out of order, as the Bahdanau bullet describes.

3 Learned queries choose the output size

Now make the queries learned parameters, a fixed table of N vectors that is the same for every input, so the output has N rows however many tokens the input has. The Perceiver (Jaegle et al. 2021) reads huge inputs this way: N learned latents cross-attend to M input elements.

With learned queries, any number of image feature rows go in and 32 rows per image come out.

import torch
import torch.nn as nn
import torch.nn.functional as F

N_QUERIES, WIDTH, INIT_STD = 32, 768, 0.02
start = torch.randn(N_QUERIES, WIDTH) * INIT_STD
queries = nn.Parameter(start)

B, M = 8, 257                  # images, features per image
img = torch.randn(B, M, WIDTH)
Q = queries.expand(B, -1, -1)  # same 32 queries, no copy
out = F.scaled_dot_product_attention(Q, img, img)
print(out.shape)  # torch.Size([8, 32, 768]) for any M

1. 132 learned queries, trained like weights, the latents of Figure 4.
2. 28 images, M = 257 features each, the M of the cost card in Section 3, line 2.
3. 3The same queries read every image, 32 rows for any M, the cost card in Section 3, line 2.

4 Cross-attention is the same PyTorch module fed two streams

Here a decoder reads a wider encoder; 5 decoder rows go in and 5 rows of width 512 come out.

import torch
import torch.nn as nn

B, n, m = 2, 5, 9  # batch, decoder rows, encoder rows
D_DEC, D_ENC, HEADS = 512, 1024, 8
cross = nn.MultiheadAttention(D_DEC, HEADS, kdim=D_ENC, vdim=D_ENC, batch_first=True)

dec = torch.randn(B, n, D_DEC)     # stream A, queries
memory = torch.randn(B, m, D_ENC)  # stream B, keys/values
mem_pad = torch.zeros(B, m, dtype=torch.bool)
mem_pad[1, 7:] = True  # last 2 rows of item 2 are padding

out, _ = cross(dec, memory, memory, key_padding_mask=mem_pad, need_weights=False)
print(out.shape)  # torch.Size([2, 5, 512])

1. 1A module with 1024-wide keys and values, the W_K and W_V of the cross card in Section 1, line 2.
2. 2Decoder rows (stream A), encoder rows (stream B) and a mask for 2 padded ones, the two streams of Figure 1.
3. 3One call; the output keeps the decoder's 5 rows and width 512, the cross card in Section 1, line 3.

• Mask only the memory's padding; the whole source is known, so no causal mask.
• A kdim unlike embed_dim turns off the fused fast path that nn.MultiheadAttention uses at inference.

From BLIP to the Q-Former uses 32 learned queries to let a frozen language model read a frozen image model.

5 References

1. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762.
2. Bahdanau, D., et al. (2015). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015. arXiv:1409.0473.
3. Jaegle, A., et al. (2021). Perceiver: General Perception with Iterative Attention. ICML 2021. arXiv:2103.03206.
4. Li, J., et al. (2023). BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models. ICML 2023. arXiv:2301.12597.