One-sentence summary: Multi-Head Attention organizes the projected width into several smaller heads, gives each head its own Q, K, and V view, then concatenates the results and mixes them with W_O. The heads may learn different patterns, but their jobs are learned—not assigned in advance.

11.1 Why Multiple Heads?

11.1.1 The Limit of a Single Head

Chapter 10 covered the complete Attention computation — but that was single-head Attention: one Query-Key score map and one Value mixture.

A single head is not limited to one simple relationship; its map can be complex. The narrower limitation is that the layer has only one Q/K/V projection view, so every relationship must share the same matching and mixing space.

Take this sentence:

“The luthier noticed the theorbo was cracked, so she returned it to its case.”

Fully understanding it requires tracking several kinds of relationships simultaneously:

• Syntactic: “returned” takes “she” as its subject
• Coreference: “she” refers to the luthier, while “it” refers to the theorbo
• Positional / phrase structure: “to its case” attaches to “returned”
• Causal: “so” connects the crack with the decision to put the instrument away

Multiple heads give the model several parallel subspaces in which to represent these relationships. That does not guarantee a clean one-head-one-job decomposition.

11.1.2 The Solution: Multiple Heads in Parallel

Multi-Head Attention's core idea is to run several narrower Attention computations in parallel, giving them an opportunity to learn different matching and routing patterns.

Head 1: might focus on syntactic structure (subject-verb-object)
Head 2: might focus on coreference (pronouns and nouns)
Head 3: might focus on local proximity (neighboring tokens)
Head 4: might focus on semantic similarity (related concepts)
...

Then their outputs are concatenated and W_O learns how to mix the resulting dimensions.

11.1.3 An Analogy

Imagine analyzing a painting. One pair of eyes may focus on color; several pairs can examine color, shape, texture, and composition at the same time. Multi-Head Attention gives the model several such views—but the views are learned, not manually assigned.

11.2 Splitting Into Heads

11.2.1 The Dimension Split

The key operation is reshaping the final dimension of each combined projection into a head axis and a per-head width. We are not cutting the raw token embedding into arbitrary pieces: Q, K, and V have already passed through learned projections.

Using K (Key) as the example, with:

• d_model = 512
• num_heads = 4
• In this classic equal-width example, d_head = d_k = d_v = d_model / num_heads = 512 / 4 = 128

The split unfolds as:

Combined K: [batch_size, seq_length, num_heads × d_k]
          = [4, 16, 512]
              ↓
Reshape:    [batch_size, seq_length, num_heads, d_k]
          = [4, 16, 4, 128]
              ↓
Transpose:  [batch_size, num_heads, seq_length, d_k]
          = [4, 4, 16, 128]

11.2.2 Why Transpose?

The transpose brings num_heads to the second axis, giving [batch, num_heads, seq_len, d_k]. This means:

• For each sequence in the batch
• We have num_heads independent Attention computations
• Each one processes seq_len positions
• Each Query and Key position uses a d_k-dimensional vector

This layout lets one batched tensor operation compute all heads in parallel. The score computations use separate slices, but the heads are not isolated little models: W_O, the residual path, and the shared loss train them jointly.

11.2.3 The Same Split Applies to Q, K, and V

Q: [4, 16, 512] → [4, 4, 16, 128]  # d_k
K: [4, 16, 512] → [4, 4, 16, 128]  # d_k
V: [4, 16, 512] → [4, 4, 16, 128]  # d_v

We now have four per-head (Q, K, V) sets, ready to be computed in parallel.

11.2.4 Two Equivalent Implementations

Without extra cross-head structure, concatenating the small matrices along their output dimension produces the same combined matrix. That is why these two views are mathematically equivalent:

Conceptual view: each head has projection slices W_i^Q, W_i^K, and W_i^V. Head i computes Q_i = X @ W_i^Q with a [d_model, d_k] matrix.

Practical view: one combined W_Q generates the full Q, then we reshape its last dimension into num_heads slices.

Real implementations use the practical view because a single large matrix multiplication is more GPU-efficient than many small ones. The GPU prefers large, contiguous operations over many small scattered ones.

11.3 Computing All Heads in Parallel

11.3.1 Each Head Computes Its Own Score Map

After the split, every head executes the same Attention formula independently:

For each head h = 1, 2, 3, 4:
    scores_h   = Q_h @ K_h^T    [4, 16, 128] @ [4, 128, 16] = [4, 16, 16]
    weights_h  = softmax(scores_h / sqrt(d_k) + M_h)
    output_h   = weights_h @ V_h    [4, 16, 16] @ [4, 16, 128] = [4, 16, 128]

11.3.2 Dimension Tracking

Q @ K^T for one head:

Q:   [4, 4, 16, 128]
     batch  heads  seq  d_k

K^T: [4, 4, 128, 16]
     batch  heads  d_k  seq

Q @ K^T: [4, 4, 16, 16]
         batch  heads  seq  seq

Softmax(Q @ K^T / sqrt(d_k) + M) @ V:

Attention Weights: [4, 4, 16, 16]
                   batch  heads  seq  seq

V: [4, 4, 16, 128]
   batch  heads  seq  d_v

Output: [4, 4, 16, 128]
        batch  heads  seq  d_v

11.3.3 What the Parallelism Gets You

For fixed total width, splitting into four heads does not multiply the leading QK^T or weights @ V arithmetic: each head is narrower and their widths add back to 512. The benefit is representational—several learned score maps and Value routes—not a guarantee that each head develops a clean specialization.

11.4 Merging the Heads Back

11.4.1 Concatenation

After all heads compute their output, we concatenate them back into the full model dimension:

Head outputs: [4, 4, 16, 128]
              batch  heads  seq  d_v
                   ↓
Transpose:    [4, 16, 4, 128]
              batch  seq  heads  d_v
                   ↓
Concatenate:  [4, 16, 512]
              batch  seq  d_model

The concatenation operation just merges the last two dimensions. In this example, 4 heads × 128 dimensions = 512 dimensions.

11.4.2 The Output Projection W_O

Concatenation is mechanical. It puts the heads' outputs next to each other but does not mix their dimensions. That is what W_O is for:

A @ W_O
[4, 16, 512] @ [512, 512] = [4, 16, 512]

W_O is a learned projection matrix. Its job:

1. Mix information across heads — what each head learned can now influence the others
2. Project the concatenated representation into a unified space
3. Let the model decide how to weight each head's contribution

11.4.3 Why W_O Matters

Before W_O, each slice of the concatenated vector comes from one head. W_O forms learned combinations across all those dimensions before the Attention branch returns to the residual stream.

11.5 Comparing the Outputs: Before and After W_O

11.5.1 A vs A @ W_O

Before W_O (A):

• Shape: [16, 512]
• Values: the raw concatenation of all heads' output vectors

After W_O (A @ W_O):

• Shape: [16, 512]
• Values: a mixed, projected representation

Same shape, different content. The post-W_O tensor is the Attention branch output; it is combined with the residual stream. Whether LayerNorm appears before or after the sublayer depends on the model's pre-norm or post-norm design.

11.6 Full Multi-Head Attention Flow

11.6.1 End to End

Input X [batch, seq, d_model]
         ↓
Generate Q, K, V (via W_Q, W_K, W_V)
         ↓
Reshape into heads [batch, num_heads, seq, d_head]
         ↓
Compute Attention independently per head
         ↓
Concatenate [batch, seq, d_model]
         ↓
Output projection (@ W_O)
         ↓
Output [batch, seq, d_model]

11.6.2 PyTorch Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        if d_model % num_heads != 0:
            raise ValueError("d_model must be divisible by num_heads")
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # This teaching example uses d_k = d_v = head_dim.
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, x, allowed_mask=None):
        batch_size, seq_len, _ = x.shape

        # 1. Generate Q, K, V
        Q = self.W_Q(x)   # [batch, seq, d_model]
        K = self.W_K(x)
        V = self.W_V(x)

        # 2. Split into heads
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim)

        # Transpose: [batch, num_heads, seq, head_dim]
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # 3. Attention per head
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)

        if allowed_mask is not None:
            # Broadcastable to [batch, heads, query_len, key_len].
            scores = scores.masked_fill(~allowed_mask, float('-inf'))

        attention_weights = F.softmax(scores, dim=-1)
        attention_output = torch.matmul(attention_weights, V)

        # 4. Merge heads
        attention_output = attention_output.transpose(1, 2)    # [batch, seq, heads, head_dim]
        attention_output = attention_output.contiguous().view(
            batch_size, seq_len, self.d_model
        )

        # 5. Output projection
        output = self.W_O(attention_output)

        return output

11.7 Key Numbers

11.7.1 Parameter Count

Multi-Head Attention has four weight matrices:

Counting only these four projection weights, the total is 4 × d_model². With biases, add 4 × d_model. Architectures such as MQA and GQA use different projection widths; Chapter 23 covers those variants.

For GPT-2 Small (d_model = 768), the four weights contain 4 × 768² = 2,359,296 parameters. Including its QKV and output biases gives 2,362,368, still about 2.36 million per Attention layer.

11.7.2 Common Configurations

These examples use head widths of 64 or 128, but that is a design choice in these models, not a universal law. Modern architectures can also decouple the number of Query heads from the number of Key/Value heads.

11.8 What Do the Heads Actually Learn?

11.8.1 Observed Patterns from Research

One BERT analysis observed fixed-offset, delimiter, syntactic, and coreference patterns, while also finding similar behaviors among heads in the same layer. The table lists possible learned patterns, not jobs assigned to head numbers in advance (Clark et al., 2019).

11.8.2 A Practical Example

For intuition, imagine possible patterns for that luthier sentence. This is not a measurement from a particular model:

Head 1 (positional):  "returned" gives weight to the nearby "she"
Head 2 (syntactic):   "returned" gives weight to "she" as its subject
Head 3 (semantic):    "cracked" and "theorbo" exchange information
Head 4 (coreference): "it" gives weight to "theorbo"

11.8.3 Head Redundancy

Not all heads are equally important. Michel et al. found that, for the translation models and BERT tasks they evaluated, many heads could be removed at test time without a significant metric drop, and some layers could be reduced to one head (Michel et al., 2019). That is evidence of model- and task-dependent redundancy—not proof that most heads in every Transformer are useless.

11.9 Multi-Head vs Single-Head

11.9.1 Compute Comparison

For d_model = 512, num_heads = 8, and classic equal-width MHA with d_head = 64:

Single head (width 512):

• Q @ K^T: [seq, 512] @ [512, seq] → O(seq² × 512)

Eight heads (width 64 each):

• Each head: [seq, 64] @ [64, seq] → O(seq² × 64)
• Total: 8 × O(seq² × 64) = O(seq² × 512)

For the QK^T term, the multiply-add total is the same; the same is true of weights @ V. With fixed total projection width, the leading QKV and W_O work also does not grow merely because the tensor is reshaped into more heads. Exact runtime can still differ with kernels, memory layout, and hardware.

11.9.2 Why Not More Heads?

In this fixed-d_model, equal-width setup, more heads means a smaller per-head width:

d_head = d_model / num_heads

If d_head becomes too small, each head may have too few dimensions for a useful subspace. Values such as 64 and 128 are choices made by the classic models above, not a universal optimum; architecture and experiments decide the trade-off.

11.10 Part 3 Checkpoint

We have now assembled Multi-Head Attention. Part 3 still has Chapter 12, which separates the Attention output from the parameters that training updates. Here is the path so far:

The complete Multi-Head Attention formula:

Where:

Chapter Checklist

After this chapter, you should be able to:

• Explain why a single Attention head has limitations.
• Derive d_head = d_model / num_heads for classic equal-width MHA.
• Trace dimension changes through split, compute, and merge.
• Explain what W_O does after concatenation.
• Describe the kinds of patterns different heads can learn.

See You in the Next Chapter

Chapter 12 closes Part 3 by separating three things that are easy to blur together: the hidden state produced during a forward pass, the embedding parameter table, and the model parameters updated by backpropagation.