Multi-Head Attention

Standard Scaled Dot-Product Attention takes a query matrix $Q$, key matrix $K$, and value matrix $V$, and computes

This produces one weighted average of value vectors per query, using one notion of "similarity." The problem is that natural language relations are heterogeneous: subject-verb agreement, coreference, negation scope, and positional adjacency are all different kinds of relevance, and a single softmax over a single projection space must compromise between them. Multi-head attention addresses this by linearly projecting the queries, keys, and values into $h$ different, lower-dimensional subspaces, running scaled dot-product attention independently in each subspace (a "head"), and then concatenating the results before a final linear projection. Formally, for head $i$,

and the overall output is

Here $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$, and $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$ are learned parameters. In the standard configuration, $d_k = d_v = d_{\text{model}} / h$, so splitting into $h$ heads keeps the total number of parameters and the total computation comparable to a single attention operation over the full $d_{\text{model}}$ dimension — the key architectural insight that makes multi-head attention essentially free relative to a single wide head. Terminology worth fixing precisely: $d_{\text{model}}$ is the model's hidden dimension (e.g., 512 or 4096); $h$ is the number of heads; $d_k$ is the per-head query/key dimension; $d_v$ is the per-head value dimension. The queries, keys, and values themselves may come from the same sequence (self-attention) or from two different sequences (cross-attention, used in Transformer Decoder Stack to attend over encoder outputs).

The computation proceeds in five stages: linear projection, splitting into heads, per-head scaled dot-product attention, concatenation, and a final output projection.

Step 1 — Project. The input sequence, represented as a matrix $X \in \mathbb{R}^{n \times d_{\text{model}}}$ ($n$ tokens), is linearly transformed into $Q = XW^Q$, $K = XW^K$, $V = XW^V$ for self-attention (or $Q = XW^Q$ with $K, V$ from a different sequence for cross-attention).

Step 2 — Split. Each of $Q, K, V \in \mathbb{R}^{n \times d_{\text{model}}}$ is reshaped/split along the feature dimension into $h$ chunks of size $d_k$ (or $d_v$), giving $h$ triples $(Q_i, K_i, V_i)$. In practice this split is implemented as a single large projection whose output is reshaped, not as $h$ separate matrix multiplies, which is what keeps the implementation efficient on parallel hardware.

Step 3 — Attend per head. Each head independently computes

where $M$ is an optional mask (e.g., the causal mask used in decoder self-attention, or a padding mask).

Step 4 — Concatenate. The $h$ outputs, each in $\mathbb{R}^{n \times d_v}$, are concatenated along the feature axis to form a matrix in $\mathbb{R}^{n \times h d_v} = \mathbb{R}^{n \times d_{\text{model}}}$.

Step 5 — Output projection. The concatenated matrix is passed through a shared linear layer $W^O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$, which mixes information across heads before the result is passed on (typically added to a residual connection and normalized, as part of a full Transformer Encoder Stack layer). ``` function MultiHeadAttention(X_q, X_kv, num_heads h, d_model, mask=None):

# X_q: queries source, shape (n_q, d_model)

# X_kv: keys/values source, shape (n_kv, d_model) (X_kv = X_q for self-attention) assert d_model % h == 0 d_k = d_model / h d_v = d_model / h Q_full = X_q @ W_Q # (n_q, d_model), W_Q: (d_model, d_model) K_full = X_kv @ W_K # (n_kv, d_model) V_full = X_kv @ W_V # (n_kv, d_model)

# reshape into heads: (n, d_model) -> (h, n, d_k) Q_heads = split_heads(Q_full, h, d_k) # h tensors of shape (n_q, d_k) K_heads = split_heads(K_full, h, d_k) # h tensors of shape (n_kv, d_k) V_heads = split_heads(V_full, h, d_v) # h tensors of shape (n_kv, d_v) head_outputs = empty list of length h for i in 0 .. h-1: scores_i = (Q_heads[i] @ transpose(K_heads[i])) / sqrt(d_k) # (n_q, n_kv) if mask is not None: scores_i = scores_i + mask # mask uses -infinity at disallowed positions weights_i = softmax(scores_i, axis=-1) # rows sum to 1 head_outputs[i] = weights_i @ V_heads[i] # (n_q, d_v) concat = concatenate(head_outputs, axis=-1) # (n_q, h d_v) = (n_q, d_model) output = concat @ W_O # (n_q, d_model), W_O: (d_model, d_model) return output ``` Two implementation details matter for correctness. First, the mask must be applied before* the softmax, as additive negative infinity (or a very large negative number), so that masked positions receive zero probability mass after normalization — applying it after softmax would leave residual probability on forbidden positions. Second, the split into heads must be a genuine partition of the feature dimension (each head sees a disjoint slice of $d_{\text{model}}$ coordinates), not an overlapping or randomly resampled subset, so that the total information available across all heads equals the total information in $Q$, $K$, $V$ with no duplication or loss.

Time complexity. For a sequence of length $n$, computing the projections $Q, K, V$ costs $O(n \, d_{\text{model}}^2)$ total (matrix multiplication of an $n \times d_{\text{model}}$ matrix by a $d_{\text{model}} \times d_{\text{model}}$ matrix), independent of how many heads it is later split into, since the split is just a reshape. Within each head, computing $Q_i K_i^\top$ costs $O(n^2 d_k)$, and applying the result to $V_i$ costs another $O(n^2 d_v)$. Summed over $h$ heads with $d_k = d_v = d_{\text{model}}/h$, this totals

exactly the same asymptotic cost as a single attention head operating at the full dimension $d_{\text{model}}$. The final output projection costs $O(n \, d_{\text{model}}^2)$. So the overall cost of multi-head attention is

identical in order to single-head attention at the same total dimension. This is the precise sense in which multi-head attention is "free": splitting into more heads does not change the asymptotic compute, it only changes what each unit of compute represents.

Space complexity. The dominant memory cost is storing the attention weight matrices, one $n \times n$ matrix per head, totaling $O(h \cdot n^2)$ entries — though since $d_k$ shrinks as $h$ grows, the total floating-point storage for weights is $O(h n^2)$ while the values stored per head shrink, so total activation memory remains $O(n^2 \cdot h + n \cdot d_{\text{model}})$, dominated in practice by the $h \, n^2$ term for long sequences. This is the well-known quadratic-in-sequence-length bottleneck of standard attention, unaffected by the head count in its scaling behavior with $n$.

Correctness / expressive-power argument. Two claims need justification: that splitting into heads does not lose expressive power compared to one big attention computation, and that it strictly gains power by allowing independent attention patterns. No loss: set $h = 1$, $d_k = d_v = d_{\text{model}}$, and $W^O = I$. This recovers exactly single-head scaled dot-product attention as a special case, so the multi-head formulation is strictly more general. Strict gain: with $h > 1$, each head computes its own softmax distribution over the sequence, $\text{softmax}(Q_i K_i^\top / \sqrt{d_k})$. Because softmax is a genuinely nonlinear, saturating function, the average of several different softmax distributions over value subspaces is not representable as a single softmax distribution applied uniformly across the same value space unless the per-head attention patterns happen to coincide. Concretely, consider two heads: head 1 attends almost entirely to position $j$, head 2 attends almost entirely to position $k \ne j$. A single softmax over the concatenated key space can produce mass on both $j$ and $k$, but it cannot simultaneously deliver head 1's exact value vector from position $j$ in one output subspace and head 2's exact value vector from position $k$ in a different, disjoint subspace, while combining both through independent weightings that a single shared softmax could not realize without conflating the two attention distributions. The output projection $W^O$ then linearly recombines these otherwise-inaccessible joint patterns, so the composed function class is strictly larger than what one shared-softmax head can express at equal total dimension. This is why multi-head attention is not merely a computational convenience but a genuine increase in the hypothesis space, consistent with the empirical finding that different heads learn qualitatively different attention patterns (positional, syntactic, coreferential) that a single head cannot simultaneously represent without averaging them away.

Tradeoffs. More heads means smaller $d_k$ per head, and if $d_k$ becomes very small (e.g., $d_{\text{model}} = 512$, $h = 64$ gives $d_k = 8$), each head's dot products become noisier and less expressive individually — there is a practical sweet spot (commonly $d_k \in \{32, \dots, 128\}$) below which adding heads degrades quality even though it does not change the asymptotic complexity. Also, the quadratic-in-$n$ term dominates for long sequences regardless of $h$, which is why long-context models turn to sparse or linear attention variants rather than tuning head count.

Take a toy setup with $d_{\text{model}} = 4$, $h = 2$, so $d_k = d_v = 2$, and a sequence of $n = 3$ tokens: "The cat sat." Suppose after embedding and projection (numbers simplified for hand-tracing), the per-head query, key, and value matrices are:

Head 1 (rows = tokens: The, cat, sat):

Head 2:

Step-by-step for Head 1, query "cat" (row 2, $Q_1 = [0, 1]$): Scores against each key ($Q_1 \cdot K_{1,j}$):

Softmax over $[0, 0.707, 0.707]$: $\exp(0)=1$, $\exp(0.707)\approx 2.028$, so weights $\approx [1/5.056,\, 2.028/5.056,\, 2.028/5.056] = [0.198, 0.401, 0.401]$. Weighted sum of $V_1$ rows: $0.198\cdot[1,0] + 0.401\cdot[0,1] + 0.401\cdot[2,2] = [0.198 + 0.802,\; 0.401 + 0.802] = [1.000, 1.203]$. So $\text{head}_1(\text{"cat"}) \approx [1.00, 1.20]$.

Step-by-step for Head 2, query "cat" (row 2, $Q_2 = [1, -1]$): Scores against each key ($Q_2 \cdot K_{2,j}$):

Softmax over $[0, 1.414, -0.707]$: $\exp(0)=1$, $\exp(1.414)\approx 4.114$, $\exp(-0.707)\approx 0.493$. Sum $\approx 5.607$. Weights $\approx [0.178, 0.734, 0.088]$. Weighted sum of $V_2$: $0.178\cdot[0,1] + 0.734\cdot[1,0] + 0.088\cdot[1,1] = [0.734+0.088,\; 0.178+0.088] = [0.822, 0.266]$. So $\text{head}_2(\text{"cat"}) \approx [0.82, 0.27]$.

Concatenate and project: the two head outputs are concatenated into a 4-dimensional vector $[1.00,\, 1.20,\, 0.82,\, 0.27]$, and this passes through $W^O \in \mathbb{R}^{4 \times 4}$ (omitted here for brevity, but any linear map) to produce the final multi-head attention output for the token "cat." Note that head 1's softmax was fairly flat (weights near 0.2, 0.4, 0.4), pooling broadly across the sentence, while head 2's softmax sharply favored "cat" itself (weight 0.734) — a concrete numerical illustration of the correctness argument above: two heads produced qualitatively different attention patterns over the same input that a single shared softmax could not have realized simultaneously.

Multi-head attention was introduced as part of the Transformer architecture, a paper written by a group at Google Brain and Google Research seeking to eliminate recurrence entirely from sequence models. The core insight — that splitting attention into parallel subspaces costs nothing asymptotically while unlocking new expressive power — built directly on top of the additive and dot-product attention mechanisms that had already proven useful for augmenting recurrent encoder-decoder models in machine translation.

Multi-head attention is the computational core of virtually every modern large language model, powering the self-attention layers in encoder architectures like BERT and RoBERTa, decoder-only architectures like the GPT family and LLaMA, and encoder-decoder architectures like T5 and the original machine translation Transformer. It is equally central outside pure NLP: vision Transformers apply the identical mechanism to image patches, and multimodal models use cross-attention variants of the same mechanism to align text tokens with image or audio features. Industrial inference systems optimize specifically around multi-head attention's computational pattern — techniques like key-value caching, grouped-query attention (sharing key/value projections across groups of heads to reduce memory bandwidth), and multi-query attention (a single shared key/value head across many query heads) are direct engineering responses to the cost structure analyzed above. Any system built on the Transformer Encoder Stack or Transformer Decoder Stack — chatbots, code assistants, translation services, search re-rankers — relies on multi-head attention as its fundamental building block, making it one of the highest-leverage ideas to understand deeply in modern NLP.

Ask the tutor: How does multi-head attention's cost and behavior change when combined with key-value caching during autoregressive decoding?