Multi-head attention

A single attention head computes one weighted blend of values — one “question” about what to look for. That single subspace (one fixed way of looking at the vectors) can’t simultaneously represent syntax (grammar roles), coreference (which words refer to the same thing, like “it” and “the cat”), and position.

Multi-head attention runs h independent attention operations in parallel. Each head is free to learn a different relationship:

• one head for syntax
• another for entity coreference
• another for positional structure

The results are then concatenated.

That’s the entire idea. The implementation is just bookkeeping around the same Q·K·V math (queries, keys and values: the three vectors each token gets in self-attention).

The shape gymnastics

If d_model = 512 and you want 8 heads, each head gets d_head = d_model / n_heads = 64. Then:

1. Project the input to Q, K, V — each shape (seq, d_model).
2. Reshape so that the d_model dim splits into (n_heads, d_head).
3. Permute so the head dim is in front.
4. Run attention per head (in parallel using batched matmul).
5. Permute and concat heads back.
6. Final linear projection back to d_model.

Let’s do it in NumPy.

import numpy as np

rng = np.random.default_rng(0)

seq_len = 4
d_model = 8
n_heads = 2
d_head = d_model // n_heads     # 4

x = rng.standard_normal((seq_len, d_model))

# 1. Three projections (in real code these are nn.Linear)
W_q = rng.standard_normal((d_model, d_model))
W_k = rng.standard_normal((d_model, d_model))
W_v = rng.standard_normal((d_model, d_model))

Q = x @ W_q          # (seq, d_model)
K = x @ W_k
V = x @ W_v

# 2. Reshape: (seq, d_model) -> (seq, n_heads, d_head)
Q = Q.reshape(seq_len, n_heads, d_head)
K = K.reshape(seq_len, n_heads, d_head)
V = V.reshape(seq_len, n_heads, d_head)

# 3. Permute: (seq, n_heads, d_head) -> (n_heads, seq, d_head)
Q = Q.transpose(1, 0, 2)
K = K.transpose(1, 0, 2)
V = V.transpose(1, 0, 2)
print(f"Per-head shapes: Q {Q.shape}, K {K.shape}, V {V.shape}")

Per-head shapes: Q (2, 4, 4), K (2, 4, 4), V (2, 4, 4)

Now the attention computation runs per-head. NumPy does this with batched matmul — the leading n_heads axis is treated as a batch.

import numpy as np

rng = np.random.default_rng(0)
seq_len, d_model, n_heads = 4, 8, 2
d_head = d_model // n_heads

x = rng.standard_normal((seq_len, d_model))
W_q = rng.standard_normal((d_model, d_model))
W_k = rng.standard_normal((d_model, d_model))
W_v = rng.standard_normal((d_model, d_model))
W_o = rng.standard_normal((d_model, d_model))

Q = (x @ W_q).reshape(seq_len, n_heads, d_head).transpose(1, 0, 2)
K = (x @ W_k).reshape(seq_len, n_heads, d_head).transpose(1, 0, 2)
V = (x @ W_v).reshape(seq_len, n_heads, d_head).transpose(1, 0, 2)

# 4. Attention per head — batched matmul over the head axis
# scores: (n_heads, seq, d_head) @ (n_heads, d_head, seq) -> (n_heads, seq, seq)
scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_head)

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

weights = softmax(scores, axis=-1)
print(f"per-head attention weights: {weights.shape}")     # (2, 4, 4)

# 5. Per-head outputs: (n_heads, seq, seq) @ (n_heads, seq, d_head)
heads_out = weights @ V                                   # (n_heads, seq, d_head) = (2, 4, 4)
print(f"per-head outputs: {heads_out.shape}")

# 6. Concat heads back: -> (seq, d_model)
#    permute: (n_heads, seq, d_head) -> (seq, n_heads, d_head)
out = heads_out.transpose(1, 0, 2).reshape(seq_len, d_model)
print(f"after concat: {out.shape}")

# 7. Final output projection
out = out @ W_o
print(f"final: {out.shape}")

per-head attention weights: (2, 4, 4)
per-head outputs: (2, 4, 4)
after concat: (4, 8)
final: (4, 8)

Every step preserves a shape we can name. If any line errors, print the shape — it’ll tell you which axis got mismatched.

Why heads help

Imagine you’re parsing the sentence “The cat sat on the mat that the dog chased.” To understand “that,” a single head has to encode both what “that” refers to (the mat) AND the syntactic role of the clause.

Two heads can split this work:

• one head specializes in coreference (“that → mat”)
• another in syntactic structure (“the clause modifies mat”)

With 12 heads, the network has 12 simultaneous “lenses” on the data, and each can specialize.

Grouped Query Attention (GQA) — the modern shortcut

In many big LLMs, you’ll see something more efficient: GQA (grouped-query attention). It shares K and V across groups of Q heads. So a model might have 32 Q-heads but only 8 K-V heads, and each K-V is reused by 4 Q-heads.

The motivation: when the model generates text, it saves every earlier token’s K and V in a KV cache so it need not recompute them. That cache is the dominant memory cost. Fewer K-V heads = 4x less KV memory = 4x longer effective context or 4x bigger batch. Examples: Llama 3 8B and Mistral 7B use GQA, and in the Llama 2 family the 70B model did (8 KV heads) while the smaller ones kept full multi-head attention. The compute math is similar, but the memory savings are huge.

The ladder: MHA → MQA → GQA → MLA

GQA is one rung on a ladder, and the whole ladder is about one thing: shrinking the KV cache without losing too much quality. The knob is how many K/V heads you keep for a given number of query heads.

• MQA is the extreme: a single K/V head shared by every query head. The cache shrinks dramatically, but collapsing all keys/values into one head measurably hurts quality — too aggressive for frontier models.
• GQA is the pragmatic middle: enough K/V heads to keep quality, few enough to slash the cache. It’s the default in most open-weight LLMs.
• MLA (multi-head latent attention, introduced by DeepSeek) takes a different route: instead of storing full K and V, it caches a small compressed latent vector per token and reconstructs K and V from it at compute time. You cache far fewer numbers than even GQA, yet each head still gets its own (reconstructed) K/V — so quality stays close to full MHA. It trades a little extra compute for a much smaller cache.

The throughline: every rung past MHA buys cheaper inference memory. Which rung a model picks is a quality-vs-cache trade — and GQA/MLA are where the frontier currently sits.

In PyTorch

import torch.nn as nn

# The easiest path — built-in
mha = nn.MultiheadAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    batch_first=True,
)

out, _ = mha(query=x, key=x, value=x, is_causal=True)

For LLM-scale work, F.scaled_dot_product_attention (with a manual Q/K/V split) plus FlashAttention is faster and more flexible than the nn.MultiheadAttention wrapper.

In one breath

• One attention head is one “question” — a single subspace can’t represent syntax, coreference, and position at once.
• Multi-head runs h attention operations in parallel: split d_model into n_heads × d_head, attend per head, concat, project back with W_o.
• Each head learns a different relationship (one for coreference, one for local syntax, …) — more lenses on the same data.
• Head dimension stays ~64–128 while head count scales with model width; below ~32 each head is too small to be useful.
• GQA shares K/V across query-head groups to shrink the KV cache (MHA → MQA → GQA → MLA is the cache-vs-quality ladder); modern open LLMs default to GQA.

Quick check