Queries, keys and values

A library holds a question, a shelf of labelled folders, and the pages inside them. You do not read every page, because you compare the question with the labels and then open the folders whose labels match.

The library fixes the three roles, and it stops at the reading. A librarian opens one folder; attention opens every folder and takes a share of each page, in proportion to how well its label matches the question. Each item produces three vectors, lists of numbers, through three learned linear maps, each a fixed matrix that turns one vector into another. A query says what a position is looking for, a key says what an item advertises, and a value says what it passes on [S:74:224]. The arithmetic is a comparison followed by an average: the query is compared with every key, and the resulting weights average the values. !The dot product of a query and a key grows with the length of the vectors, so raw scores drift upward as the model widens. Without the division the leading weight approaches one and the softmax stops responding; dividing by the square root of the key length holds the three weights at 0.665, 0.245 and 0.090 at every width. The flat line is the reason the scale is in the formula. Status: synthetic.

The three maps are separate on purpose. With one map, the score for item $i$ reading item $j$ would equal the score for item $j$ reading item $i$, because $(W{x}_i)\cdot (W{x}_j)$ is symmetric. Two maps break that symmetry, so the model can learn that a pronoun looks at its noun without the noun looking back. The maps are parameter matrices, and training determines their entries from examples.

The shapes are fixed by the number of entries in each vector. With input width $d$, every map is a matrix with $d$ rows and one column per output entry: the query and key maps have $d_k$ columns and the value map has $d_v$. The three maps hold $d(2d_k+d_v)$ parameters, which is three times $d\cdot d_k$ only in the paper's setting where the two widths agree [S:217:5]. That total is fixed once the widths are chosen, while the table of attention weights it produces grows with the square of the sequence length.

The score between a query and a key is their [[term:dot-product]]dot product, which multiplies matching entries and adds the results [S:74:224]. Two vectors pointing in the same direction score high, perpendicular vectors score zero, and opposing vectors score negative. That single number decides how much attention the item receives.

The scores become weights through the softmax function:

The line says each weight is the exponentiated score divided by the total of all the exponentiated scores [S:74:224]. Exponentiation raises the constant e, about 2.718, to each score: it turns a small lead into a larger one, it keeps every weight positive, and the division makes the weights sum to one. The output is the weighted sum of the values that the previous section described.

A discrete selection would keep the single best item and discard the rest. It is not differentiable: a small change in a score moves the choice not at all, or jumps it to another item, so no usable slope reaches the maps. Softmax keeps every item with a positive weight, so a small change in a score produces a small change in the output and the gradient can travel back into the maps [S:74:224].

Worked example. Suppose one query meets three keys and the dot products are 2, 1 and 0. Exponentiate: the score of 2 becomes e squared, about 7.389; the score of 1 becomes about 2.718; and the score of 0 becomes 1. The total is 11.107, so dividing each exponential by that total gives the weights.

The three weights add to one, so multiply each value by its weight and accumulate the products: 0.665 times 1 is 0.665, 0.245 times 2 is 0.49, and 0.09 times 3 is 0.27. The output is 1.425.

The largest similarity receives approximately two thirds of the weight, although it leads the second score by only one point. That sensitivity is advantageous, and it becomes a problem in long vectors. A dot product of vectors with many components grows with the quantity of components, so the scores become large and the softmax saturates, meaning one weight takes almost all the mass. One weight then sits near one and the rest near zero, so the gradient in that region is tiny and training stalls [S:217:4].

The remedy is a single division: divide every score by the square root of the key length before the softmax, and the operation is scaled dot-product attention [S:217:4]. With keys of length two, the divisor is about 1.414, so a raw score of 2 becomes 1.414. The competition stays moderate, and the weights stay responsive.

One set of maps is one head, and a layer runs several heads in parallel, concatenating their outputs, joining them end to end, so each head can specialize in a different relation [S:217:5]. The queries and keys of a head must have the same length, while the values may have a different one [S:74:224]. When a model predicts the next word, a mask sets some scores to negative infinity before the softmax, which prevents a position from consulting words that have not been produced yet [S:217:5].

When the queries come from one sequence and the keys and values from another, the operation is cross-attention. Translation uses it: the side that writes the output asks about the input sentence, and every position of the output can read every position of the input [S:217:5]. Self-attention is the special case where both sequences are the same.

The arithmetic has a cost: queries and keys must share a length $d_k$ [S:74:224]. Splitting the model into $h$ heads does not multiply that cost at fixed total width. With $d_k=d_v=d_{\text{model}}/h$ the heads hold $4d_{\text{model}}^2$ projection parameters between them, the same as one head of full width, and the paper notes that the total computational cost is similar [S:217:5]. Heads multiply the cost only against a single head of the same narrow width, and they add the output projection that recombines them.

A dot product measures alignment in whatever space the learned maps produce, and nothing in the formula guarantees that alignment means relevance. The maps have to learn a space, a coordinate system, in which a large dot product means relevance.

The formal statement. Three learned maps turn the position vectors into queries, keys and values, with widths $d_k$, $d_k$ and $d_v$. With one head,

$$\mathrm{attention}=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V,$$

where the softmax runs along the key index [S:217:4]. The line says each query scores every key, the scores become one row of weights, and the row averages the values. The hypotheses are: queries and keys share the width $d_k$, values may differ, the softmax normalises each row, and the divisor is the square root of $d_k$.

The derivation. Suppose the entries of a query and a key are independent, with mean zero and variance one. Their dot product is a sum of $d_k$ products, so its mean is zero and its variance is $d_k$ [S:217:4]. A typical score therefore grows like the square root of the width. Dividing by that square root gives a score of unit variance, whatever the width. The divisor is a variance normalisation, not a tuned constant.

The second fact follows from the first. The softmax is flat when the scores are close and sharp when one score leads. Sharpness is what the divisor controls, so the divisor controls how much of the gradient survives.

The edge case. The assumption fails after the first layer. The maps are trained, and their outputs need not have unit variance, so the score scale in a trained network is whatever the weights produce. A head with unusually large weights can saturate the softmax with the divisor in place: one weight near one, the rest near zero, and a gradient near zero. The divisor fixes the scale at initialisation. It does not enforce the scale afterwards.

A second edge case sits in the mask. An output position with no legal key gives every score negative infinity, and the softmax then divides zero by zero. Implementations guard that row or accept a useless output.

What the core hid. It stated that the heads share the total width, and left out that the divisor follows the head width. In the paper each head works in $d_{\text{model}}/h$ dimensions, so eight heads at width 64 hold the same total width as one head at 512 [S:217:5], and the divisor is the square root of 64. Two heads with different widths need different divisors.

Where the analogy breaks. A library label is written by a librarian and means something before the search begins. A key is a learned vector with no meaning outside the dot product, and the query is learned in the same space. Nothing guarantees that a large dot product means relevance. The pair of maps has to learn a space where it does.

Original sources. Scaled dot-product attention appears with its full definition: queries and keys of dimension $d_k$, values of dimension $d_v$, a dot product, a division by the square root of $d_k$, and a softmax [S:217:4]. The same page names additive attention as the older alternative, computes its compatibility with a small feedforward network, and reports that the two have similar theoretical complexity while the dot product is faster on matrix hardware [S:217:4]. Prince derives the same formula from the softmax and explains the names as inherited from information retrieval [S:74:224].

An open question. Which compatibility function should a head use? The paper reports that for small $d_k$ additive and dot-product attention perform similarly, and that additive attention outperforms the unscaled dot product for larger $d_k$. It suspects the cause, dot products of large magnitude pushing the softmax into a region of extremely small gradients, and applies the $1/\sqrt{d_k}$ scaling as the remedy [S:217:4]. It reports no measurement of the two mechanisms after the scaling, so the gap closing is the paper's rationale rather than its result. What the corpus does not settle is whether a learned bilinear score, with a matrix between the query and the key, buys anything at a fixed parameter and compute budget. Evidence: one matched comparison of additive, scaled dot-product and bilinear scores, swept over key width and sequence length on the same data, with the parameter count held equal.

A frontier result. Multi-head attention depends on this arithmetic and changes its cost. Each head projects into $d_{\text{model}}/h$ dimensions, runs attention there, and the outputs are concatenated and projected back. With eight heads the total cost stays near that of one full-width head [S:217:5]. The design assumes that relevance is not one relation, so one set of maps cannot serve.

What to read next. Read §3.2.1 and §3.2.2 of the transformer paper for the definition, the footnote and the head arithmetic [S:217:4][S:217:5]. Read Prince §12.3.2 and §12.3.3 for the scaling argument and for multiple heads [S:74:228]. Keep the footnote's variance calculation at hand, because it states why the divisor has the form it has.