Multi-Head Attention › How Multi-Head Attention Works

Multi-head attention divides the model's capacity across $h$ parallel attention heads. Each head operates on a reduced dimension, so the total computation remains comparable to single-head attention with the full dimension.

The process unfolds in four steps:

1. Project the input into $h$ separate QKV representations
2. Compute scaled dot-product attention independently for each head
3. Concatenate all head outputs into a single tensor
4. Project the concatenated output back to the model dimension

Let's trace through each step with concrete mathematics and then implement it in code.

Given an input sequence $\mathbf{X} \in \mathbb{R}^{n \times d_{\text{model}}}$ with $n$ tokens and model dimension $d_{\text{model}}$, we project it into QKV representations for each head. If we have $h$ heads, we typically set the head dimension to $d_k = d_{\text{model}} / h$. This division ensures that the total computation across all heads remains comparable to a single attention operation over the full dimension.

For head $i$, we compute the QKV matrices through learned linear projections:

$$\mathbf{Q}_i = \mathbf{X} \mathbf{W}^Q_i, \quad \mathbf{K}_i = \mathbf{X} \mathbf{W}^K_i, \quad \mathbf{V}_i = \mathbf{X} \mathbf{W}^V_i$$

where:

• $\mathbf{X} \in \mathbb{R}^{n \times d_{\text{model}}}$: the input sequence with $n$ tokens, each represented as a $d_{\text{model}}$-dimensional vector
• $\mathbf{W}^Q_i \in \mathbb{R}^{d_{\text{model}} \times d_k}$: the learned query projection matrix for head $i$, transforming each token into a $d_k$-dimensional query vector
• $\mathbf{W}^K_i \in \mathbb{R}^{d_{\text{model}} \times d_k}$: the learned key projection matrix for head $i$, transforming each token into a $d_k$-dimensional key vector
• $\mathbf{W}^V_i \in \mathbb{R}^{d_{\text{model}} \times d_v}$: the learned value projection matrix for head $i$ (often $d_v = d_k$)
• $\mathbf{Q}_i, \mathbf{K}_i \in \mathbb{R}^{n \times d_k}$: the resulting query and key matrices for head $i$
• $\mathbf{V}_i \in \mathbb{R}^{n \times d_v}$: the resulting value matrix for head $i$
• $d_k$: the per-head dimension, typically $d_{\text{model}} / h$

Each head gets its own projection matrices. This allows it to learn different transformations of the input. Head 1 might project tokens to emphasize syntactic features, while head 2 emphasizes semantic similarity. different projection matrices create different "views" of the same input, each potentially highlighting different aspects of the token relationships.

Each head computes standard scaled dot-product attention using its own projections. The attention mechanism compares queries against keys to determine relevance weights, then uses those weights to aggregate values:

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

where:

• $\text{head}_i \in \mathbb{R}^{n \times d_v}$: the output of attention head $i$, containing contextual representations for all $n$ tokens
• $\mathbf{Q}_i \mathbf{K}_i^T \in \mathbb{R}^{n \times n}$: the matrix of raw attention scores, where entry $(j, k)$ measures how much query at position $j$ matches key at position $k$
• $\sqrt{d_k}$: the scaling factor that prevents dot products from growing too large as $d_k$ increases (covered in the previous chapter on scaled dot-product attention)
• $\text{softmax}(\cdot)$: applied row-wise to convert scores into attention weights that sum to 1 for each query position
• $\mathbf{V}_i$: the value matrix whose rows are aggregated according to the attention weights

This produces $h$ separate outputs, each of shape $\mathbb{R}^{n \times d_v}$. The heads operate completely independently at this stage; there is no information sharing between them. This independence is what allows different heads to specialize in detecting different types of relationships.

After computing all heads, we concatenate their outputs along the feature dimension. Each head produces an $n \times d_v$ matrix, and stacking these side-by-side creates a larger matrix:

$$\text{Concat}(\text{head}_1, \text{head}_2, \ldots, \text{head}_h) \in \mathbb{R}^{n \times (h \cdot d_v)}$$

where:

• $\text{head}_1, \ldots, \text{head}_h$: the $h$ individual head outputs, each of shape $\mathbb{R}^{n \times d_v}$
• $h \cdot d_v$: the concatenated feature dimension, combining all heads' outputs

If $d_v = d_k = d_{\text{model}} / h$, then $h \cdot d_v = d_{\text{model}}$, and the concatenated output has the same dimension as the original input. This dimensional consistency is intentional: it allows multi-head attention to be used as a drop-in replacement for single-head attention. The concatenation preserves each head's contribution separately, letting the subsequent projection learn how to combine them.

Finally, we apply a learned linear projection to the concatenated output. This projection maps the combined head outputs back to the model dimension:

$$\text{MultiHead}(\mathbf{X}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \mathbf{W}^O$$

where:

• $\text{Concat}(\text{head}_1, \ldots, \text{head}_h) \in \mathbb{R}^{n \times (h \cdot d_v)}$: the concatenated outputs from all $h$ attention heads
• $\mathbf{W}^O \in \mathbb{R}^{(h \cdot d_v) \times d_{\text{model}}}$: the learned output projection matrix
• $\text{MultiHead}(\mathbf{X}) \in \mathbb{R}^{n \times d_{\text{model}}}$: the final output, matching the input dimension

This projection serves two purposes. First, it combines information from all heads, allowing the model to learn how to mix their contributions. A head specializing in syntax might contribute more heavily for certain tokens, while a semantic head dominates for others. Second, it ensures the output dimension matches the input dimension, enabling residual connections in the transformer architecture.