Multi-head attention

A single attention head computes one weighted blend of values — one “question” about what to look for. That single subspace can’t simultaneously represent syntax, coreference, and position. Multi-head attention runs h independent attention operations in parallel, each free to learn a different relationship — one head for syntax, another for entity coreference, another for positional structure — then concatenates the results.

That’s the entire idea. The implementation is just bookkeeping around the same Q·K·V math.

TryMulti-head attention

Four heads, four specializations

Each head runs separate attention over the same 7-token sentence. Click a head to enlarge it and read what it learned. Then toggle Combine heads to see how their outputs concatenate into a richer view.

Thecatchaseditquicklyonmat

Head 1

Previous-tokenHead 2

Local windowHead 3

Syntactic (verb→subj)Head 4

Coreference

Click a head above to enlarge it and read its specialization.

Combine heads

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 today’s big LLMs, you’ll see something more efficient: GQA (grouped-query attention) shares K and V across groups of Q heads. So a model might have 32 Q-heads but only 8 K-V heads. Each K-V is reused by 4 Q-heads.

The motivation: during inference, the KV-cache is the dominant memory cost. Fewer K-V heads = 4x less KV memory = 4x longer effective context or 4x bigger batch. Llama 2/3, Mistral, Qwen all use GQA. The compute math is similar; 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.

Variant	K/V heads (for 32 Q-heads)	KV cache	Quality
MHA — multi-head	32 (one per query head)	1× (baseline)	best
MQA — multi-query	1 (all queries share it)	32× smaller	noticeable drop
GQA — grouped-query	8 (groups of 4)	4× smaller	≈ MHA
MLA — multi-head latent	a small latent vector	~GQA or better	≈ MHA

• 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

Quick check

0/2

Q1Your model has `d_model=768` and you want 12 heads. What's `d_head`?

7681264256

Check answer

Q2Why does grouped-query attention (GQA) save memory during inference?

It uses smaller embeddingsIt reuses the same K and V vectors across multiple query heads — the KV-cache shrinks by the grouping factorIt runs in fp16It only attends to recent tokens

Check answer