Multi-Head Attention

In the previous chapters, we built up the complete picture of scaled dot-product attention: how queries match with keys, how values get aggregated, and how masking controls information flow. A single attention mechanism can capture one type of relationship between tokens. But language is rich with multiple simultaneous relationships. When reading "The cat sat on the mat because it was tired," you simultaneously track syntactic structure (subject-verb agreement), semantic relationships (what "it" refers to), and positional patterns (nearby words). Can we design attention to capture all of these at once?

Multi-head attention answers this question by running multiple attention operations in parallel, each with its own learned projections. Instead of one attention mechanism looking for one type of pattern, we have several "heads," each specializing in different aspects of the input. This architectural choice is central to transformer performance and appears in every major language model from BERT to GPT-4.

Think of multi-head attention as a team of specialized analysts reviewing the same document simultaneously. One analyst focuses on grammatical structure, flagging every subject-verb pair. Another scans for semantic similarity, grouping words with overlapping meanings. A third tracks discourse references, connecting every pronoun to its antecedent. When they pool their reports, the combined picture is far richer than any single analyst could produce alone. Multi-head attention works the same way: multiple learned "views" of the same input, each detecting a different kind of relationship, combined into one unified representation.

The architecture achieves this without increasing the fundamental parameter budget or computational cost of attention. By splitting the model dimension across heads, each head operates on a lower-dimensional slice of the representation. The total work done is roughly the same as running one large attention operation, but the outputs capture far more structure. This parameter efficiency is one reason multi-head attention became the standard building block for transformers rather than an expensive luxury.

Multi-head attention does not explicitly program any head to detect syntax or coreference. Those specializations emerge through gradient descent as the model learns to minimize its prediction loss. The heads discover what is useful from data alone, which means they can adapt to whatever patterns matter most for a given language and task. This learned specialization, arising without explicit supervision, is one of the more remarkable properties of the mechanism.

In this chapter we will build multi-head attention from first principles, working through the mathematics carefully, implementing it in NumPy, and then examining how different heads develop distinct behaviors. We will also look at how practitioners configure and scale multi-head attention in real systems, from the original 512-dimensional transformer to models with tens of billions of parameters.

Single-head attention computes one set of attention weights per position. Each token decides how to weight all other tokens, producing a single blended representation. This works well for capturing one dominant relationship, but it forces a choice: should "sat" attend mostly to its subject "cat," its location "mat," or its tense indicators?

Consider a more complex example: "The lawyer who the witness saw left." To parse this sentence correctly, a model must simultaneously track multiple dependencies. "Who" relates back to "lawyer" (as an antecedent). "Saw" connects to "witness" (as its subject) and to "who" (as its object). "Left" connects back to "lawyer" (not "witness") as its subject. A single attention head must somehow encode all these relationships in one set of weights.

The solution is straightforward: instead of computing attention once, compute it multiple times with different projections. Each "head" learns to focus on different aspects of the input. One head might specialize in syntactic relationships, another in coreference, and another in local context. By combining their outputs, the model builds a richer representation than any single head could achieve.

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:

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:

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:

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:

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.

Now that we've walked through each step conceptually, let's see how they combine into a complete formulation. Understanding this picture helps clarify how multi-head attention works: by running multiple specialized attention operations in parallel and learning to combine their insights.

The outer structure of multi-head attention takes the outputs from all heads and merges them:

This formula reads from inside out: first compute each head, then concatenate their outputs, then apply the output projection. But what exactly is each head doing? Each head applies its own learned projections before computing attention:

where:

• $\mathbf{Q}, \mathbf{K}, \mathbf{V}$: the QKV inputs (in self-attention, all three equal the input $\mathbf{X}$)
• $\mathbf{W}^Q_i, \mathbf{W}^K_i, \mathbf{W}^V_i$: the per-head projection matrices for head $i$, each transforming the full $d_{\text{model}}$-dimensional input into a smaller $d_k$-dimensional subspace
• $\text{Attention}(\cdot)$: the scaled dot-product attention function from the previous chapter
• $\mathbf{W}^O$: the output projection that learns how to combine all heads' contributions

The design is flexible. In self-attention, $\mathbf{Q} = \mathbf{K} = \mathbf{V} = \mathbf{X}$, so the same input sequence projects to all three roles. But the same multi-head attention mechanism also supports cross-attention in encoder-decoder models, where queries come from the decoder and keys/values come from the encoder. The mathematical structure remains identical; only the inputs differ.

Why should projecting to lower-dimensional subspaces and computing attention there be useful? Think of it this way: the full $d_{\text{model}}$-dimensional embedding space contains many different types of information about each token. Its syntactic role, semantic meaning, position in the sentence, relationship to neighboring words. A single attention operation in the full space must somehow balance all of these considerations when deciding what to attend to.

By projecting to a lower-dimensional subspace, each head can focus on a particular "slice" of the information. One head's projection matrices might learn to emphasize dimensions that capture syntactic structure, causing that head to attend based on grammatical relationships. Another head's projections might emphasize semantic dimensions, causing it to attend based on meaning similarity. The output projection then learns to weigh and combine these different perspectives.

This is analogous to how ensemble methods work in machine learning: multiple weak learners, each capturing part of the pattern, combine to form a stronger predictor than any single model could achieve.

There is a subtle but important distinction between multi-head attention and a simple ensemble of independent attention modules. In an ensemble, each member operates on the same full representation and the combination is usually an average or a vote. In multi-head attention, the combination is a learned linear transformation through $\mathbf{W}^O$. This output projection is not a passive merger; it actively learns how to mix head contributions in a context-dependent way. Heads that are more relevant for a given input can have their contributions amplified, while heads that detected spurious patterns can be suppressed. The projection also allows interactions between head outputs, so the final representation can encode relationships that no single head captured on its own.

In practice, this means you should think of the output projection as a learned integration step rather than a simple concatenation. The concatenation preserves all head outputs separately so the projection can distinguish between them. If you replaced concatenation with summation, you would lose the ability for the projection to treat different heads differently, and you would also lose the dimensional expansion that gives the projection enough freedom to mix contributions flexibly.

A natural question arises: if we're running $h$ separate attention operations, doesn't that multiply our parameter count by $h$? No. The dimension splitting ensures that multi-head attention uses approximately the same number of parameters as single-head attention with the full dimension.

Let's trace through the calculation:

1. Per head: Each head has three projection matrices ($\mathbf{W}^Q_i$, $\mathbf{W}^K_i$, $\mathbf{W}^V_i$), each of size $d_{\text{model}} \times d_k$. That's $3 \times d_{\text{model}} \times d_k$ parameters per head.

2. All heads: With $h$ heads, we multiply by $h$: total QKV parameters are $h \times 3 \times d_{\text{model}} \times d_k$. Here's the key insight: since $d_k = d_{\text{model}} / h$, we have $h \times d_k = d_{\text{model}}$. This simplifies our expression to $3 \times d_{\text{model}}^2$.

3. Output projection: The matrix $\mathbf{W}^O$ has size $(h \cdot d_v) \times d_{\text{model}}$. With $d_v = d_k = d_{\text{model}} / h$, the input dimension is $h \cdot (d_{\text{model}} / h) = d_{\text{model}}$, giving us $d_{\text{model}} \times d_{\text{model}} = d_{\text{model}}^2$ parameters.

4. Grand total: $3 \times d_{\text{model}}^2 + d_{\text{model}}^2 = 4 \times d_{\text{model}}^2$ parameters.

This is exactly what we'd have with single-head attention using the full dimension, plus one additional output projection. The key benefit of multi-head attention is that we get $h$ different attention patterns, each potentially specializing in different relationships, for essentially the same parameter budget as a single larger head.

Before turning to code, let's trace through a minimal example by hand to make the mechanics concrete. Suppose we have a sequence of three tokens: "cat chases mouse." We represent each token with a 4-dimensional embedding (a tiny model dimension chosen so the numbers stay readable), and we use $h = 2$ heads, giving each head $d_k = 2$ dimensions.

Our three token embeddings (rows of $\mathbf{X}$) are:

where:

• Row 1 ($[1, 0, 1, 0]$): the embedding for "cat"
• Row 2 ($[0, 1, 0, 1]$): the embedding for "chases"
• Row 3 ($[1, 1, 0, 0]$): the embedding for "mouse"

For head 1, suppose the learned projection matrices are:

These projections keep the first two embedding dimensions for queries and keys, and the last two dimensions for values. Computing $\mathbf{Q}_1 = \mathbf{X} \mathbf{W}^Q_1$:

The raw attention scores for head 1 are $\mathbf{Q}_1 \mathbf{K}_1^T / \sqrt{d_k}$ with $d_k = 2$:

Scaling by $\sqrt{2} \approx 1.414$ gives:

After softmax (row-wise), row 3 ("mouse") ends up attending most strongly to itself because "mouse" has the largest self-similarity score. Row 1 ("cat") splits attention roughly equally between itself and "mouse." Row 2 ("chases") splits between "chases" and "mouse."

Notice that head 2, which uses different projection matrices, will compute a completely different set of scores from the same input $\mathbf{X}$. Even if head 2's projections emphasize different embedding dimensions, it still produces an $n \times d_v$ output that captures a distinct relationship pattern. When you concatenate both head outputs and apply $\mathbf{W}^O$, the final representation for each token can draw on both perspectives simultaneously.

In practice, trained models rarely use identity-like projection matrices. The projections are dense, full-rank matrices learned to maximize task performance. But the worked example illustrates the core mechanism: the same token sequence can be "seen" through multiple lenses, and those lenses are entirely determined by what the model learns to be useful.

With the mathematics clear, let's translate these ideas into working code. We'll build multi-head attention from scratch using NumPy, making every operation explicit so you can see exactly how the formulas map to implementation. This hands-on approach will solidify your understanding and prepare you to work with production implementations in PyTorch or TensorFlow.

Our implementation needs two building blocks: first, the scaled dot-product attention we covered in the previous chapter; second, the multi-head wrapper that orchestrates multiple attention operations in parallel.

We start with a function that computes attention for tensors that already include a head dimension. This allows us to process all heads simultaneously through batched matrix multiplication:

Notice the tensor shape: (batch, heads, seq_len, d_k). By placing the head dimension second, we can compute attention scores for all heads at once. The matrix multiplication Q @ K.transpose(...) produces a (batch, heads, seq_len, seq_len) tensor containing $h$ separate attention matrices, one per head. This is the key to efficient parallel computation.

Now we build the main module. The MultiHeadAttention class encapsulates all the projection matrices and orchestrates the four-step process we described earlier.

A key implementation detail: instead of storing separate $\mathbf{W}^Q_i$, $\mathbf{W}^K_i$, $\mathbf{W}^V_i$ matrices for each head, we use a single combined projection matrix W_qkv of shape (d_model, 3 * d_model). This computes all QKV representations for all heads in one matrix multiplication. The result is then split and reshaped to separate the heads. This approach is more efficient on GPUs because it maximizes parallelism and minimizes memory transfers.

The split_heads and combine_heads methods handle the tensor reshaping that makes parallel head computation possible. split_heads takes a tensor of shape (batch, seq_len, d_model) and reorganizes it to (batch, num_heads, seq_len, d_k), separating the head dimension. combine_heads reverses this transformation after attention is computed.

Let's verify our implementation produces the expected output shapes and that the attention weights are properly normalized:

Each row sums to exactly 1.0, confirming that softmax normalization is working correctly. This is expected: for each query position, the attention weights across all key positions must form a valid probability distribution.

Let's visualize how multi-head attention changes the token representations. We'll use PCA to project the 64-dimensional embeddings down to 2D, allowing us to see how the representations shift after passing through the attention layer.

The visualization shows how multi-head attention repositions each token's representation in embedding space. The arrows indicate the direction and magnitude of change. Tokens that attended strongly to semantically related positions move toward each other, while tokens that focused on different aspects of the input shift in different directions. This is the contextual enrichment in action: each position now carries information gathered from across the sequence.

The true power of multi-head attention becomes apparent when we visualize what different heads are doing. Even with random initialization (before any training), the different projection matrices cause each head to compute different attention patterns. After training on language data, these patterns become meaningful: heads specialize to detect specific types of relationships.

Let's visualize the attention patterns for each head on a simple sentence. We'll create random embeddings and see how the four heads distribute their attention differently:

The heatmaps reveal that even with random initialization, each head produces a distinct attention pattern. Some heads distribute attention more uniformly, while others concentrate it on specific positions. In a trained model, these differences become meaningful and interpretable: one head might consistently attend to the previous word (useful for language modeling), another to syntactically related words (useful for parsing), and another to semantically similar words (useful for coreference resolution).

The key observation is that the same input, "The cat sat on the mat," produces four different perspectives on token relationships. This diversity is precisely what we wanted when we motivated multi-head attention: the ability to capture multiple types of dependencies simultaneously.

We can quantify how "focused" or "distributed" each head's attention is using entropy. A head with low entropy concentrates attention on a few positions, while a head with high entropy distributes attention more uniformly. This provides a numerical measure of the diversity we observed visually.

Heads with lower entropy relative to the maximum are more "focused," concentrating their attention on specific positions. Heads closer to the maximum entropy distribute attention more uniformly. This variation in focus is another dimension of head diversity: some heads act as sharp selectors, while others aggregate information more broadly.

Research on trained transformers reveals that attention heads develop distinct specializations. The emergence of interpretable, consistent patterns across very different models is one of the most striking properties of multi-head attention. It suggests that these specializations are not accidents of initialization or quirks of a particular training run; they reflect real structure in language that gradient descent reliably discovers.

Some commonly observed patterns include:

• Positional heads: Attend to fixed relative positions (e.g., always attend to the previous token, or always attend to the first token)
• Syntactic heads: Track grammatical relationships like subject-verb or modifier-noun connections
• Coreference heads: Link pronouns to their antecedents across the sequence
• Rare word heads: Focus on infrequent tokens that carry high information content
• Separator heads: Attend to punctuation and sentence boundaries

This specialization emerges purely from training on language modeling objectives. No explicit supervision tells heads what to focus on; they discover useful patterns through gradient descent.

What does this mean in practice? When you deploy a transformer model, its attention heads are not interchangeable. Some heads carry far more weight for specific task types: syntactic heads matter most for tasks requiring grammatical understanding, while semantic heads matter most for tasks requiring world-knowledge reasoning. This is one reason why fine-tuning a pre-trained transformer often works so well: the heads that are relevant for your task already exist in the pre-trained model, and fine-tuning adjusts their relative weights rather than learning entirely new patterns from scratch.

The three patterns illustrate different specializations:

• Previous token head: The diagonal pattern shows each token attending primarily to its predecessor. This captures local sequential dependencies, useful for language modeling where the previous word strongly predicts the next.

• [CLS] aggregation head: The first column lights up, showing all tokens attending to the special classification token. This pattern aggregates information from the entire sequence into one position, commonly used for classification tasks.

• Local context head: The banded pattern shows attention concentrated within a local window. This captures phrase-level relationships without being distracted by distant tokens.

How much does multi-head attention help? Let's compare the representational capacity of single-head versus multi-head attention on a simple task: capturing diverse pairwise relationships.

The diversity metric reveals a clear pattern: more heads lead to more varied attention patterns. With a single head, the diversity score is 0 by definition since there's nothing to compare. With 2 heads, the model can already capture distinct patterns. As we increase to 4 and 8 heads, diversity continues to grow, though the incremental gains diminish. This suggests that beyond a certain point, additional heads may begin to learn redundant patterns.

When implementing multi-head attention in production systems, several practical considerations come into play.

The original transformer paper used 8 heads with a 512-dimensional model, giving each head 64 dimensions. Modern large models often use more heads, but the per-head dimension typically stays between 64 and 128.

The visualization reveals a consistent design pattern: as model dimension increases, the number of heads scales proportionally to maintain per-head dimensions between 64 and 128. This suggests an empirically validated "sweet spot" for head expressiveness. Too few dimensions per head limit each head's ability to capture complex patterns; too many heads with too few dimensions fragment the representation unhelpfully.

Although multi-head attention computes $h$ separate attention operations, modern implementations parallelize this efficiently. The key insight is that we can stack all heads into a single tensor and use batched matrix multiplication.

This batched approach processes all heads with the same matrix operations, using GPU parallelism. The computational cost is essentially equivalent to single-head attention with some reshaping overhead.

Research has shown that not all attention heads are equally important. Some heads can be removed after training with minimal impact on performance. This head pruning technique reduces computation and memory for inference.

Studies on BERT found that many heads could be pruned without significant accuracy loss, suggesting some redundancy in the learned representations. Michel et al. (2019) found that on the WMT 2014 English-to-German translation task, up to 20 out of 48 heads could be eliminated at test time with less than 0.25 BLEU-point loss. On some GLUE tasks, models with as few as one attention head per layer performed competitively with the full 12-head BERT-base. This has practical implications for deploying models in resource-constrained environments: by identifying and pruning redundant heads before serving a model, you can reduce memory footprint and inference latency without significant quality regression.

The practical takeaway is that you should treat head count as a tunable hyperparameter rather than a fixed design choice. If you are deploying a model to an edge device or a latency-sensitive production system, structured pruning of attention heads is often one of the first compression techniques worth trying.

Multi-head attention is one of the key innovations that made transformers successful. By allowing the model to jointly attend to information from different representation subspaces at different positions, it provides a flexible mechanism for capturing complex relationships in language. Each head learns its own projection matrices, creating specialized views of the input that combine to form rich contextual representations.

The primary limitation is computational. Multi-head attention still has $O(n^2)$ complexity in sequence length $n$, and the overhead of multiple heads adds constant factors to both computation and memory. For a sequence of 4,096 tokens with 32 heads, the model must compute and store 32 separate attention matrices, each of size $4096 \times 4096$. This works out to over half a billion floating-point values per layer, just for the attention weights. At 32 transformer layers and 16-bit precision, the attention weights alone occupy roughly 32 GB of memory at inference time for a single sequence. This quadratic scaling remains the primary bottleneck for very long documents, motivating research into efficient attention variants: sparse attention attends to a structured subset of positions, linear attention reformulates the computation to avoid explicit attention matrices entirely, and FlashAttention optimizes memory access patterns on GPU hardware to reduce wall-clock time even without changing the asymptotic complexity.

Another consideration is interpretability. While we can visualize attention patterns and identify heads that appear to track specific linguistic relationships, understanding what each head learned remains challenging. The connection between attention weight magnitude and information flow is not as direct as it might seem: a head that assigns low attention weight to a token is not necessarily ignoring it, because the value vectors themselves carry information, and the output projection $\mathbf{W}^O$ can amplify or suppress contributions from each head. Jain and Wallace (2019) showed that attention weights do not reliably explain model predictions in the way that gradient-based saliency methods do, because high attention to a token does not guarantee that the token influenced the output. This makes it difficult to debug unexpected model behavior or provide faithful explanations to end users.

The observation that many heads can be pruned without significant accuracy loss adds another dimension to the interpretability challenge. If a head can be removed with negligible performance drop, it is unclear whether that head was doing nothing useful, or whether its function was being duplicated by another head. Both explanations are plausible, and distinguishing between them requires careful ablation studies rather than visual inspection of attention patterns alone.

Despite these limitations, multi-head attention works exceptionally well across virtually all NLP tasks. The ability to capture multiple types of relationships simultaneously, without explicit supervision about what those relationships should be, is a key reason why transformers generalize so well. From named entity recognition to machine translation to question answering to code generation, multi-head attention provides the representational flexibility needed to handle diverse linguistic phenomena. The mechanism also scales with model size: performance continues to improve as models grow, unlike older architectures such as recurrent networks. This scalability, combined with the parallelism that makes training tractable, is what cemented transformers as the dominant architecture for language AI.

It is worth stepping back to appreciate how far this mechanism has come from its origins. The original transformer paper introduced multi-head attention as a component of a sequence-to-sequence model for machine translation. In that context, it was one of several architectural choices competing with recurrent and convolutional baselines. Less than a decade later, scaled-up versions of the same mechanism form the core of models that write code, summarize legal documents, and pass medical licensing exams. The fundamental equations in this chapter are unchanged from Vaswani et al. (2017). The equations remained the same; progress came from training larger models on more data with more computation. Understanding multi-head attention at the level of equations and implementation is therefore not just historical curiosity: it is the foundation on which all modern large language models are built.

When configuring multi-head attention, the following parameters have the most significant impact on model behavior:

• num_heads (h): The number of parallel attention heads. Common values range from 8 to 96 depending on model size. More heads enable greater specialization but require the model dimension to be evenly divisible. Start with 8 heads for smaller models and scale proportionally with model dimension.

• d_model: The model dimension, which determines the total representation capacity. Must be divisible by num_heads. Standard values include 512 (original transformer), 768 (BERT-base), 1024 (BERT-large), and 12288 (GPT-3).

• d_k (head dimension): The dimension per head, computed as d_model // num_heads. Values between 64 and 128 work well in practice. Smaller values allow more heads but limit each head's expressiveness.

• d_v (value dimension): Often set equal to d_k. Can be set independently if you want asymmetric query-key versus value dimensions, though this is uncommon.

• Initialization scale: Projection matrices are typically initialized with variance $\frac{2}{d_{\text{model}}}$ or similar schemes. Poor initialization can cause attention weights to become too uniform or too peaked early in training.

Multi-head attention extends single-head attention by running multiple attention operations in parallel, each with its own learned projections. This architectural choice significantly impacts model capacity and performance.

Key takeaways from this chapter:

• Parallel heads: Multi-head attention computes $h$ independent attention operations, each learning different relationship patterns between tokens. Each head can specialize in detecting specific linguistic phenomena.
• Dimension splitting: The model dimension $d_{\text{model}}$ is divided among heads, with each head operating on dimension $d_k = d_{\text{model}} / h$. This keeps total computation comparable to single-head attention while enabling multiple perspectives.
• Four-step process: (1) Project input to Q, K, V for each head using learned matrices $\mathbf{W}^Q_i$, $\mathbf{W}^K_i$, $\mathbf{W}^V_i$; (2) compute scaled dot-product attention independently per head; (3) concatenate head outputs; (4) apply final output projection $\mathbf{W}^O$.
• Head specialization: Different heads learn to focus on different aspects: positional patterns, syntactic relationships, semantic similarity, or coreference links. This emerges from training without explicit supervision.
• Parameter efficiency: The total parameter count is $4 \times d_{\text{model}}^2$, similar to single-head attention with an output projection. The gain is in representational diversity, not parameter count.
• Practical implementations: Efficient implementations batch all heads into single tensor operations, using GPU parallelism through combined QKV projections.

In the next chapter, we'll examine the computational complexity of attention in detail. Understanding the $O(n^2)$ scaling (where $n$ is sequence length) and its implications is essential for working with long sequences and motivates the efficient attention variants used in modern long-context models.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about multi-head attention.