Medium Research

Cross Attention

Cross-attention lets one sequence read information from another. Each query compares itself with every key, then uses those scores to take a weighted average of the values. For example, a translation decoder can use its own queries to read an encoder's representation of the source sentence.

The calculation is the same as scaled dot-product attention. What makes it cross-attention is where the queries and key/value pairs come from.

Your task

Implement cross_attention(Q, K, V) for these floating-point arrays:

Q: shape (tgt_len, d_k), one row per query.
K: shape (src_len, d_k), one row per source key.
V: shape (src_len, d_k), one value vector per source key.

The lengths and d_k are positive. The two sequence lengths can differ; K and V must have the same number of rows. This exercise uses the same feature width for queries, keys and values.

Compute scores = Q @ K.T / sqrt(d_k), with shape (tgt_len, src_len). Apply softmax across each row, over the source positions, then multiply those weights by V. Return shape (tgt_len, d_k), one value mixture per query. Use a numerically stable softmax.

There is no mask, dropout, batch dimension or learned projection in this exercise. Every query can attend to every source position.

Example

A zero query gives every key the same score. With values [[2, 0], [0, 4]], its output is their average, [[1, 2]]. A nonzero query can give the two values different weights.

The comparison below shows allowed connections, not their weights. The cross-attention panel illustrates why the score matrix can be rectangular.

Loading visualization…

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Attention combines values using learned weights

Attention lets each output position combine information from source positions. It uses three sets of vectors:

A query for each output position.
A key for each source position, compared with the query to get a score.
A value for each source position, containing the information to combine.

Scaled dot-product attention has three steps: compute query-key scores, turn those scores into weights with softmax, then take a weighted sum of the values.

$$A=\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}+M\right),\qquad O=AV.$$

Here, d_k is the width of each query and key vector. The mask M is zero for allowed positions and negative infinity for blocked positions.

For a single head, Q has shape (query_length, key_width), K has shape (source_length, key_width), and V has shape (source_length, value_width). The weights A have shape (query_length, source_length). Softmax runs over source positions, so each query gets its own set of weights that sum to one. The output has shape (query_length, value_width).

For weights [0.25, 0.75] and scalar values [2, 6], the output is 0.25 * 2 + 0.75 * 6 = 5. Both values contribute; attention does not simply choose the value with the largest score.

Mask before softmax

A causal mask blocks future positions. A padding mask blocks positions added to make sequences the same length. Setting a blocked score to negative infinity makes its exponential zero, so it receives zero weight. Setting it to zero does not work: exp(0) = 1.

Each row needs at least one allowed, finite score. Applying ordinary softmax to an all-negative-infinity row produces undefined values. If you instead zero weights after softmax, the remaining weights no longer sum to one unless you normalise them again.

Why divide by the square root of the key width?

If query and key components are independent, with mean zero and variance one, the variance of their dot product is d_k. Dividing by sqrt(d_k) brings it back to one. This helps keep softmax from concentrating nearly all the weight on one source position, where its derivatives become small. Scaling helps under these assumptions; it does not prevent extreme weights for every possible input.

In self-attention, queries, keys, and values come from the same sequence. In cross-attention, queries come from one sequence and keys and values come from another.

Detecting runtime…

Cross Attention

Your task

Example

Attention combines values using learned weights

Mask before softmax

Why divide by the square root of the key width?

Your code, saved here

Stuck?

PyTorch reference solution