Multi-Head Attention Explained

Multi-Head Attention: Many Ways of Looking at the Same Sentence

A note on tone

Like my other notes, this is my own understanding written in plain words. For the full picture, read the original transformer paper and a proper deep learning course. Think of this as a friendly map, not the textbook.

Series: Neurons, Weights and Bias Explained → Neural Networks Explained → Activation Functions and GELU → Transformers Explained → Multi-Head Attention (this note).

In Transformers Explained we built attention: each word makes a query, compares it with every key, and blends values by the resulting weights. Now the natural question is: why do real transformers use many heads of it, and not just one?

The problem with a single head

A word usually has several relationships at once. In "The cat sat on the mat", the word "sat" cares about:

who did it (the cat),
where (the mat),
and the words right next to it.

One attention head produces one softmax, so one set of weights. It has to squeeze all those different needs into a single blended view, and averaging them tends to blur each one out.

The fix: give the model several attention heads. Each head looks at the sentence independently, with its own view.

Being honest

This picture is an illustration. Real heads are not neatly labelled "who" and "where". Researchers have found some clean, interpretable heads (for example heads that attend to the previous token, or "induction heads" that help copy patterns), but many heads are fuzzy or partly redundant, and some can be removed with little loss. The intuition is right; the tidy labels are a simplification.

How it works: split, attend, glue

You might expect that eight heads means eight times the work. It doesn't. The trick is that the model's width is shared out among the heads.

Each word arrives as a vector of size $d_{m o d e l}$ (say 512).
Project it into queries, keys and values, then split them into $h$ slices. With 8 heads, each head gets 64 dimensions.
Every head runs the same attention on its own slice, in parallel.
Concatenate the 8 outputs (8 × 64 = 512 again).
Multiply by one more matrix $W^{O}$ , which mixes what the heads found into a single vector.

{head}_{i} = Attention (X W_{i}^{Q}, X W_{i}^{K}, X W_{i}^{V})

MultiHead (X) = Concat ({head}_{1}, \dots, {head}_{h}) W^{O}

Every $W$ here is a matrix of learned weights, exactly the kind we met in Neurons, Weights and Bias Explained. And each head's slice is small, so:

d_{h e a d} = \frac{d_{m o d e l}}{h}

The original transformer used 8 heads of 64 (total 512). GPT-2 small uses 12 heads of 64 (total 768). Because the width is split rather than multiplied, many small heads cost about the same as one big head, but give many different ways of looking.

In code

Building on the attention function from Transformers Explained:

import numpy as np

def multi_head(X, Wq, Wk, Wv, Wo, h):
    n, d = X.shape
    dh = d // h                               # size of each head
    Q, K, V = X @ Wq, X @ Wk, X @ Wv          # project once, for all heads
    heads = []
    for i in range(h):
        s = slice(i * dh, (i + 1) * dh)       # this head's slice
        heads.append(attention(Q[:, s], K[:, s], V[:, s]))
    return np.concatenate(heads, axis=-1) @ Wo

In real libraries there is no Python loop: the slices are reshaped into an extra "head" dimension so all heads run as one big batched matrix multiplication (Matrix Multiplication Explained again). That is why it is fast on GPUs.

What is the extra matrix $W^{O}$ for?

After concatenation, the heads' outputs are just sitting side by side. $W^{O}$ lets the model combine them: it can mix information across heads, decide which head's finding matters for which output dimension, and return everything to width $d_{m o d e l}$ so it can flow into the next part of the block (the GELU feed-forward layer).

The memory problem, and MQA / GQA

Here is a practical issue that shaped recent models. When generating text, the KV cache stores the keys and values of every past token, for every head, in every layer. With many heads and long chats, that cache becomes huge and slows generation, because the GPU spends its time moving memory.

The observation: queries need their own heads, but keys and values can be shared.

Multi-head attention (MHA): every query head has its own key and value head. The original design.
Multi-query attention (MQA, Shazeer 2019): all query heads share one key/value head. Smallest cache, but quality can drop a bit.
Grouped-query attention (GQA, Ainslie et al. 2023): query heads are split into groups, and each group shares a key/value head. A middle ground that many recent large models use.

The queries keep their own heads, so the model still looks at the sentence in different ways. Only the memory is shared. Fewer K/V heads means a smaller KV cache and faster generation.

Putting it back in the block

Multi-head attention is the "communication" step of the transformer block (see the diagram in Transformers Explained). The full recipe for one block is:

x \to MultiHeadAttention \to Add&Norm \to FFN with GELU \to Add&Norm

Repeat that dozens of times, and you have a language model. Every piece is a neuron-style weighted sum, a bend, or a matrix multiplication. Nothing new appears at the end. It is the same small idea, wired cleverly.

Key takeaway

Multi-head attention runs several small attentions side by side, each on its own slice of the vector with its own learned Q, K and V, then glues the results together with $W^{O}$ . It doesn't cost much more than one head, because the width is shared. Different heads can pick up different relationships, and MQA/GQA later showed that keys and values can be shared to keep generation fast.

Try it yourself

Run multi_head with $d = 8$ , $h = 2$ on random data and check the output shape equals the input shape
Set $h = 1$ and confirm it behaves like plain attention
Count the parameters of the four matrices $W^{Q}, W^{K}, W^{V}, W^{O}$ for $d = 512$ (each is $d \times d$ )
Sketch how the KV cache size changes with MHA, GQA (2 groups) and MQA
Load a small model in Hugging Face and print its num_attention_heads and hidden_size