Multi-Head Attention

Multi-head attention runs several attention functions in parallel, each with its own Q, K, V projections, and concatenates their outputs. Each head works in a smaller subspace of dimension d_model / n_heads, so total parameters stay the same while the layer can express more.

The motivation is what one head can't do. A single self-attention head computes one attention matrix, which captures one kind of relationship, typically the one that lowers loss on the training signal. If your data has subject-verb agreement, co-reference, long-range discourse, and syntactic chunking tangled together, a single head smears them into one softmax distribution and loses half the signal.

The 2017 Vaswani paper introduced the fix. In 2026 it is the default in transformers, and the live questions are how many heads to use and whether keys and values share projections (Grouped-Query Attention, Multi-Query Attention, Multi-head Latent Attention).

Split. Take X of shape (N, d_model). Project to Q, K, V, each of shape (N, d_model). Reshape to (N, n_heads, d_head) where d_head = d_model / n_heads, then transpose to (n_heads, N, d_head).

Attend in parallel. Run scaled dot-product attention inside each head. Each head produces (N, d_head). The heads work on different subspaces of the embedding and don't exchange information during the attention computation itself.

Concatenate and project. Stack the heads back to (N, d_model) and multiply by a learned output matrix W_o of shape (d_model, d_model). W_o is where the heads mix.

Why it works. Each head can specialize without competing with the others for representational budget. Probing studies from 2019–2024 found distinct head roles: positional heads, heads that attend to the previous token, copy heads, named-entity heads, and induction heads (which underlie in-context learning).

The 2026 lineage of variations:

GQA is the modern default because it cuts KV-cache memory by a factor of N/G while keeping nearly full quality. MLA goes further: it compresses K/V into a latent space and projects back at compute time. That costs extra FLOPs and saves considerably more memory.

Step 1: split heads from the single-head attention we already have

Start from the standalone scaled_dot_product_attention function from Lesson 02 and wrap it with a split/combine pair. code/main.py is pure stdlib: a small Matrix class (row-major floats) carries matmul, transpose, and softmax_rows, so there is no numpy. Splitting copies each head's contiguous block of columns into its own Matrix:

def split_heads(X: Matrix, n_heads: int) -> List[Matrix]:
    d_head = X.cols // n_heads
    heads = []
    for h in range(n_heads):
        H = Matrix(X.rows, d_head)
        for i in range(X.rows):
            for j in range(d_head):
                H.set(i, j, X.get(i, h * d_head + j))
        heads.append(H)
    return heads
 
def combine_heads(heads: List[Matrix]) -> Matrix:
    n, d_head = heads[0].rows, heads[0].cols
    out = Matrix(n, d_head * len(heads))
    for h, H in enumerate(heads):
        for i in range(n):
            for j in range(d_head):
                out.set(i, h * d_head + j, H.get(i, j))
    return out

A head is a column slice of width d_head, and combining writes the slices back side by side. These are nested loops over a List[Matrix]. NumPy or PyTorch does the same split with one reshape and one transpose, which is the shape bookkeeping The Concept describes.

Step 2: run scaled-dot-product attention per head

Each head gets its own slice of Q, K, V, so attention runs once per head in a plain Python loop:

def multi_head_attention(X, Wq, Wk, Wv, Wo, n_heads):
    Q = matmul(X, Wq)
    K = matmul(X, Wk)
    V = matmul(X, Wv)
    Qh = split_heads(Q, n_heads)
    Kh = split_heads(K, n_heads)
    Vh = split_heads(V, n_heads)
    head_outs, per_head_weights = [], []
    for q, k, v in zip(Qh, Kh, Vh):
        o, w = scaled_dot_product_attention(q, k, v)
        head_outs.append(o)
        per_head_weights.append(w)
    concat = combine_heads(head_outs)
    return matmul(concat, Wo), per_head_weights

In main.py the heads run one after another. On real hardware the same computation is one batched matmul of shape (heads, N, d_head) × (heads, d_head, N) -> (heads, N, N), so adding heads doesn't add a matmul. The demo uses n_heads=2 and d_model=8, so d_head=4.

Step 3: Grouped-Query Attention variant

Only the key and value projections change. Q keeps n_heads groups. K and V get n_kv_heads < n_heads groups, which are repeated to match Q by indexing into the smaller list:

Qh = split_heads(Q, n_heads)
Kh_small = split_heads(K, n_kv_heads)
Vh_small = split_heads(V, n_kv_heads)
repeat = n_heads // n_kv_heads
Kh = [Kh_small[i // repeat] for i in range(n_heads)]
Vh = [Vh_small[i // repeat] for i in range(n_heads)]

At inference this saves memory because only n_kv_heads copies live in the KV cache, not n_heads. Llama 3 70B uses 64 query heads with 8 KV heads, an 8× cache shrink. The GQA demo in main.py uses 4 query heads and 2 KV heads and prints the KV-cache element counts for both layouts, a 2× shrink.

Step 4: probe what each head learned

Run MHA on a short sentence and print the (N, N) attention matrix for each head. The main.py demo does this for the cat sat on the mat with 2 heads at d_model=8. The GQA demo prints no attention matrices, only the output shape and cache sizes. Different heads pick out different structure even with random initialization. Part of that is signal, and part is rotational symmetry in the subspaces.

In PyTorch, MHA is one line:

import torch.nn as nn
 
mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)

GQA as of PyTorch 2.5+:

from torch.nn.functional import scaled_dot_product_attention
 
# scaled_dot_product_attention auto-dispatches Flash Attention on CUDA.
# For GQA, pass Q of shape (B, n_heads, N, d_head) and K,V of shape
# (B, n_kv_heads, N, d_head). PyTorch handles the repeat.
out = scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

How many heads? Rules of thumb from production models in 2026:

In the table, d_head lands at 64 or 128. It sets how much one head can "see." Drop below 32 and heads start fighting the scaling factor sqrt(d_head). Go above 256 and you lose the "many small specialists" benefit.

outputs/skill-mha-configurator.md is a skill that recommends head count, kv-head count, and projection strategy for a new transformer. Give it your parameter budget, sequence length, and deployment target.

1. Easy. code/main.py has no training loop, so this one is an extension you write. Build a tiny one-layer model around the MHA in main.py, train it on a synthetic copy task, and sweep n_heads from 1 to 16 with d_model=64 fixed. Plot the loss. Do more heads help, plateau, or hurt?
2. Medium. Implement MQA (one KV head shared across all query heads). Measure how much parameter count drops vs full MHA. Compute how much the KV-cache size shrinks at inference for N=2048.
3. Hard. Implement a tiny version of Multi-head Latent Attention: compress K,V to a rank-r latent, store the latent in the KV cache, decompress at attention time. At what r does cache memory cross below 1/8 of full MHA while quality stays within 1 bit of validation ppl?

• Vaswani et al. (2017). Attention Is All You Need §3.2.2: the original multi-head spec.
• Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need: the MQA paper.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints: how to convert MHA to GQA after training.
• DeepSeek-AI (2024). DeepSeek-V2 Technical Report: MLA and why it beats MHA/GQA on cache memory.
• Olsson et al. (2022). In-context Learning and Induction Heads: mechanistic look at what heads actually do.