Self-Attention

By Hao Ma  ·  haoma7.github.io

From the bottleneck in recurrent models, to what Q, K and V actually are, to the whole attention formula worked out by hand on three words.

1.1Motivation

1.1.1The problem

To see why the Transformer was a surprise, you first have to see what it removed, and why that thing was hard to give up.

The paper is about sequence transduction: read one sequence of symbols, produce another. Machine translation is the standard case — English words in, German words out. Language modelling is another. Both are transduction problems, and in 2017 both were dominated by one family of models.

That family is the recurrent neural network, or RNN. Plain RNNs, long short-term memory networks (LSTM), and gated recurrent networks were the established state of the art for this work. A large amount of research had gone into pushing them further, both as recurrent language models and as encoder–decoder architectures. This was not a weak baseline that was easy to beat. It was the best anyone had.

Two words in that last sentence are used throughout this series, so here they are once. An encoder is the half of a translation model that reads the input sentence and turns it into vectors. A decoder is the half that writes the output sentence, one token at a time, reading what the encoder produced as it goes. Section 1.2.5 needs both, and so does Section 1.3.3.

How a recurrent model spends its time

A recurrent model lines its computation up with the positions in the sequence. One position, one step. Positions are counted from 0 all through this series, so the first word of a sentence sits at position 0. At step t it produces a hidden state \(\mathbf{h}_t\). This is a vector holding everything the model has worked out so far about the sequence, up to and including position t.

The hidden state is computed from two things — the previous hidden state and the input at the current position:

To get \(\mathbf{h}_t\) you need \(\mathbf{h}_{t-1}\). To get \(\mathbf{h}_{t-1}\) you need \(\mathbf{h}_{t-2}\). The chain runs all the way back to the start of the sequence.

Think for a moment You have a GPU with thousands of cores and one training sentence of 50 words. How many of those 50 positions can the cores work on at the same time?

One. Not fifty, not eight — one. The cores are there, and they are idle, because position 1 cannot begin until position 0 has finished. This is what the paper means when it says the sequential nature of these models rules out parallelisation within a training example (Figure 1). It is not that the model is slow to compute. It is that the work refuses to be spread out.

What people tried before this paper

The field knew about this. Two lines of work had already produced real gains in computational efficiency: factorisation tricks, and conditional computation. The second of these also improved how well the models performed, not just how fast they ran.

These were genuine improvements. They made recurrent models cheaper. They did not touch the dependency itself, and the authors say so plainly:

The fundamental constraint of sequential computation, however, remains. Vaswani et al., 2017

Meanwhile a second idea had been growing: the attention mechanism. Attention lets a model relate two positions directly. The part that matters is what it ignores: the distance between them. Position 0 and position 49 are one step from each other, the same as position 0 and position 1. By 2017 attention had become a standard part of strong sequence models across many tasks.

So attention was already known, and the recurrence problem was already known. Why did nobody put the two facts together? They almost had. In all but a few cases, attention was used together with a recurrent network (Figure 2). Attention was treated as an improvement you added on top of a chain, and the chain stayed underneath, still serial, still refusing to parallelise.

Key takeaway

• Recurrent networks — RNN, LSTM, gated recurrent — were the state of the art for sequence transduction, including language modelling and machine translation.
• They compute a hidden state per position, \(\mathbf{h}_t = f(\mathbf{h}_{t-1}, \mathbf{x}_t)\), so position t cannot start until t−1 has finished. That rules out parallelising within one training example.
• Factorisation tricks and conditional computation improved efficiency, and conditional computation also improved performance, but neither removed the dependency.
• Attention could already relate positions at any distance. In almost every case it was still added on top of a recurrent network.

Try it (3 minutes): Write the loop for t in range(50): h = tanh(W @ h + U @ x[t]). In Python, @ means matrix multiply; W and U are the two weight matrices, h is the hidden state, and x[t] is the input at position t. Now try to rewrite it as a single matrix operation over all 50 positions, with no loop. Find the exact line where you get stuck. That line is the constraint the whole paper is about.

1.1.2The idea

Once the problem is stated in that way, the next step follows directly. If the chain is what blocks parallelism, and attention can already relate any two positions at any distance, then stop treating attention as a helper. Make it the whole model.

That is the Transformer: an architecture that gets rid of recurrence and relies entirely on an attention mechanism. The job attention is given there is to draw global dependencies between input and output. Global here means that any position may look at any other position directly, with no path through the steps in between.

✗ Common mistake Readers often take the title to mean attention is the only component in the model, but that is wrong. The Transformer still contains embeddings, feed-forward layers and normalisation, and it needs a positional encoding, because deleting recurrence deletes the model's only sense of word order. The title is a claim about what you can remove: what "all you need" replaces is recurrence and convolution, nothing more.

Two results follow, and they are the reason the paper was taken seriously immediately. The first is the one we have been building toward: the design allows significantly more parallelisation. The second is what that parallelisation bought. The Transformer reached a new state of the art in translation quality after training for as little as twelve hours on eight P100 GPUs.

Key takeaway

• The Transformer removes recurrence and relies entirely on attention.
• Attention's job there is to draw global dependencies between input and output.
• The payoff is significantly more parallelisation.
• It reached a new state of the art in translation after as little as twelve hours on eight P100 GPUs.

1.2Self-Attention

Before the pieces, here is the whole of self-attention on the running example, played once from start to finish (Figure 3). Nothing in it is explained yet — the rest of Section 1.2 takes it apart one step at a time, and you can come back to this figure when you want to see where a step sits.

1.2.1Where Q, K and V come from

Section 1.1 ended with a decision: delete the chain, and let attention carry the whole model. That leaves an obvious question. What does attention actually compute?

Almost everyone gets stuck in the same place on a first reading, and it is not the formula. It is three letters. Attention is written in terms of three matrices called Q, K and V — query, key and value — and the formula never says where they come from. So we start there, before any arithmetic.

Take the sentence I love LLMs. Suppose the embedding step has already turned it into a matrix 3 × 4 that we will call \(\mathbf{X}\): three tokens, four numbers each. Three rows because this series gives every word one row, as a simplification. A real tokeniser can cut one word into two — LLMs becomes LLM and s — and Part 0, Section 0.5.2 works that case through from characters to \(\mathbf{X}\). The model multiplies \(\mathbf{X}\) by three separate weight matrices, which are learned during training, and calls the three results \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) (Figure 4).

Think for a moment All three come from the same \(\mathbf{X}\), by the same kind of operation. So what makes \(\mathbf{Q}\) different from \(\mathbf{K}\)?

Only the numbers inside the three weight matrices. Nothing else. At this stage \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) are three different views of one input — three ways of re-describing the same three tokens. They start out interchangeable, and training pulls them apart because the three jobs they are used for are different.

Those three jobs are easiest to see in an old-fashioned library. Every book on the shelf has a label on its spine, and contents inside. When you go looking for something, you carry a description of what you want. The description you carry is the query. The label on the spine is the key. What is printed inside the book is the value. You compare your description against the labels to decide which books to open, and then you read the contents — never the labels — of the books you chose. Q is compared against K; V is what you come away with.

✗ Common mistake The library story sounds like three separate objects, so it is tempting to read Q, K and V as three different inputs — a question from somewhere, a set of labels from somewhere else, and a body of content from a third place. That is wrong. A self-attention layer has only one input. The sentence is compared against itself, and the three roles are three projections — three linear transformations — of that one sentence.

Key takeaway

• Q, K and V are produced from the same input \(\mathbf{X}\) by three learned weight matrices \(\mathbf{W}^Q\), \(\mathbf{W}^K\), \(\mathbf{W}^V\).
• Query = what a token is looking for. Key = what a token advertises. Value = what a token hands over once it has been chosen.
• Q is only ever compared against K. V is only ever read.

Try it (2 minutes): Take 3 × 4 for \(\mathbf{X}\) and 4 × 4 for each of the three weight matrices. Write down the shape of \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) before reading on. If you got 3 × 4 three times, you have the mechanics; the rest of this section is about what the model does with them.

1.2.2From scores to weights

We now have a query for every token and a key for every token. The question attention has to answer is: when the model encodes the word love, how much should it look at I, at love itself, and at LLMs?

That needs a number for every pair of positions. The cheapest useful measure of how much two vectors agree is the dot product: multiply the matching entries of two vectors and add up the results. Vectors with more entries tend to give bigger results, which is a fact Section 1.3.2 will need. Do it for every query against every key at once, and one matrix multiply \(\mathbf{Q}\mathbf{K}^\top\) gives the whole table of scores.

Think for a moment Write \(\mathbf{q}_1\), \(\mathbf{q}_2\), \(\mathbf{q}_3\) for the three rows of \(\mathbf{Q}\) and \(\mathbf{k}_1\), \(\mathbf{k}_2\), \(\mathbf{k}_3\) for the three rows of \(\mathbf{K}\). Using only those six vectors and the dot product, write out the 3 × 3 matrix \(\mathbf{Q}\mathbf{K}^\top\), entry by entry.

Transposing \(\mathbf{K}\) turns its rows into columns, so column \(j\) of \(\mathbf{K}^\top\) is \(\mathbf{k}_j\). A matrix product puts row \(i\) of the left factor against column \(j\) of the right one, so entry \((i,j)\) is \(\mathbf{q}_i \cdot \mathbf{k}_j\):

Read it as a table: row \(i\) is token \(i\) asking, column \(j\) is token \(j\) answering, and the cell is their score. Every query meets every key exactly once, which is why the table is square and why one multiply is enough. Section 1.2.3 fills this table with numbers.

Two things then happen to that table before it can be used as a set of weights.

First it is divided by \(\sqrt{d_k}\), where \(d_k\) is the width of one key vector — the number of columns of \(\mathbf{K}\). This is the scaling step, and it exists because raw dot products get bigger as the vectors get more entries, for reasons that have nothing to do with the sentence. Section 1.3.2 does the two-line calculation.

Then it goes through softmax, applied to each row on its own. For now all you need is this: softmax takes a row of arbitrary numbers and turns it into a row of non-negative numbers that add up to 1, keeping the order — the biggest score comes out as the biggest weight. Section 1.3.1 is a full review with worked numbers, and it is worth reading if you are not familiar with softmax.

Put together, that is the whole of scaled dot-product attention (Figure 5):

The output of the softmax is the attention weight matrix. It is worth looking at real numbers rather than symbols, because its shape says something the formula does not (Figure 6).

Those numbers are worth earning rather than being handed. The next section runs the whole of Equation 1-1 on I love LLMs, from the embeddings to the final output, one arithmetic step at a time. Nothing new appears — it is the same formula with the symbols replaced by digits.

1.2.3The formula on three words

The setup

Four matrices to start with. First \(\mathbf{X}\), the embedded sentence: three tokens, four numbers each, 3 × 4. Row \(i\) is token \(i\). Every number printed from here on is rounded to two decimals, so a row of weights sometimes adds up to 0.99 rather than 1.

Then the three weight matrices the model learned during training, each 4 × 4. These are the only things in this calculation that training changes; everything else follows from them. In a real model these matrices are not square: \(\mathbf{W}^Q\) and \(\mathbf{W}^K\) are \(d_{\text{model}} \times d_k\), and \(\mathbf{W}^V\) is \(d_{\text{model}} \times d_v\). Section 1.3.3 sets out which of those sizes are forced. Square ones here only keep the arithmetic small.

Small whole numbers, picked so the arithmetic stays in your head. A real model has messy decimals and could have 512 columns instead of 4. Every step below would be identical.

Step 1 — three projections give Q, K and V

This is Figure 4 with numbers in it. Entry \((i,j)\) of \(\mathbf{X}\mathbf{W}^Q\) is row \(i\) of \(\mathbf{X}\) dotted with column \(j\) of \(\mathbf{W}^Q\). Take the top-left entry: row 1 of \(\mathbf{X}\) is \((0,0,1,1)\), column 1 of \(\mathbf{W}^Q\) is \((0,0,1,1)\), so

Twelve such dot products fill in \(\mathbf{Q}\) — one per entry of a 3 × 4 matrix. Do the same with \(\mathbf{W}^K\) and \(\mathbf{W}^V\):

All three are 3 × 4, the same shape as \(\mathbf{X}\). Row 2 of \(\mathbf{Q}\) is the query for love; row 2 of \(\mathbf{K}\) is its key; row 2 of \(\mathbf{V}\) is its value. Three descriptions of one token, and the only reason they differ is that three different weight matrices produced them. Since \(\mathbf{K}\) has four columns, \(d_k = 4\) — a fact we need in step 3.

Think for a moment Every token in this example has a 1 in the last column of \(\mathbf{X}\), and the three tokens differ only in which of the first three columns is set. If you gave two tokens identical rows in \(\mathbf{X}\), what would their rows in \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) look like?

Identical as well. The projection has no idea where a row sits in the sentence — it multiplies each row by the same matrix, independently. Two identical embeddings produce identical queries, keys and values, and therefore identical attention behaviour. This is exactly why position has to be injected separately, before attention ever runs.

Step 2 — one dot product per pair of positions

Write \(\mathbf{q}_i\) for row \(i\) of \(\mathbf{Q}\), and \(\mathbf{k}_j\) for row \(j\) of \(\mathbf{K}\). The score for "how much should token \(i\) look at token \(j\)" is \(\mathbf{q}_i\) dotted with \(\mathbf{k}_j\). Take \(\mathbf{q}_2\), the query for love, against \(\mathbf{k}_3\), the key for LLMs:

Nine pairs, nine dot products, and doing all nine at once is what the single matrix multiply \(\mathbf{Q}\mathbf{K}^\top\) means. The transpose is there so that row \(i\) of \(\mathbf{Q}\) meets column \(j\) of \(\mathbf{K}^\top\), which is row \(j\) of \(\mathbf{K}\):

The 3 in row 2, column 3 is the score we just did by hand. The result is 3 × 3: one number per ordered pair of positions.

Think for a moment Row 2 column 3 is 3, the score from love to LLMs. Row 3 column 2 is 2, the score from LLMs to love. Same pair of tokens. Why are the two numbers different?

Because the two cells do not use the same pair of vectors. The first one is \(\mathbf{q}_2 \cdot \mathbf{k}_3\), the query for love against the key for LLMs. The second is \(\mathbf{q}_3 \cdot \mathbf{k}_2\), the query for LLMs against the key for love. The two tokens have swapped roles, and queries and keys come out of different weight matrices, so nothing forces the two numbers to agree. This is where the asymmetry in the mistake box below is born: right here, at step 2.

Step 3 — divide by √d_k

Here \(d_k = 4\), so \(\sqrt{d_k} = 2\), and every score is halved:

Notice what this step does not do. Every number shrank by the same factor, so the order within each row is untouched: 5 was the largest score in row 2 before, and 2.5 is the largest after. Scaling changes the spread of the scores, not their ranking, and the spread is what softmax reacts to. Section 1.3.2 works out why that matters.

✗ Common mistake Here the two are the same number, because the toy model is four numbers wide and \(\mathbf{K}\) has four columns, so \(d_k\) and \(d_{\text{model}}\) are both 4. It is easy to leave with the divisor is \(\sqrt{d_{\text{model}}}\), and that is wrong. They are different sizes, and the formula uses \(d_k\), the width of one key vector. Part 2 pulls them apart: there each head gets a slice of the width, so \(d_k\) shrinks while \(d_{\text{model}}\) stays where it was, and the divisor shrinks with it.

Step 4 — softmax, one row at a time

Softmax takes a row, raises \(e\) to the power of each entry, and divides each result by their total. Take row 2, the row for love, which is \(1.0,\ 2.5,\ 1.5\). First the three exponentials:

Their total is \(2.718+12.182+4.482=19.382\). Now divide each one by that total:

Those three add up to 1, and they are row 2 of Figure 6. The gap between the scores did the work: 2.5 beat 1.5 by a single point, but after exponentiating, the winner takes almost three times the weight.

Rows 1 and 3 follow the same recipe. Row 1 is \(2.0,\ 1.0,\ 1.0\) — two equal scores must give two equal weights, so it comes out \(0.58,\ 0.21,\ 0.21\). Row 3 is \(1.5,\ 1.0,\ 2.5\), the same three numbers as row 2 in a different order, so it gives the same three weights in that order. Stacking all three rows gives the attention weight matrix \(\mathbf{A}\):

Figure 7 lines the three stages up side by side. It is the same 3 × 3 grid throughout — only the numbers in the cells change.

Step 5 — spend the weights on V

The weight matrix is not the answer. It is a set of instructions for mixing the value vectors, and Equation 1-1 finishes by carrying them out: multiply \(\mathbf{A}\) by \(\mathbf{V}\), as Figure 8 shows. Row 2 of the result is row 2 of \(\mathbf{A}\) — the weights \(0.14,\ 0.63,\ 0.23\) — applied to the three rows of \(\mathbf{V}\):

Here we define \(\mathbf{Z} = \mathbf{A}\mathbf{V}\). The subscripts are labels, not new objects: \(\mathbf{v}_{\text{I}}\), \(\mathbf{v}_{\text{love}}\) and \(\mathbf{v}_{\text{LLMs}}\) are rows 1, 2 and 3 of \(\mathbf{V}\), the same vectors written \(\mathbf{v}_1\), \(\mathbf{v}_2\) and \(\mathbf{v}_3\) elsewhere. So \(\mathbf{z}_{\text{love}}\) is row 2 of \(\mathbf{Z}\).

Figure 9 takes that one row apart, column by column, so you can see which numbers produce each of the four output numbers.

Take the first column of \(\mathbf{V}\). The three value vectors have 1, 2 and 1 there, so

Repeat for the other three columns, then do the same for rows 1 and 3, and the layer is finished:

The output is 3 × 4 — the same shape we started with. Three tokens in, three tokens out, four numbers each. That is what makes attention stackable: the next layer cannot tell whether its input came from an embedding table or from another attention layer.

Key takeaway

• Three weight matrices turn one 3 × 4 input into \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\), each 3 × 4.
• \(\mathbf{Q}\mathbf{K}^\top\) fills a 3 × 3 grid — one score per ordered pair of tokens.
• Dividing by \(\sqrt{d_k} = 2\) and running softmax along each row turns that grid into the attention weight matrix \(\mathbf{A}\).
• \(\mathbf{A}\mathbf{V}\) brings the shape back to 3 × 4: same in, same out, which is what lets attention layers stack.

1.2.4What the whole operation buys

The arithmetic is done. Now read what it produced — first for one word, then for the layer as a whole.

Think for a moment Row 2 of \(\mathbf{X}\) was \((0,1,0,1)\). Row 2 of the output is \((1.63,\ 0.14,\ 0.23,\ 0.63)\). What did the layer actually do to love?

It replaced a description of love on its own with a description of love in this sentence. The old row knew nothing about its neighbours. The new one is 63% its own value vector and 37% borrowed from I and LLMs, so the surrounding words are now baked into the number itself. That mixing is the entire point of the layer, and it is why Figures 8 and 9 are worth a second look now that you know what the mixing does.

✗ Common mistake The weight matrix is the part everyone visualises and the part papers publish heatmaps of, so it is tempting to think the output of attention is the weight matrix, but that is wrong. It is 3 × 3 — one entry per pair of positions, and not even the right shape to be passed on. What is handed to the next layer is \(\mathbf{A}\mathbf{V}\), which is 3 × 4. Weights are a plan; \(\mathbf{V}\) is what gets spent.

Try it (4 minutes): Compute row 1 of \(\mathbf{A}\mathbf{V}\) yourself. The weights are \(0.58,\ 0.21,\ 0.21\) and the rows of \(\mathbf{V}\) are \((1,1,0,0)\), \((2,0,0,1)\) and \((1,0,1,0)\). You should get \((1.21,\ 0.58,\ 0.21,\ 0.21)\). If your first entry is 1.21, the whole of Equation 1-1 is now something you can run by hand.

Think for a moment Softmax was applied along the rows, so every row adds up to 1. Why the rows and not the columns?

Because a row is one token's budget. Row \(i\) says how token \(i\) divides its attention among all the positions available to it, and a budget has to add up to the whole. A column would be "how much attention did everyone pay to token \(j\)", and there is no reason that should come to 1 — a very informative token should be allowed to receive more total attention than a boring one.

✗ Common mistake The score started life as a dot product, and a dot product genuinely is symmetric, so it is tempting to expect that if love pays 0.23 to LLMs, then LLMs pays 0.23 back. That is wrong, and Figure 6 says so: LLMs pays only 0.14 back. The score is not \(\mathbf{x}_i \cdot \mathbf{x}_j\); it is \(\mathbf{q}_i \cdot \mathbf{k}_j\), and \(\mathbf{q}_i\) came through \(\mathbf{W}^Q\) while \(\mathbf{k}_j\) came through \(\mathbf{W}^K\). Two different matrices, so swapping \(i\) and \(j\) gives a different number. Attention is directional, and that is deliberate: a verb needing its object is not the same relation as an object needing its verb.

Key takeaway

• \(\mathbf{Q}\mathbf{K}^\top\) produces one score for every (query, key) pair.
• Dividing by \(\sqrt{d_k}\) keeps those scores at a workable size — see Section 1.3.2.
• Softmax runs along each row, so each row is one token's attention budget and sums to 1.
• The weight matrix is not symmetric, because Q and K come from different learned matrices.

Now step back and look at the whole operation at once, because the reason for all of this is visible only from a distance (Figure 10).

A recurrent model reaches position 2 by walking through positions 0 and 1. Self-attention reaches position 2 from position 0 directly, in one matrix multiply, and it does the same for every other pair at the same time. The distance between two tokens has stopped being a number of steps.

Key takeaway

• The weight matrix is applied to V, so every output vector is a blend of all the value vectors.
• Every output position is computed in the same pass — no position waits for another.
• Any two positions are one operation apart, whatever the distance between them in the sentence.

Try it (3 minutes): Using Figure 6's numbers, compute the output row for LLMs by hand — the one row that was not worked out above. Multiply each value vector by its weight and add the three results. Then check each of the four entries against the matching entries of the three value vectors: it should sit between the smallest and the largest of them, because a weighted average never leaves the range of the things it averages.

1.2.5Three names people mix up

Three terms get used as if they were interchangeable, and they are not. Sorting them out now prevents a lot of confusion later, because the decoder uses two of them side by side.

Scaled dot-product attention is the calculation — Equation 1-1. It is a function. Give it any three matrices \(\mathbf{Q}\), \(\mathbf{K}\), \(\mathbf{V}\) of compatible shapes and it returns an output. It does not know or care where they came from.

Self-attention is a wiring decision: send the same sequence in three times, so \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) are all projections of one input, in the sense of Section 1.2.1. That is what the encoder does, and it is what every figure so far has shown.

Cross-attention is the other wiring: \(\mathbf{Q}\) comes from one sequence and \(\mathbf{K}\) and \(\mathbf{V}\) from another. In the decoder, the queries come from the text generated so far, while the keys and values come from the encoder's output — which is how a translation model lets each output word look at the whole input sentence (Figure 11).

Figure 11 also shows something the earlier figures hid. Because \(\mathbf{Q}\) can come from a different sequence than \(\mathbf{K}\), the score matrix \(\mathbf{Q}\mathbf{K}^\top\), and so the attention weight matrix that softmax makes from it, need not be square. Its height is the number of queries and its width is the number of keys. In self-attention those happen to be equal, so it looks square — but that is a coincidence of the wiring, not a rule of the mechanism. Section 1.3.3 sets out which shapes are actually forced.

Key takeaway

• Scaled dot-product attention is the algorithm; self-attention and cross-attention are two ways of wiring data into it.
• Cross-attention proves the terms are not synonyms: same engine, not self-attention.
• The weight matrix is queries-by-keys. Square only when the two sequences are the same length.

1.3The Math Behind It

This section collects the three pieces of mathematics the chapter leaned on. Nothing here is new material — it is the detail behind claims already made, gathered in one place so the main argument could run without stopping. Read it whenever a forward reference sent you here, or read it straight through as a review.

1.3.1Softmax, row by row

Softmax turns a list of arbitrary scores into a list of weights. It appears in Equation 1-1 because the scores coming out of \(\mathbf{Q}\mathbf{K}^\top\) can be any real numbers — positive, negative, large — and what the model needs instead is a set of mixing proportions.

For a row \(\mathbf{x}\) with entries \(x_1, \dots, x_n\), the \(i\)-th output is:

Equation 1-2 says: exponentiate every entry, then divide each by the total. In the attention mechanism this is applied to a matrix one row at a time, so row 1 is computed without reference to row 2. Two queries looking at three keys give a 2 × 3 matrix, and softmax turns it into two independent distributions over the same three keys.

A worked example. Take the row [1.0, 2.0, 3.0]. The exponentials are about 2.72, 7.39 and 20.09, which add up to 30.20. Dividing through gives [0.09, 0.24, 0.67]: the largest input takes two-thirds of the weight. Now take the row [1.0, 1.0, 1.0]. All three exponentials are equal, so the output is [0.33, 0.33, 0.33] — equal scores produce an even split, and note that it is the differences between scores that matter, not their absolute size.

Four properties are worth holding on to, because the rest of the chapter uses all four:

• Non-negative. Every output is at least 0, because \(e^x\) is never negative.
• Sums to 1. By construction: the denominator is the sum of the numerators. This is what makes a row an attention budget.
• Order-preserving. A bigger score always gets a bigger weight, so softmax never reorders anything.
• Shift-invariant. Adding the same constant \(c\) to every entry leaves the output unchanged, since \(e^{x_i + c} = e^{c}e^{x_i}\) and the \(e^{c}\) cancels top and bottom. Every real implementation uses this: subtract the row maximum before exponentiating, and \(e^{x}\) can never overflow.

✗ Common mistake Softmax is often learned as outputs are strictly between 0 and 1, and that is wrong. Two ordinary cases break it. A row of length 1 comes out as exactly 1.0. And when a position is masked — the decoder does this so a token cannot see the future — that position is set to \(-\infty\) before the softmax, and \(e^{-\infty} = 0\) makes its weight exactly 0. Non-negative and summing to 1 is the version that survives contact with a real model.

One more property is not on the list because it is a failure rather than a feature, and it is the reason Section 1.3.2 exists at all. Softmax cares about the gaps between scores, and it cares about them exponentially. Stretch every score by a factor of 8 and the weights do not stretch by 8 — they collapse (Figure 12).

Key takeaway

• Softmax exponentiates and normalises, one row at a time.
• Outputs are non-negative and sum to 1; they are not strictly positive once masking is involved.
• Only the differences between scores matter, and they matter exponentially.
• Large scores saturate the output toward one-hot, which flattens the gradient.

1.3.2Why divide by the square root of \(d_k\)

Section 1.2.2 divided the scores by \(\sqrt{d_k}\) and promised a reason. Here it is. It rests on one assumption and two short steps.

Assume the entries of a query vector \(\mathbf{q}_i\) and a key vector \(\mathbf{k}_j\) are independent, with mean 0 and variance 1. Their dot product is a sum of \(d_k\) products:

Take that as a simplifying assumption, not a fact about trained models. Real queries and keys do not have mean 0 and variance 1 exactly — the Common mistake at the end of this section says why the argument is still worth having.

Each term is a product of two such entries. Two independent entries with mean 0 give a product with mean 0, and the variance of that product is \(\operatorname{Var}(q)\operatorname{Var}(k) = 1\). Different terms are independent as well, and variances of independent terms add, so:

Equation 1-3 says the variance is \(d_k\), which means a typical magnitude of \(\sqrt{d_k}\). At \(d_k = 64\) the raw scores sit around ±8; at \(d_k = 512\), around ±23. Nothing about the sentence changed — the scores grew because the vectors have more entries. Feed scores of that size into softmax and you land in the right-hand half of Figure 12: one weight near 1, the rest near 0, and almost no gradient to learn from.

Dividing by \(\sqrt{d_k}\) undoes exactly that growth. The scaled scores have variance about 1 whatever \(d_k\) is. A score therefore has about the same typical size whatever \(d_k\) the model uses, and so whatever head count it uses.

✗ Common mistake It is tempting to describe this as an experimental finding, as if the authors tried several divisors and picked the best, but that is wrong. The paper states it in a footnote as a suspicion supported by the calculation above, under an explicit assumption about the distribution of the entries. The distinction matters when you meet a model where that assumption does not hold: the argument is a piece of reasoning you can check, not a measurement you have to trust.

Key takeaway

• Under mean-0, variance-1, independent entries, \(\mathbf{q} \cdot \mathbf{k}\) has variance \(d_k\).
• So raw scores grow like \(\sqrt{d_k}\) for reasons unrelated to the content.
• Dividing by \(\sqrt{d_k}\) restores variance about 1 and keeps softmax out of its saturated region.
• This is an argument under an assumption, not an experimental result.

1.3.3Which shapes are forced

Four constraints on the shapes of \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) get treated as one rule and memorised wrongly. Two of them are forced by matrix multiplication, one by the meaning of the mechanism, and one is not forced at all. The subscripts do all the work below. \(L_Q\), \(L_K\) and \(L_V\) are lengths — how many rows a matrix has. \(d_q\), \(d_k\) and \(d_v\) are widths — how many columns (Figure 13).

1. Q and K must have the same width. Forced by the multiply. To compute \(\mathbf{Q}\mathbf{K}^\top\), the inner dimensions must agree, so \(d_q = d_k\). From here on that shared width is written \(d_k\), which is how the paper writes it. The resulting \(\mathbf{Q}\mathbf{K}^\top\) is a matrix of shape \(L_Q \times L_K\).

2. K and V must have the same length. Forced by meaning. Keys and values come in pairs — a key is the label, its value is the contents. Weight number \(j\) is computed from key \(j\) and applied to value \(j\), so there must be exactly as many values as keys: \(L_K = L_V\).

3. Q need not have the same length as K. This is the one people get wrong, and Figure 11 already showed why. In self-attention all three come from one sequence, so \(L_Q = L_K = L_V\) and everything looks square. In the decoder's cross-attention the queries come from the output built so far and the keys and values from the encoder, so \(L_Q \neq L_K\) is completely normal — and necessary, since a translation rarely has the same number of tokens as its source.

4. K and V need not have the same width. Not forced by anything. \(d_k\) controls how scores are computed; \(d_v\) controls how wide the output is. The two never meet in a multiplication. The paper sets \(d_k = d_v = 64\) for symmetry and for convenience in the multi-head split, but that is a choice, and later architectures have made other ones.

Putting it together: \(\mathbf{Q}\) is L_Q × d_k, \(\mathbf{K}\) is L_K × d_k, \(\mathbf{V}\) is L_K × d_v. The scores are L_Q × L_K, and so are the weights after softmax. Multiplying by \(\mathbf{V}\) gives an output of L_Q × d_v: one row per query, as wide as one value.

Key takeaway

• \(d_q = d_k\) is forced by the multiply; \(L_K = L_V\) is forced by key–value pairing.
• \(L_Q\) is free, which is exactly what makes cross-attention possible.
• \(d_v\) is free; the paper sets it equal to \(d_k\) by choice.
• The output has one row per query and one column for each entry of a value vector: \(L_Q \times d_v\).

Try it (3 minutes): A decoder is translating into a sentence that is 7 tokens long so far, attending to a 12-token source, with \(d_k = 64\) and \(d_v = 64\). Write down the shape of the score matrix and of the layer's output. Then say what goes wrong at the next multiply if you write the score matrix as \(L_K \times L_Q\) instead: which two inner dimensions stop agreeing?