What is multi-head attention?
Multi-head attention runs several self-attention operations, called heads, in parallel on the same input and combines their results. Each head has its own query, key, and value projections into a smaller subspace, so each can learn a different pattern of which words to look at. The head outputs are concatenated and passed through a final linear projection, W_O, to produce one vector per token.
Every Transformer layer uses multi-head attention, from the original translation model to GPT, BERT, Llama, and vision transformers. A single head can only form one weighted average per query, so it blurs together different kinds of context. Splitting the model width into heads lets one layer track position, syntax, and long-range links at the same time for about the same compute as one wide head.
Key ideas
- Multi-head attention runs h attention heads in parallel on the same input.
- Each head works in a subspace of size d_k = d_model / h.
- Head outputs are concatenated back to width d_model.
- The output projection W_O mixes information from all heads.
- The total cost is about the same as one attention head of full width.
- Many heads in trained models are redundant and can be pruned.
Common questions
How many attention heads should a Transformer have?
Most models choose a head size of 64 or 128 and set the number of heads to d_model divided by that size. That gives 12 heads in BERT-base (768 wide) and 32 heads in a 4,096-wide model. Very small head sizes hurt, because each head's dot products become too low-rank to separate keys well.
Do more heads always make attention better?
No. With d_model fixed, adding heads makes each one narrower, and past a point quality drops. Experiments in the original Transformer paper found 8 heads better than 1, but 16 and 32 slightly worse than 8 at the same width.
What is the output shape of multi-head attention?
The same as its input: one vector of width d_model for each token, or (batch, length, d_model) for a batch. Inside, the tensors take the shape (batch, heads, length, d_k), and each head's weight matrix is (length × length).
Is multi-head attention the same as stacking attention layers?
No. Heads run side by side within one layer and all read the same input, while layers run one after another and each reads the previous layer's output. Models use both: for example, 12 layers with 12 heads each in BERT-base.
Related explainers
- How Self-Attention Works — Query, key, value — the core of self-attention
- The Transformer Architecture — A visual guide to 'Attention Is All You Need'
- Cross-Attention — Attending across sequences and modalities
- Positional Encoding — How transformers know word order
- Decoder-Only Transformers — The architecture behind GPT and modern LLMs
- BERT & Masked Language Modeling — Mask a word, predict it from both sides — how BERT learns deep bidirectional representations
References
- Vaswani et al. (2017), "Attention Is All You Need", NeurIPS
- Michel, Levy & Neubig (2019), "Are Sixteen Heads Really Better than One?", NeurIPS
- Voita et al. (2019), "Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned", ACL
- Olsson et al. (2022), "In-context Learning and Induction Heads"
Last updated
Spotlight