Cross-Attention › The Cross-Attention Formulation

To understand cross-attention deeply, we need to think about what problem it solves. Imagine you're translating "The cat sat on the mat" into French. You've already generated "Le chat" (the cat), and now you need to produce the next word. Which part of the English sentence should you focus on? The answer is "sat," because that's the verb that follows the subject. The model doesn't know this in advance; it has to figure it out by examining the decoder's current state (which encodes the fact that you've generated a subject) and comparing it against every position in the encoder output to find the best match.

This is precisely what cross-attention computes: for each position in the target sequence, it determines which positions in the source sequence are most relevant, then gathers information from those positions. The mechanism needs to answer two questions simultaneously:

1. Where should I look? Each decoder position needs to identify which encoder positions contain relevant information.
2. What should I extract? Once the relevant positions are identified, the decoder needs to pull out the appropriate information.

The query-key-value framework elegantly separates these concerns. Queries encode what the decoder is looking for, keys encode what each encoder position offers, and values encode the actual information to transmit. The separation between keys and values is conceptually similar to the separation between an index and the data it points to: the key is a compact description used for matching, while the value is the full content that gets transmitted once a match is found.

Let's construct the cross-attention formula piece by piece, understanding why each component is necessary.

Step 1: Create queries from the decoder. Each decoder position needs to express what information it's seeking. We project the decoder representations into a "query space." Think of each query as a search query that the decoder sends out: "I'm looking for information about X." We compute this with a linear projection:

$$Q = X_{\text{dec}} W^Q$$

where $X_{\text{dec}} \in \mathbb{R}^{n_{\text{dec}} \times d}$ contains the decoder's current representations and $W^Q \in \mathbb{R}^{d \times d_k}$ is a learned projection matrix. The resulting $Q \in \mathbb{R}^{n_{\text{dec}} \times d_k}$ has one row per decoder position, each encoding "what am I looking for?" The projection $W^Q$ is learned: during training, the model learns to map decoder states into a query space where semantically or syntactically related queries and keys produce high dot products.

Step 2: Create keys from the encoder. Each encoder position needs to advertise what information it contains. We project the encoder output into a "key space" using a different projection:

$$K = X_{\text{enc}} W^K$$

where $X_{\text{enc}} \in \mathbb{R}^{n_{\text{enc}} \times d}$ is the encoder output and $W^K \in \mathbb{R}^{d \times d_k}$ is another learned matrix. The resulting $K \in \mathbb{R}^{n_{\text{enc}} \times d_k}$ has one row per encoder position, each encoding "here's what I offer." The key space is designed to match the query space: a query that's looking for a verb should produce a high dot product with keys that correspond to verb positions in the source. Critically, the projection $W^K$ is separate from $W^Q$: queries and keys live in the same $d_k$-dimensional space, but they are projected there by different matrices, giving the model flexibility to learn asymmetric matching.

Step 3: Create values from the encoder. When attention flows to an encoder position, we need to specify what information transfers. The value projection captures this:

$$V = X_{\text{enc}} W^V$$

where $W^V \in \mathbb{R}^{d \times d_v}$ projects into a "value space." The resulting $V \in \mathbb{R}^{n_{\text{enc}} \times d_v}$ contains the content that will be aggregated. The value space is separate from the key space: the key is what you use to find a match, and the value is what you receive once the match is made. This separation allows the model to learn different representations for retrieval (keys) and content transmission (values).

Step 4: Compute similarity scores. Now we need to measure how well each query matches each key. The dot product is ideal for this: when two vectors point in similar directions (when they are highly correlated in the embedding space), their dot product is large. We compute all pairwise scores with a single matrix multiplication:

$$S = Q K^T$$

The resulting score matrix $S \in \mathbb{R}^{n_{\text{dec}} \times n_{\text{enc}}}$ has entry $(i, j)$ equal to the dot product between decoder query $i$ and encoder key $j$. This is where the rectangular shape emerges: we're comparing $n_{\text{dec}}$ queries against $n_{\text{enc}}$ keys. Entry $(i,j)$ being large means decoder position $i$ is likely to attend strongly to encoder position $j$, because their query and key representations are well-aligned.

Step 5: Scale to prevent saturation. In high dimensions, dot products tend to have large magnitudes, which would push softmax into saturation. When a softmax input contains very large values, the output becomes nearly one-hot: almost all probability mass concentrates on the maximum entry, and gradients for non-maximum entries vanish. This makes learning very slow or unstable. Dividing by $\sqrt{d_k}$ normalizes the scores:

$$S_{\text{scaled}} = \frac{Q K^T}{\sqrt{d_k}}$$

Why $\sqrt{d_k}$ specifically? The reasoning comes from the expected magnitude of a dot product between two random unit vectors in $d_k$ dimensions. If the components of $Q$ and $K$ are independent with mean zero and unit variance, the dot product $Q_i \cdot K_j$ has variance equal to $d_k$, so its standard deviation is $\sqrt{d_k}$. Dividing by $\sqrt{d_k}$ normalizes this to unit variance, keeping the scores in a range where softmax gradients remain healthy throughout training. The original "Attention Is All You Need" paper noted this instability explicitly and proposed the scaling factor as the fix.

The histograms illustrate why scaling matters. With $d_k = 64$, raw dot products have a standard deviation around 8, producing scores that can easily reach $\pm 20$ or more. After scaling by $\sqrt{64} = 8$, the standard deviation drops to approximately 1, keeping scores in a range where softmax produces meaningful gradients. This is not a minor detail: without the scaling factor, training deep transformer models becomes significantly more difficult. The gradients from the softmax layer become nearly zero for most attention entries, meaning the model barely learns to adjust its attention patterns. The fix is elegantly simple, and it works.

Step 6: Convert scores to attention weights. We apply softmax row-wise to convert each row of scores into a probability distribution over encoder positions. Each row of the score matrix represents one decoder position asking questions of all encoder positions. The softmax converts these raw scores into a probability distribution, interpreting the scores as un-normalized log-probabilities:

$$A = \text{softmax}(S_{\text{scaled}})$$

Each row of $A$ sums to 1.0, representing how one decoder position distributes its attention across all encoder positions. An entry $A_{ij} = 0.7$ means decoder position $i$ is directing 70% of its "attention budget" to encoder position $j$. The softmax is applied row-wise, so each decoder position independently computes its own distribution over the encoder. This means different decoder positions can attend to completely different parts of the encoder simultaneously.

Step 7: Aggregate values. Finally, we use the attention weights to compute a weighted sum of encoder values for each decoder position. This is the step where information flows from encoder to decoder:

$$\text{Output} = A V$$

The output has shape $(n_{\text{dec}} \times d_v)$: one row per decoder position, each containing information gathered from the encoder. If decoder position $i$ attends strongly to encoder position $j$, then the value vector $V_j$ contributes heavily to the output for position $i$. The weighted average captures the decoder's selective reading of the encoder: the output is not a copy of any single encoder representation, but a blend computed according to the decoder's current needs.

Why does this formula make sense? Notice that the softmax attention weights $A$ are non-negative and sum to 1 along each row, so the output $AV$ is literally a convex combination of the encoder's value vectors. Position $i$ in the decoder output is a weighted mixture of all encoder values, with weights determined by how well each encoder key matched that decoder's query. The model learns to assign high weights to relevant encoder positions and low weights to irrelevant ones, entirely through gradient descent on the end task.

Combining all steps, we arrive at the cross-attention formula:

$$\text{CrossAttention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

where:

• $Q = X_{\text{dec}} W^Q \in \mathbb{R}^{n_{\text{dec}} \times d_k}$: queries from the decoder, encoding "what information am I looking for?"
• $K = X_{\text{enc}} W^K \in \mathbb{R}^{n_{\text{enc}} \times d_k}$: keys from the encoder, encoding "what information do I have to offer?"
• $V = X_{\text{enc}} W^V \in \mathbb{R}^{n_{\text{enc}} \times d_v}$: values from the encoder, encoding "what content should I contribute?"
• $Q K^T \in \mathbb{R}^{n_{\text{dec}} \times n_{\text{enc}}}$: raw similarity scores measuring alignment between decoder queries and encoder keys
• $\sqrt{d_k}$: scaling factor that maintains healthy gradients during training
• $\text{softmax}(\cdot)$: applied row-wise to produce probability distributions
• $d_k$: query/key dimension (must match for the dot product)
• $d_v$: value dimension (determines output size)

The attention weight matrix has shape $(n_{\text{dec}} \times n_{\text{enc}})$, fundamentally different from self-attention's square $(n \times n)$ matrix. This rectangular shape reflects the asymmetry of sequence-to-sequence tasks: we have $n_{\text{dec}}$ positions that need to attend to $n_{\text{enc}}$ positions, and these lengths are typically different. For translating a 10-word English sentence into a 12-word French sentence when 8 French words have been generated, the attention matrix is $8 \times 10$: eight rows for the eight generated target positions, ten columns for the ten source positions.

Let's make this concrete by tracing the shapes through a realistic example. We'll simulate translating a 6-word English sentence into French, where we've generated 4 words so far.