Why Eight Attention Heads and Not One Wide One

Last timeThe One Equation in It

Eight heads of sixty-four cost exactly what one head of five hundred and twelve costs, in weights and in multiplications, so the split has to be buying something else.

Section 3.2.2 of the paper is four sentences long and one of them carries the

whole argument. Multi-head attention, it says, allows the model to jointly

attend to information from different representation subspaces at different

positions, and with a single attention head, averaging inhibits this. The rest

of the section is shapes.

The shapes are worth doing first anyway, and not out of pedantry. Until the

weight count is on the table it is natural to read eight heads as eight times

something, and to assume the gain is extra capacity bought with extra

parameters. It is not. At the configuration the paper actually trained, width

512 with eight heads of 64 for both keys and values, the split changes no

weight count and no multiplication count anywhere in the block. Whatever the

heads buy, they buy for nothing, which is a far stronger claim than a larger

model doing better.

FIG 1One attention block at width 512, stage by stage

step	stage	one head	all eight heads together	weights	what happened
1	project the queries	512 in, 64 out	512 in, 512 out	262,144	eight slabs of 512 by 64 side by side are one matrix of 512 by 512
2	project the keys	512 in, 64 out	512 in, 512 out	262,144	the same shape again, and the same reshape
3	project the values	512 in, 64 out	512 in, 512 out	262,144	64 is the value width in the base configuration, equal to the key width
4	score every pair of positions	1,000 by 1,000 scores	eight grids of 1,000 by 1,000	none	no weights live here, which is where the quadratic term comes from
5	take the weighted sum of values	1,000 by 64 out	1,000 by 512 once laid side by side	none	eight outputs of 64 concatenate back to the width they came from
6	project the concatenation	not per head	512 in, 512 out	262,144	the output projection, the one matrix that is shared across heads

6 steps

Shapes at the working scale of this course, a sequence of 1,000 positions at width 512. The fourth column is the one to read: four matrices of 512 by 512, so 1,048,576 weights in the block, and the head count appears nowhere in that total. The two rows with no weights at all are where every cost that grows with the square of the sequence length is incurred.

The shapes, counted once

Read the fourth column. Three input projections and one output projection, each

512 by 512, so the attention block holds $4 \times 512 \times 512 = 1{,}048{,}576$

weights. That total does not contain the head count and cannot. The per-head

query matrix is 512 by 64, there are eight of them, and $8 \times 64 = 512$ , so

the eight slabs stacked side by side are indistinguishable from one 512 by 512

matrix. In the code the paper links to, the split is a reshape.

Keep the rest of the layer in view for scale. The inner feed forward layer of

the same block goes from 512 out to 2,048 and back, which is

$2 \times 512 \times 2{,}048 = 2{,}097{,}152$ weights, exactly twice the

attention block. An encoder layer therefore holds 3,145,728 weights and two

thirds of them sit in the part that has nothing to do with attention at all.

FIG 2Why the score cost does not depend on the number of heads

: the number of positions in the sequence, 1,000 in the worked case
: the number of heads, 8 in the base configuration
: the width of the keys and queries inside one head, 64 in the base configuration
: the width of the stream the block reads and writes, 512

The paper sets the per-head width to the model width divided by the head count, and the moment it does that the head count cancels out of the score cost. At 1,000 positions both sides are 512 million multiply-adds. The reader who was expecting eight heads to cost eight times as much was reading the first factor and not the third.

The split is free, and that is the whole point

Do the arithmetic at the standing scale. One head of 512 forms a 1,000 by 1,000

score matrix whose every entry is a dot product of two vectors of length 512, so

the scores cost $10^6 \times 512 = 5.12 \times 10^8$ multiply-adds. Eight heads

of 64 form eight 1,000 by 1,000 score matrices whose entries are dot products of

length 64, so the scores cost $8 \times 10^6 \times 64 = 5.12 \times 10^8$ . The

same number, for the same reason the weight count was the same number. The

weighted sum of the values comes out equal on both sides by the identical

argument, and the four projections are $4 \times 1{,}000 \times 512 \times 512

\approx 1.05 \times 10^9$ multiply-adds either way, because they are literally

the same matrices.

One quantity is not neutral, and the paper does not raise it. The score entries

that have to exist at once go from $10^6$ to $8 \times 10^6$ , because there are

now eight grids of scores rather than one. Arithmetic is flat in the head count

and the storage for the attention weights is linear in it. With sentences of a

few dozen tokens that was a rounding error, which is why it is absent from a

2017 paper, and it is the term that keeps head count in the conversation when

the sequence is a document rather than a sentence.

The lesson stops here

4 more paragraphs to go

You have read the opening. The rest of the argument, the problems that check whether it landed, and the lines worth keeping at the end all come with a plan.

The first lesson of every course in the library reads the whole way through, free, so you can see exactly what the rest of them are.

See the plan The contents

This is the reading half

Starting the course gives you your own copy of it. Every idea on every page has problems standing under it, marked with a reason rather than a tick, and any sentence you do not believe can be opened and argued with. None of that can happen on a page nobody owns.

The contents

Why Eight Attention Heads and Not One Wide One

The shapes, counted once

The split is free, and that is the whole point

4 more paragraphs to go

This is the reading half

Read alongside