Why use multiple attention heads instead of one large attention operation?

How to think about it

Multiple attention heads are useful because each head is an independently parameterized attention calculation, so a token can make several separate decisions about which other tokens matter. A single full-width head has one set of attention weights and one weighted mixture; standard multi-head attention gets several independent routing patterns without multiplying the leading parameter and attention-compute cost.

Why one head is a bottleneck

A token representation often needs more than one kind of context at once. In a sentence, it may need to find the noun that a pronoun refers to, the verb that governs it, and a nearby phrase that supplies local syntax. These relationships may point to different positions.

An attention head uses three vectors:

• A query is what the current token is looking for.
• A key is the matching signal that each source token offers.
• A value is the information copied from that source token if it is selected.

For a query at position i and a source token at position j, the head computes a similarity score:

score_ij = q_i · k_j / sqrt(d_k)

A softmax converts all scores for the query into weights that sum to one:

alpha_ij = softmax(score_i)_j

The output at position i is then a weighted average of the value vectors:

z_i = sum_j alpha_ij v_j

The important detail is that alpha_ij is one scalar shared by every coordinate of v_j. If a single head gives token A a weight of 0.7 and token B a weight of 0.3, every output channel receives a 0.7 and 0.3 mixture of those tokens.

That is not the same as saying one head can look at only one token. One head can spread its weight across many positions. The limitation is that all of its value channels share the same spread.

Multi-head attention changes this by creating several smaller heads:

head_r = Attention(Q W_Q,r, K W_K,r, V W_V,r)

Each head has its own query, key, and value projections. The usual design sets the per-head dimension to:

d_h = d_model / h

The head outputs are concatenated and mixed by an output projection:

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W_O

The separate softmax operations are the key. Head 1 can give most of its weight to one source token while head 2 gives most of its weight to another. The heads run in parallel; one head does not read another head’s output until concatenation and W_O.

A single full-width head has a larger dot product for each score. That gives it a richer scoring space, but still only one attention distribution per query. Multiple heads trade some width in each scoring space for several independently normalized distributions.

A concrete example

Consider the sentence:

The engineer handed the manager the report because she approved it.

At the token she, the model may need evidence about an antecedent and about the surrounding event. The correct interpretation is not guaranteed to be cleanly separated, but the example shows the mechanism.

Use a toy model with d_model = 4 and h = 2. Each head therefore has d_h = 2. Focus on four source positions: engineer, manager, report, and because.

Suppose training produces these illustrative attention rows for the query at she:

• Head 1: [0.85, 0.05, 0.05, 0.05]
• Head 2: [0.05, 0.10, 0.75, 0.10]

The first head can put most of its two output channels under evidence from engineer. The second can put most of its two output channels under evidence from report. After concatenation, the representation carries both mixtures before W_O remixes them.

A single full-width head could assign weights such as [0.45, 0.05, 0.45, 0.05]. That lets it use both engineer and report, but every value channel receives that same source mixture. A later linear projection can rearrange and combine those channels, but it cannot go back and request a second independently normalized attention pattern at that layer.

The numbers are illustrative, not a claim that every trained model develops a neat “pronoun head” and “report head.” Real heads often mix roles.

Why the cost is roughly the same

Suppose d_model = 512 and there are h = 8 heads. Each head has dimension d_h = 64.

For one full-width head, each of the query, key, and value projections is roughly a 512 × 512 matrix. For eight heads, each projection is eight 512 × 64 matrices. The total is the same:

8 × 512 × 64 = 512 × 512

The usual output projection is another 512 × 512 matrix. Ignoring biases, the standard attention block therefore has approximately:

3 × 512 × 512 + 512 × 512 = 1,048,576

projection weights whether it uses one full-width head or eight 64-dimensional heads.

The leading attention arithmetic also matches. For a sequence of 128 tokens, the query-key score products require roughly:

• One head: 128 × 128 × 512
• Eight heads: 8 × 128 × 128 × 64

Both equal 8,388,608 scalar multiply-accumulate positions for that product. The value aggregation has the same shape of comparison.

That does not mean multi-head attention has identical wall-clock cost. A naive implementation stores eight 128 × 128 attention maps rather than one, and multiple softmax operations add overhead. Fused kernels and memory-efficient attention reduce that cost substantially, but hardware efficiency still depends on head dimension, batch size, sequence length, and the implementation.

What this looks like in PyTorch

Here is a standard eight-head self-attention layer:

import torch

from torch import nn

mha = nn.MultiheadAttention(
    embed_dim=512,
    num_heads=8,
    batch_first=True,
)

x = torch.randn(2, 128, 512)

y, weights = mha(
    x,
    x,
    x,
    need_weights=True,
    average_attn_weights=False,
)

print(y.shape)       # torch.Size([2, 128, 512])
print(weights.shape) # torch.Size([2, 8, 128, 128])

The input has shape (batch, sequence, model_dimension). The output returns to the same model dimension after the heads are concatenated and projected. The attention weights have one 128 × 128 map for each of the eight heads.

A common debugging failure appears when someone expects eight maps but sees a tensor shaped like torch.Size([2, 128, 128]). PyTorch averages the head weights by default when weights are requested. That averaged map can look diffuse and hide the fact that two heads are focusing on different positions. Use average_attn_weights=False when inspecting individual heads. Do not mistake an averaged visualization for evidence that all heads behave identically.

Common misconception: Head 1 is not guaranteed to be “the syntax head,” and head 2 is not guaranteed to be “the coreference head.” Specialization can emerge, but heads may overlap, change behaviour across layers, or serve several purposes. Attention weights show routing, not a complete explanation of the model’s reasoning.

The senior-level nuance

More heads are not automatically better. With d_model fixed, increasing h makes each head narrower. A 512-dimensional model with 64 heads gives every head only eight dimensions. That may leave each head too little capacity to form a useful query-key similarity, while adding softmax and memory overhead.

The best head count depends on model width, task, sequence length, hardware, and the quality-latency trade-off. A small latency-sensitive model may get the same validation quality with fewer heads. The right answer is an experiment, not “use as many heads as possible.”

Heads are also not guaranteed to be indispensable. Pruning studies have found that many trained heads can sometimes be removed with little change to an aggregate benchmark score. That usually means the model learned redundant routes or had spare capacity. It does not invalidate the reason for multi-head attention: during training, the model had the option to create multiple independent routes. Before pruning, some heads may still matter for rare inputs, long contexts, or particular output classes.

The number of heads does not change the basic quadratic dependence on sequence length. A sequence of length n still creates attention interactions proportional to n²; the heads divide the feature dimension, not the token pairs.

There is also a serving trade-off in decoder-only language models. Standard multi-head attention stores separate key and value vectors for every head in the KV cache, which is the saved context used during autoregressive decoding. Multi-query attention shares keys and values across query heads, while grouped-query attention shares them across groups. These variants reduce KV-cache memory and memory bandwidth, often at some quality cost. They are useful production choices, but they do not provide the same independent key-value subspaces as ordinary multi-head attention.

What they’ll ask next

Can one head attend to multiple tokens?

Yes. Its softmax row can assign weight to many positions. The difference is that one head uses the same distribution for all of its value channels. Multiple heads provide several distributions, so different channel groups can retrieve different mixtures.

Does using eight heads multiply the parameter count by eight?

Not when d_model is fixed and the head dimension is reduced proportionally. Eight 512 × 64 projections contain the same number of weights as one 512 × 512 projection. If you keep the head dimension fixed while increasing the number of heads, then d_model grows and the parameter and compute costs do increase.

Do heads always correspond to interpretable linguistic relationships?

No. Some heads show clear local, positional, copying, or syntactic patterns, but interpretation is model- and layer-dependent. Test a head by masking or ablating it and measuring task performance, including important slices. A visually interesting attention map is not enough.

Say this in the interview

“Multi-head attention gives each query several independently learned attention distributions, so different subspaces can retrieve different relationships in parallel; because each head is narrower, the standard design keeps the leading parameter and attention-compute cost close to one full-width head.”