Lensa ML

How Multi-Head Attention Works

Step 1 of 6: One Head's Limit

A single head captures only one attention pattern per query

single head: same attention scores as self-attentionThecatsatonthematcat0.200.220.160.100.190.13Thecatsatonthemat

Query

cat

Top key

cat

Entropy (bits)

2.54

"cat" puts its largest weight on itself, only 22%, and spreads the rest almost evenly (entropy 2.54 bits, near the 2.58-bit maximum for 6 words). One head has to fold every relationship into this single blurred pattern.

Click each word to change the query and watch the Entropy readout. Entropy H = −Σ p_i log₂ p_i measures how spread out a row is: it is 0 bits when all weight sits on one word and log₂ 6 ≈ 2.58 bits when the 6 weights are equal. Every query here lands between 2.53 and 2.58 bits, so this head barely focuses at all. The deeper limit of one head is that its output is a single weighted average. If a word needs its subject from one position and its object from another, one softmax must split its weight between them, and the two get blended into one vector.

Step 2 of 6: Multiple Heads

Split the model dimension across parallel attention heads

d_model = 64H1d_k=16H2d_k=16H3d_k=16H4d_k=16

d_model

64

Heads

4

d_k

16

Params (4·d_model²)

16,384

4 heads each get d_k = 16 dimensions. Each head can learn a different attention pattern in its own subspace. The parameter count stays at 4·d_model² for any number of heads: more heads means narrower heads, not a bigger layer.

Click 1, 2, 4, and 8 heads and watch d_k shrink. Here d_model = 64, so d_k = 64 / h. The original Transformer used d_model = 512 with 8 heads of 64, and GPT-3 uses 12,288 with 96 heads of 128. Because each head is smaller, total compute and parameter count stay about the same as for one full-width head. In code there are no separate head modules: one projection produces a 64-wide Q, which is reshaped to (h, length, d_k) so all heads run in a single batched matrix multiply.

Step 3 of 6: Different Patterns

Each head learns a distinct type of relationship

ThecatsatonthematThe0.550.270.060.040.050.03cat0.210.470.190.050.040.04sat0.040.210.520.160.040.03on0.030.040.160.570.170.03the0.030.040.040.190.480.21mat0.030.040.040.040.250.60

Head

H1

Pattern

Positional

Top key

cat

Entropy

2.00

Head 1 (Positional): "cat" puts 88% of its attention on itself and its immediate neighbours. This captures local context and word order. These eight patterns are hand-set to show typical behaviours; in a real model each comes from that head's own W_Q and W_K, and many heads are less tidy.

Click each head label, then click words to change the query. These eight patterns are hand-built to illustrate common head types. In trained models, roles are found only by inspection, and many heads turn out to be redundant. Michel et al. (2019) removed a large share of the heads from trained models at test time with little loss in accuracy, and Voita et al. (2019) found that a few specialized heads do most of the work. Removing heads is one way to make a model smaller and faster.

Step 4 of 6: Concatenation

Stack all head outputs back together

H1:16H2:16H3:16H4:16Concat → [64]W_O → [64]Output: d_model = 64

Per head

16

Concat

64

Output

64

4 heads of size 16 concatenate to 64 = d_model. The linear projection W_O combines all head outputs back to 64 dimensions.

Click 1, 2, 4, and 8 heads and check that the concatenated width always equals d_model = 64. Concatenation only places the head outputs side by side; nothing is learned at this step. With 4 heads, dimensions 0 to 15 of the result come from head 1, 16 to 31 from head 2, and so on. In code it reverses the earlier reshape: move the head axis back and flatten it. The width has to stay d_model because the residual connection adds this output to the block's input.

Step 5 of 6: Output Projection

W_O mixes head outputs into the final representation

Concatenated headsH1:16H2:16H3:16H4:16W_O[64 x 64]Output [64]Head 1 share of output (illustrative): 35%

Concat

64

W_O shape

64x64

Output

64

W_O is a learned linear projection that mixes all head outputs. In this simplified picture Head 1 (Positional) contributes 35%, and its local patterns are blended with the other heads' signals.

Click each head button to see its share. The percentages are a simplified picture. W_O is a full d_model × d_model matrix, 64 × 64 = 4,096 weights here, and every output dimension mixes every dimension of every head. Splitting W_O into one block of rows per head shows that multi-head attention equals the sum of each head's output times its own slice of W_O. That is why heads can be studied or pruned one at a time. Together, W_Q, W_K, W_V, and W_O hold 4 × d_model² weights per layer.

Step 6 of 6: Multi-Head in Action

Compare attention patterns across heads for a full sentence

H1: PositionalH2: SyntacticH3: GlobalH4: Backward

Head

H1

Query

cat

Top key

cat

Weight

0.47

Comparing all heads for "cat": Head 1 captures positional proximity, while other heads capture different relationships. Together they give the model several views of the same sentence.

Click the word buttons and head buttons to compare all four heads for one query. Heads also work together across layers. A head in a later layer attends over vectors that already mix in what earlier heads gathered, so two heads in sequence can perform a lookup neither could do alone. A well-studied example is the induction head: one head marks the token that followed an earlier copy of the current word, and a later head copies it forward. Olsson et al. (2022) linked these heads to in-context learning in language models.

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

References

  1. Vaswani et al. (2017), "Attention Is All You Need", NeurIPS
  2. Michel, Levy & Neubig (2019), "Are Sixteen Heads Really Better than One?", NeurIPS
  3. Voita et al. (2019), "Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned", ACL
  4. Olsson et al. (2022), "In-context Learning and Induction Heads"

Last updated