What is multi-head attention in a transformer?

In this article

1. Before you start
2. Why interviewers ask this
3. Multi-head attention from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one step of the transformers and attention learning path.

Why interviewers ask this

Multi-head attention is a core part of the transformer that interviewers probe with questions such as: why several heads, what each one learns, what the shapes are, and why modern models share keys and values across heads. A candidate who can only say "multiple heads look at different things" has memorised a slogan.

A weak answer stops there. A strong answer gives the shapes (width split into heads, scores per head, outputs concatenated and mixed), says the parameters and the compute are about the same after splitting, and then names the inference problem that led to multi-query, grouped-query and latent attention: the size of the key-value cache. The worked example computes what grouping does to the projections, to the attention arithmetic and to the bytes read for each decoded token.

Multi-head attention from first principles

Multi-head attention

Multi-head attention is an attention layer that splits its projections into several heads, lets each head compute attention separately over its own slice, and concatenates the heads' outputs before a final linear mixing step.

Attention head

An attention head is a parallel attention computation inside a multi-head layer that has its own query, key and value projections and produces one output vector per token.

In plain words

Imagine a committee reading one sentence. Each member looks for something different. One tracks which noun a pronoun refers to, another tracks nearby words, another tracks the verb. Each reports a summary, and a chair combines the reports. Multi-head attention works the same way: each head reads the same tokens through its own learned lens, and a final matrix combines the reports. The analogy breaks in one place: nobody assigns the lenses in advance; training discovers whatever split of work reduces the loss.

In one precise sentence: the layer cuts the model width into h slices, runs ordinary scaled attention in each slice, glues the h results back together and multiplies by one more matrix.

Step 1: Each token becomes a query, a key and a value

Every token's vector is multiplied by three learned matrices to give a query, a key and a value. The query says what the token is looking for, the key says what it offers, and the value is the information handed on when a key matches. The layer in Vaswani et al. (2017) does this with separate matrices per head, which is the same as one large matrix whose output is cut into slices.

Input: the token vectors. Output: three vectors per token. Nothing in this step compares tokens; it only prepares the three roles each token can play. In a model with 64 heads the same step simply produces 64 sets of the three vectors, one per head, each of the smaller head size, so the layer stays one matrix multiplication per role.

Step 2: One head scores every pair and mixes the values

Inside a head, the dot product of one token's query with every key gives a score for each token. The scores are divided by the square root of the head size, then passed through a function called softmax, which turns scores into positive weights that sum to one. The weights then mix the value vectors. In a toy example with three tokens and a head of size 2, the first token's weights for head 0 are 0.529, 0.287 and 0.184, which sum to one (illustrative: random weights, seed 1).¹ The output for that token is the weighted mix of the three value vectors. This is single-head attention, the building block of everything below.

Step 3: Split the width into heads

With model width d and h heads, each head works on a slice of d divided by h dimensions. The Transformer paper uses d = 512 and h = 8, so each head has 64 dimensions (Vaswani et al., 2017, section 3.2.2). In the toy example d = 4 and h = 2, so each head has 2.

Each head still reads the whole token vector through its own projection; only its query, key and value are d divided by h wide. Each head produces its own weights, and they differ: for the first token head 0 puts 0.529 on token one while head 1 spreads weight almost evenly, 0.339, 0.313 and 0.348 (illustrative).¹ Different heads can therefore weigh the same tokens differently, which is the whole point of having several.

Text description of the diagram

• For the first token head 0 puts 0.529 on token one.
• Head 1 spreads weight almost evenly, 0.339, 0.313 and 0.348.

Step 4: Splitting into heads keeps parameters and compute about the same

Eight heads of 64 dimensions use the same parameters as one head of 512 dimensions. The four projection matrices (query, key, value and output) have 512 times 512 entries each whichever way the width is sliced: 1,048,576 parameters in total for both layouts.¹ The paper says the compute is close too: "Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality" (Vaswani et al., 2017).

For a sequence of n tokens, the score and value-mix steps cost 2 times h times n squared times d over h multiply-adds (a multiplication followed by an addition), which is 2 times n squared times d for any h. At n = 128 and d = 512 that is 16,777,216 multiply-adds for 8 heads and for 1.¹

Two things are not equal. The h heads together hold h times as many scores: at n = 128 that is 131,072 for 8 heads against 16,384 for one, and softmax runs over h times as many numbers.¹ Heads add score storage, not multiply-adds.

Step 5: Concatenate the heads and mix them

The h head outputs, each of size d over h, are glued side by side into one vector of size d and multiplied by an output matrix, usually written W_O. In the toy example the glued result has shape 3 tokens by 4 dimensions, and so does the layer's output, the same shape as the input, so layers can be stacked.¹

W_O maps the joined vector, h times the value size wide, back to the model width and lets every head write into every dimension of the output. Without it, each slice would carry one head's result until the feed-forward layer that follows mixes them. Each head's value matrix and its block of W_O can also be multiplied into one matrix, so W_O adds no new kind of computation.

1 · Q, K, V2 · One head3 · h heads4 · Parameters5 · JoinReplay

Step-by-step description

1. Each token gets a query, key and value.
2. A head scores, applies softmax, then mixes values.
3. The other heads run in parallel on their own slices.
4. The projections hold 1,048,576 parameters either way.
5. Every head feeds the concatenation and W_O.

• Every token's vector is multiplied by three learned matrices to give a query, a key and a value.
• Each head works on a slice of d divided by h dimensions.
• The four projection matrices have 1,048,576 parameters in total for both the 8-head layer and a single full-width head.
• The head outputs are concatenated and multiplied by an output matrix.

Worked example

When a model decodes one new token, the new token's query attends to every earlier key and value, which come from the KV cache. The sizes of a real model show what sharing keys and values changes in that step. Llama 3 70B has a width of 8,192, 80 layers, 64 attention heads and 8 key-value heads, so each head has 8,192 divided by 64, which is 128 dimensions (config.json, as of October 2026). The numbers below are per layer, for one decoded token at an 8,192-token context.

First the projections. The query matrix has 67,108,864 parameters and the output matrix has 67,108,864. With 8 key-value heads the key matrix and the value matrix have 8,388,608 each, so the four total 150,994,944, against 268,435,456 if all 64 heads had their own keys and values.¹ Each parameter is used once per token, so the projection multiply-adds fall by the same amount.

Then the attention itself. The score and value-mix multiply-adds are 2 times 64 query heads times 8,192 tokens times 128, which is 134,217,728 in every variant, because they depend on the 64 query heads and not on how many key-value heads exist.¹ Grouped-query attention therefore does not change the attention floating-point operation (FLOP) count.

What falls is memory traffic. Keys and values read from the cache for this token, in 16-bit numbers, are 268,435,456 bytes for multi-head attention, 33,554,432 for grouped-query attention with 8 heads and 4,194,304 for multi-query attention.¹

Decoding is limited by how fast bytes move (Shazeer, 2019), which is why the bytes column matters more than the arithmetic columns. The whole cache for this model is 320 KiB (1 KiB is 1,024 bytes) per token with grouped-query attention, which is 327,680 bytes;¹ the KV cache page works out that sizing.

Text description of the diagram

• The score and value-mix multiply-adds stay at 134,217,728 per decoded token per layer.
• Key-value bytes read go 268,435,456 for multi-head attention, 33,554,432 for grouped-query attention and 4,194,304 for multi-query attention.

Trade-offs and edge cases

The trade-off is memory and speed at inference against modelling capacity. The query heads still number 64 in the grouped variant; only the stored keys and values are shared, which is why the attention arithmetic does not shrink.

Multi-query attention. Shazeer (2019) shares one key and value head across all query heads and reports that incremental decoding becomes much faster, with minor quality loss, because decoding is limited by the memory bandwidth spent loading large key and value tensors.

Grouped-query attention. Ainslie et al. (2023) use an intermediate number of key-value heads, more than one but fewer than the query heads, and report quality that approaches multi-head attention at nearly the speed of multi-query; they also show how to convert existing multi-head checkpoints with about 5% of the original pre-training compute.

Multi-head latent attention. DeepSeek-V2 compresses the key-value cache into a latent vector and reports a 93.3% smaller cache than their earlier 67B model, a different route to the same goal of a small cache.

Masking. In a language model each token may attend only to itself and earlier tokens, so a causal mask sets scores for later positions to minus infinity before softmax. In the toy example, for head 0, the first token then has weights 1, 0, 0 and the second 0.4236 and 0.5764, with the last row unchanged because it already sees everything before it.¹ The mask is also why the cache works: see the key-value (KV) cache page for the full sizing story and FlashAttention for how attention is computed efficiently.

For serving context, the deep learning hub gives the wider map, the large language model (LLM) inference hub puts these choices in order, and low-rank adaptation (LoRA) fine-tuning adapts these same projection matrices cheaply.

How to answer it in an interview

The spoken structure: say the split and the join, say the parameters are unchanged and the compute about the same, then name the cache problem and its fix.

Common mistakes

• "More heads means more parameters." (wrong) At a fixed model width, splitting the width keeps the projection parameters the same. (right) "Eight heads of 64 use the same parameters as one head of 512."
• "Each head is told what to look for." (wrong) The roles emerge in training. (right) "Heads specialise because the loss rewards it, not by design."
• "Grouped-query attention reduces the number of query heads." (wrong) Only the key and value heads are shared. (right) "Query heads stay, key and value heads are fewer."
• "The causal mask is applied after softmax." (wrong) Masking after softmax would leave rows that do not sum to one. (right) "I add minus infinity to masked scores before softmax."

Follow-up questions

What is grouped query attention?

Grouped-query attention lets several query heads share one key head and one value head. Llama 3 70B has 64 query heads and 8 key-value heads, so each key-value pair serves 8 query heads, and the cache is one eighth of full multi-head attention.¹

What is masked multi-head attention?

Masked multi-head attention is multi-head attention with a causal mask that sets scores for later positions to minus infinity before softmax. A model can then train on every position of a sequence at once while each token sees only earlier tokens, exactly as when it generates text one token at a time. See self-attention for the single-head version.

What is the difference between single-head and multi-head attention?

Single-head attention computes one set of weights per token over the full width. Multi-head attention computes several sets over slices of the width, then joins them. The cost is about the same, but the model can attend in several ways at once.

What is grouped-query attention versus multi-query attention?

Multi-query attention shares one key-value head across all query heads, giving the smallest cache, 40 KiB per token for a 1-key-value-head version of Llama 3 70B. Grouped-query attention keeps a few, 8 here, which costs 320 KiB per token and keeps quality closer to full multi-head attention.¹

Does grouped-query attention change the FLOPs?

Grouped-query attention leaves the attention FLOPs unchanged and cuts memory traffic. For Llama 3 70B at an 8,192-token context, the score and value-mix multiply-adds stay at 134,217,728 per decoded token per layer, because they depend on the 64 query heads. Only the key and value projection work falls, taking the projection total from 268,435,456 to 150,994,944 multiply-adds, and the key-value bytes read fall from 268,435,456 to 33,554,432.¹

What is multi-head latent attention in DeepSeek?

Multi-head latent attention (MLA) compresses each token's keys and values into one small latent vector, caches that vector instead of the full keys and values, and derives the keys and values from it. DeepSeek-V2 reports a 93.3% smaller key-value cache than its earlier 67B model (DeepSeek-V2, 2024).

Sources

1. Attention Is All You Need — arXiv (Vaswani et al.), 12 Jun 2017
2. Fast Transformer Decoding: One Write-Head is All You Need — arXiv (Shazeer), 6 Nov 2019
3. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — arXiv (Ainslie et al.), 22 May 2023
4. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model — arXiv (DeepSeek-AI), 7 May 2024
5. Meta-Llama-3-70B config.json (mirror of the model configuration) — Hugging Face (date = read date), 24 Oct 2026