Cross-Attention and Position Encoding

The earlier pages of this chapter built self-attention. Every token of a sequence

X

(

T \times d

) is projected into a query, a key and a value,

Q = XW^q, \qquad K = XW^k, \qquad V = XW^v,

and each output row is a weighted sum of the values, with weights given by how well that row’s query matches every key:

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V.

Two questions were left open. The first is about order: nothing in this formula knows which token came first. The second is about sources:

Q

K

and

V

all come from the same

X

, but a translation model must relate an English sentence to a French one, and a vision language model must relate text to an image. This section answers both. It covers the causal mask that separates a decoder from an encoder, cross-attention between two sequences, and the ways of encoding position. The vision chapters rely on all three: a vision transformer is an encoder with learned position embeddings, and the connectors of vision language models use cross-attention.

import math

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

torch.manual_seed(0)


def attention(Q, K, V, mask=None):
    scores = Q @ K.transpose(-2, -1) / math.sqrt(Q.shape[-1])
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ V, weights

Attention ignores order

Shuffle the rows of

X

and the rows of

Q

K

and

V

are shuffled the same way. Each query still meets the same set of keys, so each output row is unchanged, and it lands in the shuffled position. Self-attention is permutation equivariant: it treats the input as a set, not a sequence. The cell checks this with random matrices.

T, d = 6, 8
X = torch.randn(T, d)
Wq, Wk, Wv = (torch.randn(d, d) for _ in range(3))

out, _ = attention(X @ Wq, X @ Wk, X @ Wv)

perm = torch.randperm(T)
out_shuffled, _ = attention(X[perm] @ Wq, X[perm] @ Wk, X[perm] @ Wv)

print("permutation:", perm.tolist())
print("outputs of the shuffled input equal the shuffled outputs:",
      torch.allclose(out_shuffled, out[perm], atol=1e-5))

permutation: [4, 1, 0, 2, 5, 3]
outputs of the shuffled input equal the shuffled outputs: True

“dog bites man” and “man bites dog” contain the same tokens, so without extra information a transformer would produce the same outputs for both, only reordered. The model needs the position of each token as part of its input. The last part of this section shows how it gets it.

Encoders and decoders

An encoder reads a whole input at once, and every token may attend to every other token, in both directions. BERT and the vision transformer are encoders. Their outputs are representations of the input, used for classification or as input to another model. A decoder generates a sequence one token at a time, and each new token is predicted from the tokens before it. During training the whole target sequence is available, and all positions are trained in parallel. Token

i

must then be prevented from seeing tokens

i+1, \dots, T

, otherwise it could copy the answer. The causal mask does this: before the softmax, the score of every key that comes after the query is set to

-\infty

, so its weight becomes zero.

M_{ij} = \begin{cases} 0 & j \le i \\ -\infty & j > i \end{cases}, \qquad \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V.

GPT and the language model inside every vision language model are decoders. The cell builds the mask and checks its effect: changing the last token changes the last output, and none of the outputs before it.

causal = torch.tril(torch.ones(T, T, dtype=torch.bool))   # True where attention is allowed
out, weights = attention(X @ Wq, X @ Wk, X @ Wv, mask=causal)

X2 = X.clone()
X2[-1] = torch.randn(d)                                    # change only the last token
out2, _ = attention(X2 @ Wq, X2 @ Wk, X2 @ Wv, mask=causal)

print("attention weights (row = query, column = key):")
print(weights.round(decimals=2))
print("outputs 1 to T-1 unchanged:", torch.allclose(out[:-1], out2[:-1]))
print("output T changed:", not torch.allclose(out[-1], out2[-1]))

attention weights (row = query, column = key):
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.9000, 0.1000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.4300, 0.0400, 0.5300, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.8900, 0.0300, 0.0700, 0.0100, 0.0000],
        [0.4200, 0.0000, 0.5800, 0.0000, 0.0000, 0.0000]])
outputs 1 to T-1 unchanged: True
output T changed: True

The weights above the diagonal are zero, and each row still sums to one over the tokens it may see.

Cross-attention

In cross-attention the queries come from one sequence and the keys and values from another. Let

Y

(

T_y \times d

) be the sequence that asks and

Z

(

T_z \times d

) the sequence it reads from:

Q = YW^q, \qquad K = ZW^k, \qquad V = ZW^v, \qquad \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V.

The weight matrix is

T_y \times T_z

: one row per query, one column per key. The output has one row per query,

T_y \times d

, whatever the length of

Z

. Each row is a mixture of the values of

Z

, chosen by the query of the corresponding row of

Y

. The original transformer used cross-attention to translate (Vaswani et al., 2017). The encoder reads the source sentence, and each block of the decoder has a self-attention layer with a causal mask over the target words written so far, followed by a cross-attention layer whose queries come from those target words and whose keys and values come from the encoder outputs. The decoder decides which source words to look at while it writes each target word. Vision language models use the same layer to connect an image to text. The image encoder produces a sequence of patch features

Z

, and a sequence

Y

reads from them. In Flamingo, the text tokens are

Y

, through cross-attention layers inserted into a frozen language model (Alayrac et al., 2022). In BLIP-2,

Y

is a small set of learned query vectors (Li et al., 2023): 32 queries read from the 257 features of the image encoder and return 32 outputs, a summary of fixed size that is then passed to the language model. The decoder of the Segment Anything Model uses cross-attention in both directions between prompt tokens and the image embedding. The cell builds the BLIP-2 pattern on a small scale: 4 learned queries attend to 64 image features and return 4 vectors.

D = 32
image_features = torch.randn(1, 64, D)                  # Z: 64 patch features from an image encoder
queries = nn.Parameter(torch.randn(1, 4, D))            # Y: 4 learned query vectors

cross_attn = nn.MultiheadAttention(D, num_heads=4, batch_first=True)
summary, weights = cross_attn(query=queries, key=image_features, value=image_features)

print("queries:", tuple(queries.shape), " image features:", tuple(image_features.shape))
print("output:", tuple(summary.shape), " attention weights:", tuple(weights.shape))

# The output size does not depend on the number of image features
summary_256, _ = cross_attn(queries, torch.randn(1, 256, D), torch.randn(1, 256, D))
print("output for 256 image features:", tuple(summary_256.shape))

queries: (1, 4, 32)  image features: (1, 64, 32)
output: (1, 4, 32)  attention weights: (1, 4, 64)
output for 256 image features: (1, 4, 32)

The four outputs keep their shape when the image produces 256 features instead of 64. The queries decide how many tokens leave the layer, and the image decides only what they contain. LLaVA connects the two models differently: a small MLP maps each image feature to a token of the language model, so the number of image tokens equals the number of patches, and the language model’s own self-attention relates them to the text. The LLaVA deep dive compares the two designs.

Encoding position

A transformer gets order back by adding information about position to its input or to its attention scores. Three methods are in common use.

Learned absolute embeddings

The simplest is a table of learned vectors, one per position, added to the token embeddings:

X + P

, with

P \in \mathbb{R}^{T_{\max} \times d}

. BERT, GPT-2 and the vision transformer use it. It is flexible, but the model has no vector for a position beyond

T_{\max}

, and nothing in it says that positions 5 and 6 are neighbors. That has to be learned.

Sinusoidal encodings

The original transformer used a fixed encoding instead (Vaswani et al., 2017). Position

t

gets a vector whose entries are sines and cosines of

t

at geometrically spaced frequencies:

p_{t, 2i} = \sin\!\left(\frac{t}{10000^{2i/d}}\right), \qquad p_{t, 2i+1} = \cos\!\left(\frac{t}{10000^{2i/d}}\right), \qquad i = 0, \dots, d/2 - 1.

The first dimensions oscillate quickly and tell neighboring positions apart. The last ones change slowly and tell the beginning of a long sequence from its end, like the hands of a clock. Each pair of dimensions is a point on a circle that turns by a fixed angle per step, so the encoding of position

t + k

is a fixed rotation of the encoding of position

t

. As a consequence, the dot product

\mathbf{p}_t \cdot \mathbf{p}_{t+k}

depends only on the offset

k

, not on

t

def sinusoidal(T, d):
    t = torch.arange(T, dtype=torch.float32)[:, None]
    freq = 1.0 / 10000 ** (torch.arange(0, d, 2, dtype=torch.float32) / d)
    p = torch.zeros(T, d)
    p[:, 0::2] = torch.sin(t * freq)
    p[:, 1::2] = torch.cos(t * freq)
    return p


P_sin = sinusoidal(128, 64)
for t in (10, 50, 90):
    print(f"p_{t} . p_{t + 5} = {P_sin[t] @ P_sin[t + 5]:.4f}")

p_10 . p_15 = 23.5040
p_50 . p_55 = 23.5040
p_90 . p_95 = 23.5040

The dot product with position 64 is largest at 64 and falls off on both sides, so nearby positions have similar encodings without any training.

Rotary position embeddings

Most current language models, including Llama, Qwen and SmolLM2, use rotary position embeddings (RoPE, arXiv:2104.09864). Instead of adding a vector to the input, RoPE rotates the query and the key inside every attention layer. The dimensions of a query

\mathbf{q}_m

at position

m

are grouped into pairs, and pair

i

is rotated by the angle

m\theta_i

, with the same frequencies

\theta_i = 10000^{-2i/d}

as the sinusoidal encoding. The key

\mathbf{k}_n

at position

n

is rotated by

n\theta_i

. A rotation by

m\theta

followed by the inverse of a rotation by

n\theta

is a rotation by

(n-m)\theta

, so

(R_m \mathbf{q})^\top (R_n \mathbf{k}) = \mathbf{q}^\top R_{n-m}\, \mathbf{k}.

The attention score depends on the content of the two tokens and on their offset

n - m

, and not on where the pair sits in the sequence. Position enters exactly where it is used, in the comparison of a query with a key. The cell checks the property: the same query and key give the same score at positions (3, 7) and (103, 107).

def rope(x, pos):
    # Rotate each consecutive pair of dimensions of x by pos * theta_i
    d = x.shape[-1]
    theta = 1.0 / 10000 ** (torch.arange(0, d, 2, dtype=torch.float32) / d)
    angle = pos * theta
    cos, sin = torch.cos(angle), torch.sin(angle)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    out = torch.empty_like(x)
    out[..., 0::2] = x1 * cos - x2 * sin
    out[..., 1::2] = x1 * sin + x2 * cos
    return out


q, k = torch.randn(64), torch.randn(64)
for m, n in [(3, 7), (103, 107), (3, 8)]:
    print(f"query at {m:>3}, key at {n:>3}: score {rope(q, m) @ rope(k, n):.4f}")

query at   3, key at   7: score 11.2033
query at 103, key at 107: score 11.2033
query at   3, key at   8: score 11.3851

The first two scores are equal because the offset is 4 in both cases. The third uses an offset of 5 and differs.

Positions in two dimensions

Image patches have a row and a column. The vision transformer numbers its patches row by row and learns one embedding per number, and the learned embeddings end up encoding the 2D layout. The vision transformer section shows this, and shows how the table is interpolated when the image size changes. Some vision language models apply RoPE to the row and the column separately. Qwen2-VL splits the rotary dimensions into three groups, for time, height and width, so that the same mechanism handles text, images and video (arXiv:2409.12191).

Exercises

Add the sinusoidal encodings to $X$ in the permutation test above. Is the self-attention output still permutation equivariant? Explain.
Implement a single decoder block for translation: causal self-attention over the target, then cross-attention to an encoder output, then an MLP, each with a residual connection and layer norm. Check the output shapes for a source of 10 tokens and a target of 7.
In the BLIP-2 example, compute the attention weights of each of the 4 queries over the 64 image features and check that every row sums to one. How many tokens would the language model receive from a LLaVA-style MLP connector for the same 64 features?
Show that the sinusoidal encoding of position $t + k$ is a linear function of the encoding of position $t$ : find the $2 \times 2$ matrix that maps the pair $(p_{t,2i}, p_{t,2i+1})$ to $(p_{t+k,2i}, p_{t+k,2i+1})$ .

Key references

Vaswani et al. introduced the encoder and decoder, cross-attention and the sinusoidal encoding (Vaswani et al., 2017). Flamingo (Alayrac et al., 2022) and BLIP-2 (Li et al., 2023) use cross-attention to connect an image encoder to a language model.

References

Alayrac, J., Donahue, J., Luc, P., Miech, A., Barr, I., et al. (2022). Flamingo: a Visual Language Model for Few-Shot Learning. arXiv [cs.CV].
Li, J., Li, D., Savarese, S., Hoi, S. (2023). BLIP-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. arXiv [cs.CV].
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., et al. (2017). Attention Is All You Need. arXiv [cs.CL].

Found an error on this page? Report an issue.

Large Language Models

NLP Foundations

Recurrent Neural Networks

Language Models

Neural Machine Translation

Transformers

LLM SFT Lab

Cross-Attention and Position Encoding

Attention ignores order

Encoders and decoders

Cross-attention

Encoding position

Learned absolute embeddings

Sinusoidal encodings

Rotary position embeddings

Positions in two dimensions

Exercises

Key references

References