Multi-Head Attention in transformers

Module 05 · Transformers · 19 concepts

In depth / 15 lines

Multi-Head Attention, in more detail

A single attention layer computes one set of weights over the sequence.

That forces every kind of relationship to share one blend.

Grammatical agreement, coreference and topical relevance would all have to average together.

Multi-head attention runs several attention operations side by side instead.

Each head has its own query, key and value projections, so each looks for something different.

The model dimension is split across the heads rather than multiplied.

With 512 dimensions and eight heads, each head works in 64.

So the total cost is about the same as one full-width head.

Their outputs are concatenated and passed through a final linear layer that mixes them.

Studies of trained models find heads with recognisable specialities.

Some track the previous token, some link verbs to subjects, some resolve pronouns.

Many heads appear to do very little and can be pruned with small loss.

The specialisation is not designed, since nothing tells a head what to become.

Eight to sixteen heads is typical in a mid-sized model, and large models use more.

Grouped-query attention lets heads share keys and values to shrink the cache during inference.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Multi-Head Attention

Split, attend separately, join

One attention operation produces one set of weights, so every position must average the sequence exactly one way. Multi-head attention runs several in parallel.

\mathrm{head}_i = \mathrm{Attention}\left( X W_Q^{(i)}, \; X W_K^{(i)}, \; X W_V^{(i)} \right)

\mathrm{MHA}(X) = \left[ \mathrm{head}_1 ; \ldots ; \mathrm{head}_h \right] W_O

$h$ : the number of heads.
$W_O$ : the output projection, which mixes the heads back together.
The semicolons denote concatenation along the feature axis.

The width bookkeeping

d_k = d_v = \frac{d_{\text{model}}}{h}

Each head works in a narrower space, so the concatenation of all $h$ heads comes back to exactly $d_{\text{model}}$ .

h \times \frac{d_{\text{model}}}{h} = d_{\text{model}}

That is why more heads does not mean more computation. With $d_{\text{model}} = 512$ , eight heads of width 64 and sixteen of width 32 cost the same.

The choice is between many narrow views and a few wide ones, not between cheap and expensive.

Why several are better than one

A single softmax row must sum to one, so attending strongly to a syntactic dependency means attending weakly to everything else.

With eight heads, one can track subject-verb agreement while another follows coreference and a third simply looks at the previous token.

Empirically the heads do specialise, though far less cleanly than the tidy diagrams suggest, and many can be pruned after training with almost no loss.

The parameter count

4 \, d_{\text{model}}^2

Three projections in and one out, each $d_{\text{model}} \times d_{\text{model}}$ once the heads are stacked.

At $d_{\text{model}} = 4096$ that is 67 million parameters per attention layer, and it is independent of $h$ .

One layer, in order

Project the input once into $Q$ , $K$ and $V$ at full width.
Reshape each into $h$ heads of width $d_{\text{model}} / h$ .
Run scaled dot-product attention on every head at once, as a batched matrix multiply.
Concatenate the heads back to full width.
Multiply by $W_O$ .

Step three is why heads are free: they are a reshape, not a loop.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

Why run several attention heads instead of one large one?

Transformers

19 concepts in this module, then a graded test with every answer explained.

Study this module →