Cross-attention: reading from another set

Queries come from one set; keys and values come from the other

In self-attention, one set of tokens makes all three vectors. In cross-attention, the roles split. For each token a of set A and each token b of set B, learned matrices make

q = W_Qa,   k = W_Kb,   v = W_Vb

Stack A's n queries as the rows of Q, and B's m keys and values as the rows of K and V. The formula is the same as before:

CrossAttention = softmax( QK^T / √d_k ) V

Only the shapes change, because n and m can differ:

Each row of the weights belongs to one A token and adds up to 1. The output has n rows: one new vector per A token. Set B is read, not changed. Self-attention is the special case where A and B are the same set.

Worked by hand: two tokens read from three

Set A has two tokens with queries q₁ = [1, 0] and q₂ = [0, 2]. Set B has three tokens with keys [2, 0], [0, 1] and [1, 1], and values [6, 0], [0, 6] and [3, 3]. The key size is d_k = 2.

Predict first. Which B token will A's first token read most? Which will the second read most?

The order of set B does not matter, and its size can change

Each A token takes a weighted sum over all of B, and a sum does not depend on the order of its terms. So reordering B's tokens leaves every output the same: cross-attention is permutation invariant in B. Reordering A's tokens reorders the outputs in the same way. Set B can also hold 3 tokens or 3,000; the output still has one vector per A token.

cross_attention.py

import numpy as np

def softmax(s):
    e = np.exp(s - s.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def cross_attention(Qa, Kb, Vb):
    w = softmax(Qa @ Kb.T / np.sqrt(Qa.shape[-1]))   # n rows, m columns
    return w @ Vb, w

Qa = np.array([[1.0, 0.0], [0.0, 2.0]])               # queries from set A
Kb = np.array([[2.0, 0.0], [0.0, 1.0], [1.0, 1.0]])   # keys from set B
Vb = np.array([[6.0, 0.0], [0.0, 6.0], [3.0, 3.0]])   # values from set B
out, w = cross_attention(Qa, Kb, Vb)
print(np.round(w, 3))
print(np.round(out, 2))

perm = [2, 0, 1]                                      # reorder set B
# np.allclose is True when two arrays match, up to tiny rounding errors
print(np.allclose(cross_attention(Qa, Kb[perm], Vb[perm])[0], out))

[[0.576 0.14  0.284]
 [0.108 0.446 0.446]]
[[4.31 1.69]
 [1.99 4.01]]
True

A model uses cross-attention to let each token choose what to read

A common block for set A runs three steps, each with a layer normalization and a residual connection: self-attention, so A's tokens can coordinate with each other; cross-attention, so each A token can read the condition B; and an MLP. The original transformer for translation worked this way: the words of the sentence being written read the words of the sentence being translated. The transformer block unit shows this pattern without the cross-attention step.

There are two other common ways to bring a condition in. A model can join A and B into one set and run self-attention over all n + m tokens. Or it can squeeze B into one summary vector and add it to every A token, which is cheap but gives every token the same information. Cross-attention lets each A token pick the parts of B it needs, at a cost of n × m scores.

Cross-attention can also shrink a large set. A model can learn a handful of query tokens that read a set of thousands of tokens, and so get a fixed-size summary of it. The Perceiver model works this way; PointNet pooling is a simpler order-free summary.