Multi-head attention

One head can only ask one question

The calculation in the last lesson gave "it" a single set of weights: 0.5871 on server, 0.1859 on request, 0.2270 on itself. Look at what that vector is now committed to. It resolved a pronoun. It did not record that "it" is the subject of "was overloaded", or that the clause is causal, or that the request is the thing that got rejected. There was one budget and coreference spent it.

The limitation is structural, not a matter of capacity. A query vector points in one direction, so a dot product against it ranks keys along one axis of similarity. Ask "which noun does this stand for?" and you cannot simultaneously ask "what is being described here?" — the same numbers cannot encode both questions.

So run the calculation more than once. Each head gets its own W_Q, W_K and W_V, which means its own query, its own keys, its own notion of what counts as a match. Same input, same arithmetic, different learned weights — and a different word wins.

Two heads over the same three tokens: one lands on “server”, the other on “request”, because each head projects the input through its own learned weights

Head 1 is the one we traced: its query asks what noun "it" stands for, and server takes 0.5871. Head 2 projects the same three tokens through different weights, and its keys rank request first at 0.5880 — nearly the same magnitude, a different word entirely. Neither head is right or wrong. One is tracking coreference, the other is tracking what was acted on, and a representation of "it" that carries both is strictly better than one carrying either.

The transformer uses eight of these. Nobody assigns them jobs; the specialisations are whatever gradient descent finds useful, and in a trained model they are often not cleanly interpretable at all. What is guaranteed is only that the eight sets of weights start out random and different, so they cannot collapse into eight copies of the same question.

That leaves an arithmetic problem. Eight heads produce eight output vectors per token, and the feed-forward layer above expects one. Concatenating gives a vector eight times too wide — so the concatenated result is multiplied by one more learned matrix, W_O, which projects it back to the model width. W_O is not bookkeeping: it is where the model learns how to weigh what each head found.

One detail makes the whole thing nearly free. Each head works in d_model / 8 dimensions rather than the full width — 64 instead of 512 in the base transformer. Eight heads at 64 dimensions cost about what one head at 512 would. Multi-head attention buys eight perspectives at roughly the price of one, which is why it is not an optional refinement but the default.

The multi-head attention mechanism runs H independent attention operations in parallel, each with its own set of learned projection matrices, then concatenates and linearly transforms their outputs. The operation is $\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_{H}) @ W_{O}$ where $\text{head}_{h} = \text{Attention}(Q @ W_{Q}^{h}, K @ W_{K}^{h}, V @ W_{V}^{h})$ .

Why multiple heads

A single attention head learns one linear projection of queries and keys, which defines one notion of token-to-token relevance. Consider the sentence "The cat sat on the mat because it was tired." A single attention head resolving "it" must decide: does "it" refer to "cat" or "mat"? The head could learn a coreference pattern (pronouns attending to their antecedents), but then it cannot simultaneously learn a syntactic pattern (verbs attending to their subjects) or a positional pattern (attending to adjacent tokens). One set of Q/K projections imposes one linear similarity metric — one definition of "these two tokens are related."

Multi-head attention solves this by running H separate attention operations, each with independent projection matrices. Head 3 might learn to track subject-verb agreement while head 7 learns coreference and head 12 learns positional adjacency. The model does not assign these roles — they emerge from training. Clark et al. (2019) analyzed BERT's 144 attention heads (12 layers × 12 heads) and found heads that specialize in syntactic dependencies, heads that attend to the previous/next token, heads that attend to sentence separators, and heads with broad, nearly uniform attention.

The dimension split

Multi-head attention does not increase the total computation relative to a single head operating on the full dimension. The standard approach splits the model dimension across heads:

d_{\text{head}} = \frac{d_{\text{model}}}{H}

If d_model = 4096 and H = 32, each head operates in a 128-dimensional subspace. The key/query/value projections for each head map from d_model down to d_head:

W_Q^h: (d_model, d_head) — 4096 × 128 = 524,288 parameters per head
W_K^h: (d_model, d_head) — same
W_V^h: (d_model, d_head) — same

All H heads' projection matrices are typically implemented as a single large matrix. Instead of 32 separate (4096, 128) matrices, you store one (4096, 4096) matrix for W_Q, one for W_K, one for W_V, and slice the output into H chunks of d_head dimensions each. This is computationally identical but more efficient on GPU hardware because it becomes a single large matrix multiply instead of H small ones.

Head configurations in real models

GPT-2 small — d_model = 768, H = 12 heads, d_head = 64. 12 layers. 117M parameters total.
GPT-2 XL — d_model = 1600, H = 25 heads, d_head = 64. 48 layers. 1.5B parameters.
Llama 3 8B — d_model = 4096, H = 32 heads, d_head = 128. 32 layers. 8B parameters.
Llama 3 70B — d_model = 8192, H = 64 heads, d_head = 128. 80 layers.
GPT-3 — d_model = 12288, H = 96 heads, d_head = 128. 96 layers. 175B parameters.

Notice that d_head is almost always 64 or 128. This is not a coincidence — it is a hardware-driven choice. GPU tensor cores operate most efficiently on tile sizes that are multiples of 16, and 64/128 provide good utilization on A100 and H100 hardware. The number of heads scales with model size, while the per-head dimension stays fixed. Doubling the model from Llama 3 8B to 70B doubles d_model (4096 to 8192) and doubles the head count (32 to 64), but d_head remains 128.

Parameter count

Each multi-head attention layer has four weight matrices:

W_Q: (d_model, d_model) — all heads' query projections concatenated
W_K: (d_model, d_model) — all heads' key projections
W_V: (d_model, d_model) — all heads' value projections
W_O: (d_model, d_model) — output projection that recombines the heads

Total: 4 × d_model² parameters per attention layer (ignoring biases, which many modern models omit).

For Llama 3 8B: $4 \times 4096^{2} = 67, 108, 864$ parameters per attention layer. With 32 layers: ~2.15 billion attention parameters. The remaining ~5.85B parameters are in the feed-forward layers (which are 3 × d_model × d_ff per layer, with d_ff = 14336 for Llama 3 8B, using the SwiGLU architecture that has three weight matrices instead of two).

The output projection

After computing all H attention heads independently, their outputs are concatenated and projected:

python

head_h = Attention(Q @ W_Q^h, K @ W_K^h, V @ W_V^h)    # (N, d_head)
concat = [head_1 | head_2 | ... | head_H]                # (N, d_model)
output = concat @ W_O                                      # (N, d_model)

The output projection W_O is a (d_model, d_model) matrix. It serves a specific function: it allows the model to learn interactions between what different heads have found. Head 5 might detect that "it" refers to "cat" while head 9 detects that the current position requires an animate noun. The output projection can combine these signals: "position 8 needs the representation of an animate referent, and head 5 found that the referent is cat."

Without W_O, each head's contribution to the output would be confined to its own d_head-dimensional slice. The output projection enables cross-head information mixing.

Grouped-query attention (GQA)

Standard multi-head attention requires storing separate key and value tensors for each head in the KV cache during inference. For Llama 3 8B with a 128K context: $2 (K \text{and} V) \times 32 \text{heads} \times 128 d_{\text{head}} \times 128K \text{tokens} \times 2 \text{bytes}(\text{float16}) = ~2 \text{GB} \text{per} \text{layer}$ . Across 32 layers: ~64 GB just for the KV cache — dwarfing the 16 GB model weights.

Grouped-query attention (Ainslie et al. 2023) reduces this by sharing key/value heads across multiple query heads. Instead of 32 unique KV heads, Llama 3 8B uses 8 KV heads shared across 32 query heads (4 query heads per KV group). This cuts KV cache memory by 4x with minimal quality loss — the keys and values carry less head-specific specialization than the queries.

Multi-head attention (MHA) — H query heads, H KV heads. Full expressiveness, maximum KV cache.
Grouped-query attention (GQA) — H query heads, G KV heads (G < H). Llama 3 uses G = 8.
Multi-query attention (MQA) — H query heads, 1 KV head. Minimum KV cache, used in Falcon and PaLM. Some quality degradation on long-context tasks.

PyTorch implementation

python

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

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

        Q = self.W_Q(x)  # (B, N, d_model)
        K = self.W_K(x)
        V = self.W_V(x)

        # Split into heads: (B, N, d_model) -> (B, H, N, d_head)
        Q = Q.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)

        # Scaled dot-product attention per head
        scores = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5)  # (B, H, N, N)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        weights = F.softmax(scores, dim=-1)  # (B, H, N, N)
        attn_output = weights @ V             # (B, H, N, d_head)

        # Concatenate heads: (B, H, N, d_head) -> (B, N, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)

        # Output projection
        output = self.W_O(attn_output)  # (B, N, d_model)
        return output


d_model = 512
n_heads = 8
seq_len = 20
batch_size = 2

mha = MultiHeadAttention(d_model, n_heads)
x = torch.randn(batch_size, seq_len, d_model)
output = mha(x)

print(f"Input shape:  {x.shape}")       # (2, 20, 512)
print(f"Output shape: {output.shape}")  # (2, 20, 512)

total_params = sum(p.numel() for p in mha.parameters())
print(f"Parameters: {total_params:,}")  # 1,048,576 = 4 × 512²
print(f"Expected:   {4 * d_model ** 2:,}")

The key implementation detail is the reshape-and-transpose pattern. The single large matrix multiply W_Q(x) produces all heads' queries at once as a (B, N, d_model) tensor. The view-and-transpose reshapes this into (B, H, N, d_head), placing the head dimension before the sequence dimension. This layout lets the batched matrix multiply Q @ K.T compute all H heads' attention scores simultaneously as a single (B, H, N, N) operation.

← Previous