Many Heads.
One attention can only look in one way at a time. Split it into several narrow heads and each can look for something different.
A single attention produces, for each position, one weighted average of the values. If a token needs two things at once — the previous word for grammar and the subject of the sentence for meaning — one set of weights cannot point at both cleanly. Multi-head attention splits the width d into H heads of dimension d/H, gives each head its own query, key and value projections, runs them in parallel, concatenates their outputs and mixes them with one more matrix, W_o.
The cost is about the same as one wide head: the projections have the same total size, and the score computation is split, not multiplied. What changes is expressiveness: each head has its own attention pattern. The course model micro has four heads of 32 dimensions in each of its four layers. Production models have tens of heads per layer — Llama 3 8B has 32 query heads of 128 dimensions.
Below: the four heads of one of micro’s layers reading the same sentence. Choose the layer; compare the patterns. The readings give each head’s average look-back distance and how spread its attention is.
- Describe multi-head attention: split, attend in parallel, concatenate, project.
- Relate head count and head dimension to the model width.
- Compare the attention patterns of different heads.
- Explain why many heads cost no more than one.
Division of labour
micro’s second layer shows the point. Its fourth head is a near-pure previous-token head: it puts 60% of its weight on the token just before, looks back 1.0 tokens on average, and has an entropy of only 1.09 bits. Its other heads are less specialised but also local. The heads of the first layer read far more widely: head 1 there looks back 3.6 tokens on average, with an entropy of 2.40 bits. In the third layer, head 1 sends 58% of its weight to the very first token — a pattern discussed on the next plate. Four heads of 32 dimensions cost the same as one head of 128, but they can do four different jobs.
Nobody assigns those jobs; training finds them, and not every head finds a useful one — many heads in large models can be removed with little effect (plate IV.8). The number of heads is a design choice traded against head dimension at fixed width: more, narrower heads give more patterns, but each with less room for the information it carries. Grouped-query attention (plate VII.2) changes this trade by letting several query heads share one key and value head, which saves memory during generation.
Terms
- Multi-head attention
- Several attention heads with separate projections run in parallel, then concatenated and projected.
- Head
- One attention with its own queries, keys and values of dimension d/H.
- Output projection
- The matrix W_o that mixes the concatenated head outputs back into d dimensions.
- Previous-token head
- A head that attends almost entirely to the token just before.