ML//transformer//attention//multi-head attention

Multi-head attention is the attention layer of the transformer that runs self-attention several times in parallel with different learned projections, each head operating on a reduced subspace, so that the model can track different types of relationship between tokens at once. The outputs of all heads are concatenated and projected back to the model dimension through \(W_O\); GPT-3 has 96 heads per layer, each of 128 dimensions.

Multi-head attention is the attention layer of the transformer that runs self-attention several times in parallel with different learned projections, each head operating on a reduced subspace, so that the model can track different types of relationship between tokens at once. The outputs of all heads are concatenated and projected back to the model dimension through WOW_OWO​; GPT-3 has 96 heads per layer, each of 128 dimensions.

The dimension is reduced so that each head can specialize (syntax, coreference, positional patterns, semantic similarity, induction or pattern copying), with different types of relationships in different subspaces.

No explicit mechanism forces this: specialization emerges from training. Two heads learning the same thing do not reduce the loss, so backprop pushes them to diversify. BERT-base, the textbook case, has 12 heads of 64 dimensions per layer.

No head is given a grammatical role. The subject-verb or verb-object heads reported in studies are descriptions found after training, many heads do nothing a person can name, and removing a sizeable share of them often costs little, so calling one head the syntax head claims more than is known (mechanistic interpretability).

GQA shares KKK and VVV across groups of heads to shrink the KV cache, the modern optimization.

Heads can form cross-layer circuits: induction heads are two heads in consecutive layers that coordinate. One finds matching prefixes, the other copies what followed.