Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Cross-Attention › Multi-Head Cross-Attention

Michael BrenndoerferUpdated June 18, 20256 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 114Cross-Attention › Multi-Head Cross-Attention

Reading section 8 of 14 · Multi-Head Cross-Attention · 6 minSection 8 of 14 · 6 min

Section 8 of 14 · 6 min read

Cross-Attention › Multi-Head Cross-AttentionOne section · Sections in this chapter8 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Just like self-attention, cross-attention benefits enormously from multiple attention heads. Each head runs an independent cross-attention computation with its own set of QKV projection matrices, and the results from all heads are concatenated and projected back to the model dimension. Multiple heads let the model add capacity while specializing each head in different aspects of source-target alignment.

Think of multi-head attention as asking multiple questions simultaneously. One head might specialize in semantic alignment (which source word has the same meaning as the current target word?), another in syntactic alignment (which source word plays the same grammatical role?), and a third in positional proximity (which source word is near the corresponding position in the sequence?). Because each head has its own independent projection matrices, these specializations emerge naturally from training. The final concatenation and projection combine all these perspectives into a single rich representation.

The mathematical formulation of multi-head attention applies cross-attention $h$ times in parallel, then combines the results:

\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O

where each head computes:

\text{head}_i = \text{CrossAttention}(Q W^Q_i, K W^K_i, V W^V_i)

and:

$W^Q_i \in \mathbb{R}^{d \times d_k}$ : query projection for head $i$
$W^K_i \in \mathbb{R}^{d \times d_k}$ : key projection for head $i$
$W^V_i \in \mathbb{R}^{d \times d_v}$ : value projection for head $i$
$W^O \in \mathbb{R}^{h d_v \times d}$ : output projection combining all heads

In practice, we set $d_k = d_v = d / h$ , so each head operates on a $d/h$ -dimensional subspace, and concatenating $h$ heads of dimension $d/h$ gives back $d$ dimensions. This keeps the total computation comparable to single-head attention while providing the diversity of perspectives.

In[21]:

Code

class MultiHeadCrossAttention:
    """
    Multi-head cross-attention for encoder-decoder transformers.
    """

    def __init__(self, d_model, n_heads):
        """
        Initialize multi-head cross-attention.

        Args:
            d_model: Model dimension (must be divisible by n_heads)
            n_heads: Number of attention heads
        """
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.scale = 1.0 / np.sqrt(self.d_k)

        # Projections for all heads combined
        self.W_Q = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )
        self.W_K = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )
        self.W_V = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )
        self.W_O = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )

    def __call__(self, decoder_state, encoder_output, encoder_mask=None):
        """
        Apply multi-head cross-attention.

        Args:
            decoder_state: (n_dec, d_model)
            encoder_output: (n_enc, d_model)
            encoder_mask: Optional (n_enc,)

        Returns:
            output: (n_dec, d_model)
            attention_weights: (n_heads, n_dec, n_enc)
        """
        n_dec = decoder_state.shape[0]
        n_enc = encoder_output.shape[0]

        # Project and reshape to (n_heads, n, d_k)
        Q = (
            (decoder_state @ self.W_Q)
            .reshape(n_dec, self.n_heads, self.d_k)
            .transpose(1, 0, 2)
        )
        K = (
            (encoder_output @ self.W_K)
            .reshape(n_enc, self.n_heads, self.d_k)
            .transpose(1, 0, 2)
        )
        V = (
            (encoder_output @ self.W_V)
            .reshape(n_enc, self.n_heads, self.d_k)
            .transpose(1, 0, 2)
        )

        # Compute attention for all heads: (n_heads, n_dec, n_enc)
        scores = np.matmul(Q, K.transpose(0, 2, 1)) * self.scale

        # Apply mask if provided
        if encoder_mask is not None:
            mask = encoder_mask.reshape(1, 1, -1)
            scores = np.where(mask, scores, -1e9)

        attention_weights = softmax(scores)

        # Aggregate values: (n_heads, n_dec, d_k)
        head_outputs = np.matmul(attention_weights, V)

        # Concatenate heads and project: (n_dec, d_model)
        concatenated = head_outputs.transpose(1, 0, 2).reshape(n_dec, -1)
        output = concatenated @ self.W_O

        return output, attention_weights

The implementation projects all heads simultaneously using large combined projection matrices, then reshapes and transposes to separate the heads. This is more efficient than running $h$ separate projections, because a single large matrix multiplication is generally faster than many small ones on modern hardware.

In[22]:

Code

# Test multi-head cross-attention
n_heads = 4
d_model_mh = 32

# Create larger representations
encoder_output_mh = np.random.randn(n_source, d_model_mh)
decoder_state_mh = np.random.randn(n_target, d_model_mh)

mh_cross_attn = MultiHeadCrossAttention(d_model_mh, n_heads)
output_mh, weights_mh = mh_cross_attn(decoder_state_mh, encoder_output_mh)

Out[23]:

Console

Multi-head cross-attention with 4 heads:
  Input encoder shape:  (4, 32)
  Input decoder shape:  (2, 32)
  Output shape:         (2, 32)
  Attention weights:    (4, 2, 4) (heads × target × source)

The output maintains the same shape as the decoder input (2 tokens, 32 dimensions), but now each position has gathered information from the encoder through 4 independent attention computations. The attention weights tensor has shape (4, 2, 4), meaning each of the 4 heads produces its own 2×4 attention pattern. The output projection $W^O$ then combines all four heads' outputs into a single unified representation.

Out[24]:

Visualization

Heatmap showing attention weights for head 1. — Head 1 attention pattern. Each head independently attends to source positions, potentially specializing in different alignment patterns.

Heatmap showing attention weights for head 2. — Head 2 attention pattern. Variation across heads allows the model to capture multiple aspects of source-target alignment simultaneously.

Heatmap showing attention weights for head 3. — Head 3 attention pattern. In trained models, heads often develop interpretable specializations such as semantic or syntactic alignment.

Heatmap showing attention weights for head 4. — Head 4 attention pattern. The output projection combines all four heads into a unified representation.

Each head develops its own attention pattern. In trained translation models, researchers have observed heads that focus on positional alignment (source position $i$ attends to target position $i$ ), heads that track syntactic relationships, and heads that capture semantic similarities. The diversity across heads allows the model to capture multiple aspects of source-target alignment simultaneously. Research by Michel et al. (2019) found that many heads in trained models are prunable without significant performance loss, suggesting that the learned specialization tends to be concentrated in a subset of heads, though which heads are important varies by task.