What the Attention Equation Actually Computes

Last timeWhat the Paper Was Replacing

Equation one of the paper read matrix by matrix, with the square root in its denominator derived from the variance of a dot product and the mask put where the paper puts it.

Equation one of the paper is a single line, and nearly everything else in the

architecture is bookkeeping arranged around it. It is also short enough to be

memorised without being understood, which is the usual outcome, so it is worth

reading one piece at a time with the shapes carried along.

FIG 1The equation, with its four pieces pointed at

: applied along each row separately, turning one row of scores into coefficients that sum to one
: every query meets every key, giving one score per ordered pair of positions and nothing else
: the standard deviation of one of those scores, so the division is a change of units
: the coefficients mix the value rows, so each output is a weighted average of values and never of inputs

Take from this that the line has two matrix products in it and one elementwise step, and that the two products do entirely different jobs. The first decides who is listened to. The second decides what is heard, and it reads from the values, which are a third projection of the input and not the input itself.

What the letters are and what shape they have

$Q$ , $K$ and $V$ are matrices whose rows are vectors, one row per position.

There are $n$ queries of width $d_k$ , and $m$ keys of the same width $d_k$

paired with $m$ values of width $d_v$ . Section 3.2.2 fixes the sizes for the

base model: a model width of $d_{\text{model}} = 512$ , $h = 8$ heads, and

$d_k = d_v = 512 / 8 = 64$ for each of them. The projection that makes the

queries of one head is therefore a matrix of $512 \times 64$ weights, and the

same for the keys and the values.

The two lengths are separate on purpose. In the self-attention inside either

stack, the queries, keys and values are three projections of the same sequence,

so $n = m$ . In the attention that joins the stacks, described in section 3.2.3,

the queries come from the decoder and the keys and values from the output of the

encoder, so $n$ is the length of the output and $m$ the length of the input and

the score matrix is not square.

$Q K^{T}$ multiplies an $n$ by $d_k$ matrix into a $d_k$ by $m$ one, giving an

$n$ by $m$ matrix. Entry $(i, j)$ of it is the dot product of query $i$ with key

$j$ , which is a sum of $d_k = 64$ products. Put the standing sizes in. At

$n = m = 1000$ positions, one head holds

$1000 \times 1000 = 1{,}000{,}000$ scores and spends

$1{,}000{,}000 \times 64 = 64{,}000{,}000$ multiply-adds making them. Across

eight heads that is $8 \times 64{,}000{,}000 = 512{,}000{,}000$ , which is the

quadratic term in the paper's opening table written out, and across the six

layers of one stack it is $6 \times 8{,}000{,}000 = 48{,}000{,}000$ scores alive

at the same time, every one of them needed again on the backward pass.

The widths are separate too, and the base model setting both to 64 hides it. A

row of the weight matrix has length $m$ , the number of keys, and knows nothing

about $d_v$ . The output row has width $d_v$ . Nothing in the equation requires

the two to match.

From a row of scores to a weighted average

Write $s_{i1}, \ldots, s_{im}$ for row $i$ of the scaled score matrix. The

softmax of that row is

The lesson stops here

9 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

What the Attention Equation Actually Computes

What the letters are and what shape they have

From a row of scores to a weighted average

9 more paragraphs to go

This is the reading half

Read alongside