What is multi-head attention and why use multiple heads instead of one?

How to think about it

Multi-head attention runs several scaled dot-product attention operations in parallel, each with its own learned projections of queries, keys, and values. Multiple heads let the model make several independent routing decisions at once, so different feature slices can focus on different tokens or relationships without multiplying the total model width.

The mechanism

Start with ordinary attention. Suppose a sequence has n tokens, and each token is represented by a vector of width d_model.

A query is what the current token is looking for. A key describes what each candidate token offers for matching. A value is the information actually retrieved after a match. The names are slightly theatrical, but the mechanism is straightforward: compare a query with every key, turn those scores into weights, and use the weights to mix the values.

The core operation is:

Attention(Q, K, V) = softmax((QK^T) / sqrt(d_k)) V

QK^T produces one compatibility score for every query-key pair. The softmax turns the scores for each query into weights that sum to one. The result is therefore a weighted mixture of value vectors. A token can attend mostly to one other token, or spread its attention across several.

The division by sqrt(d_k) is not decoration. If the components of the query and key vectors have roughly unit variance, their dot product tends to grow in magnitude as d_k grows. Without scaling, the softmax can become nearly one-hot: one score wins decisively, the others receive almost no weight, and useful gradients become harder to learn. Scaling keeps the logits in a more workable range.

For self-attention, the input sequence is projected into all three roles. For head i, the operation is conceptually:

head_i = Attention(XW_Q_i, XW_K_i, XW_V_i)

Multi-head attention then concatenates the head outputs and applies one final learned projection:

MHA(X) = Concat(head_1, ..., head_h) W_O

The heads do not each see a separate copy of the task. They all start from the same input, but their projection matrices learn different feature spaces. One head may compare grammatical subjects and verbs. Another may connect a pronoun to an earlier noun. A third may preserve local neighbouring information. Those roles are tendencies, not labels stamped onto the model; learned heads are often messy and input-dependent.

In a decoder, a causal mask is applied before softmax so a token cannot attend to future tokens. Multi-head attention changes how many attention patterns are computed, not that masking rule.

A concrete example

Take a small encoder processing:

The bank raised rates

Imagine looking at the query for bank. Two of the model’s heads might produce these illustrative attention weights:

Token receiving weight	The	bank	raised	rates
Head 1	0.05	0.10	0.10	0.75
Head 2	0.05	0.85	0.05	0.05

These are not claimed outputs from a particular trained model. They show the useful possibility. One head routes information from rates, while another keeps a strong local representation of bank. Each head produces its own value mixture. The model then concatenates those mixtures, preserving both routes until the output projection combines them.

A single head has only one attention distribution for that query. It could place some weight on both bank and rates, and its value vectors could be engineered to carry several facts at once. So one head is not incapable of solving the task. The limitation is that every output channel must use the same routing weights and the same projected value space. Multiple heads provide several independent weighted mixtures before information is combined.

Now use realistic transformer dimensions. Let:

• d_model = 512
• h = 8 heads
• d_head = 512 / 8 = 64
• sequence length n = 128

Each head produces queries, keys, and values with shape 128 × 64. Its score matrix has shape 128 × 128. After eight heads are computed, their outputs are concatenated back to 128 × 512.

This keeps the main parameter and arithmetic budget close to a single head of width 512. Across all eight heads, the query projection contains:

8 × 512 × 64 = 262,144 weights

The key and value projections contain the same number each. The output projection contains 512 × 512 = 262,144 weights. Ignoring biases, the four projections total 1,048,576 weights. A single-head design using width 512 has the same four-matrix total.

The attention arithmetic has the same property. Eight heads perform:

8 × 128 × 128 × 64 = 8,388,608

scalar query-key products. One head with width 512 performs:

128 × 128 × 512 = 8,388,608

The work is divided into smaller independent spaces rather than simply multiplied by eight. Real latency is not identical, because kernels, memory movement, softmax operations, and hardware utilisation matter. The arithmetic comparison explains the design; profiling decides the winner.

Why one wide head is not the same thing

A single wide head has more dimensions available inside one value mixture, but it still produces one score distribution per query. That is the bottleneck multiple heads address.

Think of a query token as needing two separate answers: “Which earlier token refers to me?” and “Which token tells me what semantic topic matters here?” One attention distribution has to compromise between those goals. Multiple heads can make both selections independently, then preserve both results in separate slices of the output.

This is why the phrase “different representation subspaces” matters. A subspace is simply a learned set of directions in the model’s vector space. One head may make similarity depend strongly on number or tense. Another may project the same token vectors so that entity identity matters more. Their keys and queries are not comparing exactly the same features.

Common misconception: multi-head attention does not guarantee one interpretable linguistic job per head. Researchers often find heads with useful patterns, but heads can be redundant, split a pattern across layers, or change behaviour with the input. Treat “this is the coreference head” as a hypothesis to test with ablations, not as an architectural promise.

Multiple heads also do not improve the asymptotic complexity of dense attention. The attention part remains roughly O(n² d_model), and the projection part remains roughly O(n d_model²). If a 100,000-token context is too expensive, increasing the head count will not rescue it. The quadratic sequence interaction is still there.

The production trade-off

More heads are not automatically better. With d_model fixed, increasing the head count makes each head narrower. In the example, 32 heads would give d_head = 16. That creates more independent attention distributions, but each distribution has less room to represent useful features. Some heads may become redundant, and extra reshaping or kernel overhead can hurt latency.

A single head can be a sensible choice for a small model, a simple task, or a latency-constrained system if validation confirms that quality is sufficient. The textbook answer is not “always maximise the number of heads.” The practical choice depends on model width, sequence length, hardware, quality targets, and whether inference memory is the real bottleneck.

That last point matters in autoregressive generation. Standard multi-head attention caches keys and values for every head. Consider one sequence of length 2,048, with 32 layers, 32 heads, head width 128, and two-byte floating-point cache entries:

2 tensors × 2,048 tokens × 32 heads × 128 values × 32 layers × 2 bytes
= 1,073,741,824 bytes

That is about 1 GiB for one sequence’s key-value cache. Multi-query attention keeps multiple query heads but shares one key-value head. Grouped-query attention shares key-value heads across smaller groups of query heads. With eight key-value heads instead of 32, the same example uses roughly 256 MiB. The trade-off is lower cache cost and often better serving throughput, in exchange for less independent key-value capacity.

A common configuration failure is changing h without checking divisibility. With d_model = 512, seven equal-width heads cannot be formed. Common transformer libraries reject the model during construction rather than silently inventing uneven heads. Even when the division works, changing head count should be treated as a model change and evaluated, not as a free performance knob.

What they’ll ask next

Does multi-head attention use eight times as much computation when there are eight heads?
Not when the model width is held fixed and the head width is reduced proportionally. The main dot-product arithmetic is roughly unchanged. Memory use, kernel overhead, and actual latency can still change, and standard multi-head key-value caches are larger than grouped or shared variants.

Why divide by the square root of the key dimension?
Dot products tend to grow in magnitude as the key dimension grows. Dividing by sqrt(d_k) keeps the softmax from becoming excessively sharp, which preserves useful gradients during training.

Are multiple heads always more expressive than one head?
They provide more independent attention distributions and usually make the architecture easier to train for varied relationships, but one head can solve some tasks. Head count is a capacity and optimisation choice, not a guarantee of better accuracy. Compare configurations on held-out data and measure serving costs.

Say this in the interview

“Multi-head attention runs several scaled dot-product attention operations on different learned projections, then concatenates their outputs; multiple heads let the model route different feature subspaces to different positions at once, while using smaller per-head dimensions keeps the total projection and attention arithmetic roughly comparable to one wide head.”