Cross-Attention

Cross-attention operates on two sequences: a query sequence $X_Q \in \mathbb{R}^{T_Q \times d_{model}}$ and a key-value sequence $X_{KV} \in \mathbb{R}^{T_{KV} \times d_{model}}$. Three learned linear projections map these into the attention's internal space:

Explanation: The query matrix Q comes from one source sequence, while K and V both come from a separate source — this asymmetry is the entire definition of cross-attention versus self-attention. Here $W^Q, W^K \in \mathbb{R}^{d_{model} \times d_k}$ and $W^V \in \mathbb{R}^{d_{model} \times d_v}$ are learned weight matrices. The raw attention logits (affinities) are computed as a scaled dot product:

Explanation: Each of the $T_Q$ query vectors is compared against each of the $T_{KV}$ key vectors, producing a full affinity matrix whose $(i,j)$ entry scores how relevant the $j$-th key-value position is to the $i$-th query. The $\frac{1}{\sqrt{d_k}}$ scaling prevents dot products from growing large in magnitude (which would push softmax into saturation regions with near-zero gradients). Softmax normalizes each row independently:

Explanation: Each query position now holds a valid probability distribution over all key-value positions, expressing how much attention that query pays to each part of the conditioning sequence. Finally, the output is a weighted mixture of value vectors:

Explanation: Every output token is a convex combination of all value vectors, with weights given by the corresponding row of the attention matrix — the output length equals $T_Q$, not $T_{KV}$. For multi-head cross-attention, $h$ independent attention heads run in parallel, each with its own projections $W^Q_i, W^K_i, W^V_i$, and their outputs are concatenated and projected:

Explanation: Multiple heads allow the model to simultaneously route different types of information from the conditioning sequence — for example, in translation one head might track syntactic correspondence while another tracks semantic meaning. A minimal PyTorch implementation clarifies the mechanics: ```python import torch import torch.nn as nn import torch.nn.functional as F class CrossAttention(nn.Module): def __init__(self, d_model: int, d_k: int, d_v: int, num_heads: int): super().__init__() self.num_heads = num_heads self.d_k = d_k

# Separate projections: Q from query seq, K/V from kv seq self.W_Q = nn.Linear(d_model, d_k num_heads, bias=False) self.W_K = nn.Linear(d_model, d_k num_heads, bias=False) self.W_V = nn.Linear(d_model, d_v num_heads, bias=False) self.W_O = nn.Linear(d_v num_heads, d_model, bias=False) def forward( self, x_q: torch.Tensor, # (B, T_Q, d_model) — query source x_kv: torch.Tensor, # (B, T_KV, d_model) — key/value source mask: torch.Tensor | None = None, ) -> torch.Tensor: B, T_Q, _ = x_q.shape T_KV = x_kv.shape[1] h = self.num_heads

# Project and split into heads Q = self.W_Q(x_q).view(B, T_Q, h, self.d_k).transpose(1, 2) # (B, h, T_Q, d_k) K = self.W_K(x_kv).view(B, T_KV, h, self.d_k).transpose(1, 2) # (B, h, T_KV, d_k) V = self.W_V(x_kv).view(B, T_KV, h, -1).transpose(1, 2) # (B, h, T_KV, d_v)

# Scaled dot-product attention scale = self.d_k ** -0.5 logits = torch.matmul(Q, K.transpose(-2, -1)) * scale # (B, h, T_Q, T_KV) if mask is not None: logits = logits.masked_fill(mask == 0, float('-inf')) attn = F.softmax(logits, dim=-1) # (B, h, T_Q, T_KV) out = torch.matmul(attn, V) # (B, h, T_Q, d_v)

# Merge heads and project out = out.transpose(1, 2).contiguous().view(B, T_Q, -1) return self.W_O(out) # (B, T_Q, d_model) ``` Note how `x_q` and `x_kv` are kept strictly separate — they can differ in length, token type, and even modality (text vs. image patches).

Self-Attention · Scaled Dot-Product Attention · Positional Encoding · Diffusion Models · Vision Transformers