Module 3 — Multi-head attention

The five-line function from module 2 works. Yet no production model uses it alone. It runs in parallel copies called heads, each with its own projections, and their outputs are concatenated before a final projection. This module explains why, sizes the heads correctly, and adds the multi-head module to our red-thread Transformer.

One head sees one thing at a time

Look again at the attention matrix from module 2. Every row sums to one. A single row expresses how strongly the current token attends to every other token in one distribution. That distribution is inevitably a compromise: if a token wants to attend both to the subject of the sentence and to a punctuation mark far away, one softmax cannot express two peaks cleanly without smearing them.

Multiple heads solve the problem by giving the model several independent distributions per token. In practice a base Transformer uses 8 or 12 heads. Each head has its own $W_Q$, $W_K$, $W_V$, so each learns to look at something different: syntactic dependencies, coreference, positional patterns, semantic similarity. Post-hoc analyses on BERT and GPT show that heads do specialise, but rarely in ways you can name — one head follows the previous token, another jumps to the sentence subject, several look mostly redundant and can be pruned.

The formula, with several heads

Let $h$ be the number of heads and $d$ the model dimension (the embedding size). The head dimension is $d_k = d / h$, chosen so the total width stays the same:

$\text{head}_i = \text{Attention}(X W_Q^i, X W_K^i, X W_V^i), \quad i = 1, \dots, h.$

The heads run independently, then their outputs are stacked side by side and mixed by a final linear layer:

$\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) \, W_O.$

For $h = 8$ and $d = 512$, each head works in dimension 64, exactly the toy size we picked in module 2. That is not a coincidence: the paper's numbers were chosen to keep each head small enough that its dot products remain interpretable.

Parameter count, done properly

Miscounting parameters is a rite of passage. Do it slowly, once, and never worry about it again.

For one head with input dimension $d$ and head dimension $d_k$:

• $W_Q^i$: shape $d \times d_k$
• $W_K^i$: shape $d \times d_k$
• $W_V^i$: shape $d \times d_k$

Multiply by $h$ heads and add the output projection $W_O$ of shape $h \cdot d_k \times d = d \times d$:

$3 \cdot h \cdot d \cdot d_k + d^2 = 3 \cdot h \cdot d \cdot (d/h) + d^2 = 4 d^2.$

For $d = 512$, that is $4 \times 512^2 = 1\,048\,576$ parameters per multi-head module — about a million. Multiply by twelve layers and you already see the parameter mass of a small Transformer, before the feed-forward blocks that dominate the total. Module 5 adds those blocks and completes the accounting.

One matrix, three projections

Every serious implementation stores $W_Q$, $W_K$, $W_V$ stacked into a single matrix of shape $d \times 3d$. A single matmul produces $Q$, $K$, $V$ in one shot, then splits the output into three slices. The math is unchanged; the accelerator prefers one large matmul over three small ones.

In PyTorch: from single head to multi-head

We build multi-head attention on top of the function from module 2. The move is mostly reshaping: split the projected tensor across a "head" axis, run attention per head in one batched call, then merge back.

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0
        self.h = num_heads
        self.d_k = d_model // num_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, mask=None):
        B, N, D = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.h, self.d_k)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)   # (3, B, h, N, d_k)
        scores = q @ k.transpose(-2, -1) / self.d_k ** 0.5
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        out = (weights @ v).transpose(1, 2).reshape(B, N, D)
        return self.proj(out), weights

Save this file under mha.py. Module 5 wraps it in a residual block, module 6 stacks residual blocks into an encoder, and module 10 assembles everything.

What different heads learn

Interpretation is more art than science, but a few patterns recur across models trained on English text.

Kind of head	Behaviour	Detection
Previous token	large weight on position $t-1$ for every $t$	first sub-diagonal is bright
Positional / neighbourhood	weight decays with distance	narrow bright diagonal band
Punctuation	attention collapses on commas or periods	bright vertical stripes on those tokens
Coreference	pronouns look at their antecedent	long-range spikes
Content-based	matches semantic role rather than position	patterns that no simple rule captures

Module 10 plots the attention weights of our trained Transformer on the date-translation task. You will recognise a positional head that reads the day, a content head that copies the year, and a mostly redundant head that could probably be pruned.

Why not just add heads without bound

More heads sound better, so why does the paper stop at 8 or 12?

• Head dimension shrinks with head count. For fixed $d$, doubling $h$ halves $d_k$. Below $d_k = 32$ or so, each head becomes too narrow to represent useful patterns and adding more no longer helps.
• Total parameter count stays $4d^2$, whatever $h$. The gain from more heads has to come from specialisation, not from capacity. Beyond a point, the extra heads duplicate each other.
• Pruning studies on BERT show that removing half the heads costs less than one point on GLUE. That confirms the redundancy and reframes the design choice: pick the smallest $h$ that reaches your target, not the largest.

$d$ must be divisible by $h$

Almost every clean implementation asserts $d \bmod h = 0$. Choose $d = 512$ and $h = 7$ and the module will refuse to build, or worse, will silently round. If you tweak $h$, adjust $d$ so the head dimension stays a power of two, which the accelerator prefers.

In summary

• Multi-head attention runs $h$ parallel copies of the single-head operation, each with its own $W_Q$, $W_K$, $W_V$, and concatenates their outputs before a final projection $W_O$.
• The head dimension is $d_k = d / h$: the total width is preserved, and the parameter count for the whole module is $4d^2$.
• Different heads specialise: some follow neighbours, some track punctuation, some resolve coreference — but many are redundant, which is why pruning research is a live topic.
• In code, store $W_Q, W_K, W_V$ as a single matrix, reshape into a head axis, and reuse the module 2 attention function unchanged.

Next module: positional encoding, which fixes the fact that our attention operation, as it stands, does not know the order of the tokens.