PyTorch Cross-Attention in Depth

Cross-attention is attention where the queries come from one sequence and the keys and values come from another. It lets an output sequence read an input sequence while the output is being built.

Scores cross the two sequences

The target holds two positions and the source holds three. The score table therefore has two rows and three columns.

import torch
source = torch.tensor([[1.0, 0.0, 0.5, 0.0],
                       [0.0, 1.0, 0.0, 0.5],
                       [0.5, 0.0, 1.0, 0.0]])
target = torch.tensor([[1.0, 0.5, 0.0, 0.0],
                       [0.0, 0.0, 0.5, 1.0]])
scores = target @ source.T / (4 ** 0.5)
print("Scores shape:", tuple(scores.shape))

Softmax makes the weights

Softmax runs along each row, across the source positions. Every row then sums to 1.

weights = torch.softmax(scores, dim=1)
for row in weights.tolist():
    print([round(v, 3) for v in row])
print("Row sums:", [round(sum(r), 3) for r in weights.tolist()])

The output mixes the values

The weights multiply the source itself here, since this small example uses the source as its values. The output holds one row per target position.

output = weights @ source
print("Output shape:", tuple(output.shape))
print("First output row:", [round(v, 3) for v in output[0].tolist()])

Example

import torch
from torch import nn

torch.manual_seed(0)

dim = 8
embedding = nn.Embedding(6, dim)
to_q = nn.Linear(dim, dim)
to_k = nn.Linear(dim, dim)
to_v = nn.Linear(dim, dim)

source = embedding(torch.tensor([1, 2, 3]))
target = embedding(torch.tensor([4, 5]))
print("Source shape:", tuple(source.shape))
print("Target shape:", tuple(target.shape))

queries = to_q(target)
keys = to_k(source)
values = to_v(source)
scores = queries @ keys.T / (dim ** 0.5)
weights = torch.softmax(scores, dim=1)
print("Weights (target rows, source columns):")
for row in weights.tolist():
    print([round(v, 3) for v in row])
print("Each row sums to:", [round(sum(r), 3) for r in weights.tolist()])

output = weights @ values
print("Output shape:", tuple(output.shape))

Key points

Query: the target vector that asks what to read
Key: the source vector that a query is matched against
Value: the source vector that gets mixed into the output
Use scores = queries @ keys.T to match each target position against every source position
Cross-attention lets each target position take a weighted mix of the source values

In this example, a source of three words and a target of two words are embedded at width 8. Three linear layers turn the embeddings into queries, keys and values. The queries come from the target, and the keys and values come from the source. The scores divide by the square root of 8, which keeps the softmax steady. The weights print as two rows, one per target position. Each row sums to 1, and the output mixes the values by those weights. The output holds one row per target position.

Note

Self-attention is cross-attention where the two sequences are the same one

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn

torch.manual_seed(0)

dim = 8
embedding = nn.Embedding(6, dim)
to_q = nn.Linear(dim, dim)
to_k = nn.Linear(dim, dim)
to_v = nn.Linear(dim, dim)

source = embedding(torch.tensor([1, 2, 3]))
target = embedding(torch.tensor([4, 5]))
print("Source shape:", tuple(source.shape))
print("Target shape:", tuple(target.shape))

queries = to_q(target)
keys = to_k(source)
values = to_v(source)
scores = queries @ keys.T / (dim ** 0.5)
weights = torch.softmax(scores, dim=1)
print("Weights (target rows, source columns):")
for row in weights.tolist():
    print([round(v, 3) for v in row])
print("Each row sums to:", [round(sum(r), 3) for r in weights.tolist()])

output = weights @ values
print("Output shape:", tuple(output.shape))

Exercise

In cross-attention, where do the queries come from?

The target sequence The source sequence The loss function The padding mask