Cross-Attention (Encoder-Decoder)

Background

Self-attention queries within one sequence; cross-attention lets the decoder attend over the encoder's outputs. It's the transformer's encoder-decoder bridge — used in machine translation, image-conditioned text generation, and most multimodal architectures.

Problem statement

Implement cross_attention(Q, K, V) returning:

$$\text{Attention}(Q, K, V) \;=\; \mathrm{softmax}\!\Bigl(\frac{Q K^\top}{\sqrt{d_k}}\Bigr) V$$

where $Q$ comes from the decoder and $K, V$ from the encoder. Use numerically stable row-wise softmax.

Input

• Q — np.ndarray shape $(T_q, d_k)$, decoder queries.
• K — np.ndarray shape $(T_k, d_k)$, encoder keys.
• V — np.ndarray shape $(T_k, d_v)$, encoder values. Note: $K$ and $V$ share their first dimension.

Output

Returns np.ndarray shape $(T_q, d_v)$.

Examples

Example 1 — basic shape

Input:  Q shape=(3, 4), K shape=(5, 4), V shape=(5, 6)
Output: shape=(3, 6)

Example 2 — hard attention via large logits → output equals one row of V

Input:  Q designed to overwhelmingly favour K[0]
Output: ≈ V[0]

Example 3 — uniform Q over K → output is the mean of V

Input:  Q=0 (uniform softmax)
Output: row-wise mean of V

Constraints

• Q.shape[-1] == K.shape[-1] (shared $d_k$).
• K.shape[0] == V.shape[0] (same $T_k$).
• Softmax must subtract the row max before exp for stability.

Notes

• No masking here. Decoder self-attention typically masks future tokens; cross-attention does not (the decoder can attend over the full encoder output).
• Multi-head. Production transformers split $Q, K, V$ into $h$ heads, apply attention per head, concatenate. This problem isolates one head.