6. Multi-head attention

A single attention head produces, per position, one set of weights summing to 1 — one weighted average. But a token in a real sentence needs several things at once. In "The keys to the cabinet are on the table", the position of are needs the subject keys for agreement, and separately the syntactic structure, and separately the topic. One averaging operation cannot deliver all of that without blurring them into mush.

The fix is straightforward: run several heads in parallel, each with its own W_Q, W_K, W_V, so each learns to attend on a different criterion.

Crucially, the heads do not each get the full width. d_model is divided among them:

d_head = d_model / n_heads

With d_model = 4096 and 32 heads, each head works in 128 dimensions. Its Q, K and V projections are [128, 4096], and it produces a 128-dimensional output per position. In code you will rarely see those 32 small matrices: implementations fuse them into one [4096, 4096] projection and slice the result into heads afterwards, which is the form the calculator below prints.

Then:

1. Concatenate all 32 head outputs → back to 4096 dimensions.
2. Multiply by an output projection W_O of shape [4096, 4096], which lets the heads' results mix and be rescaled before rejoining the residual stream.

The total compute is about the same as one full-width head, which is the elegant part: you get many relationship types for free, paid for by giving each less room. That trade-off has a limit — heads that are too narrow cannot represent anything useful — which is why head count scales with model width rather than growing indefinitely.

Heads are not assigned roles; specialisation emerges from training. Interpretability work has identified recurring types:

• Positional heads — attend to the previous token, or to a fixed offset.
• Syntactic heads — verbs attending to their subjects, adjectives to their nouns, closing brackets to their opening partners.
• Coreference heads — pronouns attending to the entity they refer to.
• Induction heads — the copy-the-pattern behaviour introduced in Lesson 5. These reliably appear during training at the same moment as a sharp jump in in-context learning ability.
• Duplicate-token heads, name-mover heads, and other narrow circuits found by tracing specific tasks.

Study the per-head patterns below on the same sentence and notice how different the same layer's heads look.

At generation time you cache the keys and values of every past token so you do not recompute them (Lesson 11). That cache is proportional to n_heads × d_head × n_layers × sequence_length, and it becomes the dominant memory cost of serving a model with long context.

Two variants attack it:

• Multi-query attention (MQA) — keep all the query heads, but use a single shared key head and value head. The cache shrinks by a factor of n_heads. Quality degrades a little.
• Grouped-query attention (GQA) — the middle ground now used almost everywhere. Group the query heads (say 32 queries into 8 groups); each group shares one key/value pair. A 4× or 8× cache reduction with quality close to full multi-head.

This is a good example of a pattern you will see repeatedly: an architecture choice driven not by accuracy but by what it costs to serve the model.