Query, Key, Value — Dot-Product Attention from Scratch

The last read ended on a formula the non-local paper wrote in one line, $y=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(x^{\top }{W}_{\theta }^{\top }{W}_{\varphi }x)g(x)$, and called "the self-attention form" of the transformer [WANG17]. This read takes that line apart. The definition we are after is the paper's own, §3.2 of [VASW17]: "An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key."

Three steps, the same three as before: score (query · key), normalise (softmax), average (weights × values). Everything in this read is a name, a scale or a repetition of those three steps.

Brush up on Attention Before Transformers — Score, Squash, Multiply first, Eq. 3 of it in particular: Bahdanau's three lines are reused below for a side-by-side. Softmax itself I assume you know; what it does to a dot product you may not, and that is Section 4.

1. Three Names for the Same Input

The paper's vocabulary is a lookup table. A dictionary takes a query, finds the one key that matches exactly, and returns its value. Attention is the soft version: every key matches a little, and what comes back is a blend of all the values, weighted by how well each key matched. The new idea compared with the last read is where the three vectors come from. In self-attention they all come from the same input $X\in {ℝ}^{n\times d_{model}}$, one row per token, through three different learned matrices:

• Query $q_i$, one row of $Q$: what token $i$ is looking for.
• Key $k_j$, one row of $K$: what token $j$ advertises that it contains. Same width $d_k$ as the query, so the two can be dotted.
• Value $v_j$, one row of $V$: what is actually copied out if $j$ is chosen. Width $d_v$, which need not equal $d_k$.

In the last read I called these "the thing being compared" and "the things being averaged". The trick that turns two things into three is that the token being averaged advertises itself with one projection (the key) and delivers itself with another (the value).

NOTE: $W^Q\ne W^K$. If they were the same matrix, $Q{K}^{\top }$ would be symmetric: "it" would match "cat" exactly as much as "cat" matches "it", and a token could not ask for something it is not. The two projections are what let "it" look for "cat" without "cat" looking back for "it".

2. One Token, Step by Step

Let's devise it for a single query $i$ against $m$ keys ($m=n$ in self-attention):

• Score. A dot product between the query and each key, divided by $\sqrt{d_k}$. Why the division is Section 4; park it.
• Normalise. Softmax across the $m$ scores: the weights are positive and sum to 1.
• Average. The output is the weighted average of the value rows. Nothing else happens.

Now put Bahdanau's version from the last read next to it (Eqs. 5–6 of [BAHD14]):

Lines two and three are identical. Only the first line changed: a one-hidden-layer MLP became a dot product, as Luong had already tried [LUON15], plus a scale. That is the whole story of "from additive to dot-product attention". The paper's reason for preferring the dot product is practical: the two "are similar in theoretical complexity", but the dot product "is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code" [VASW17].

Let's watch it on eight tokens. Pick a query row and turn the temperature:

query token (0 = the … 7 = it)

7

temperature τ (× √d_k)

1

In the figure, the row for "it" puts its largest weight on "cat": the pronoun's query matched the animate noun's key, because I built the toy projections to do so. Change the query index and read the other rows; turn the temperature down and every row becomes an arg-max, turn it up and every row flattens to 1/8. That knob is exactly what Section 4 is about.

3. All Tokens at Once

Stack the $n$ queries as rows and the three lines of Eq. 2 become Eq. 1 of the paper:

• $Q{K}^{\top }\in {ℝ}^{n\times m}$ holds every query's score against every key; row $i$ is the $e_{ij}$ of Eq. 2.
• The softmax runs per row, so each row of the result is a probability over the keys.
• Multiplying by $V\in {ℝ}^{m\times d_v}$ averages the values, one output row per query: the result is $n\times d_v$.

Two matrix multiplications and a row-wise softmax. The non-local block of the last read is this exact formula on pixels, minus the $\sqrt{d_k}$, with $W_{\theta },W_{\varphi },W_g$ playing $W^Q,W^K,W^V$ [WANG17].

4. Why Divide by √d_k?

The paper says it in one sentence: "for large values of $d_k$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients" [VASW17]. Footnote 4 gives the reason, and it deserves to be written out. Assume the components of $q$ and $k$ are independent with mean 0 and variance 1:

• Line 1: a score is a sum of $d_k$ terms.
• Lines 2–3: each term has mean 0 and variance 1 (independent factors).
• Line 4: variances of independent terms add, so the score's spread grows like $\sqrt{d_k}$.
• Line 5: dividing by $\sqrt{d_k}$ puts the spread back to 1 whatever the width.

Why does spread matter? Softmax is exp-then-normalise, so it exaggerates differences before dividing, and its derivative for a weight $p$ is $p(1-p)$. Take $d_k=64$, where a raw score has std 8:

• Two keys, raw scores 8 and 0 → softmax (0.9997, 0.0003). Scaled by 8 → scores 1 and 0 → (0.731, 0.269).
• Four keys, raw (16, 0, 0, 0) → (0.9999997, 10⁻⁷, …). Scaled → (2, 0, 0, 0) → (0.711, 0.096, 0.096, 0.096).
• The gradient factor $p(1-p)$ at the top weight of the two-key case: 0.0003 raw against 0.197 scaled, almost 600× larger after scaling. That is the "extremely small gradients" sentence in numbers.

Now let the figure run the same experiment at every width, with random queries and keys:

keys m

10

marker at d_k = 2^k

6

Unscaled, the mean largest weight climbs from 0.49 at $d_k=4$ to 0.86 at 64 and 0.97 at 1024 while the entropy collapses towards zero: the layer turns into a hard arg-max before any learning has happened. Scaled, both curves are flat at every width (0.32 and 1.93 nats for 10 keys). Temperature, if you like, is a volume knob between "ask everyone equally" and "ask only the loudest"; $\sqrt{d_k}$ sets the knob to the same place regardless of how many dimensions are shouting.

NOTE: Footnote 4 assumes independent, zero-mean, unit-variance components. That is true at initialisation and false after training, so $\sqrt{d_k}$ is the right scale to start from, not a law. Also, $d_k$ is the per-head key width, 64 in the base model, not $d_{model}=512$.

The paper cites an earlier study (its ref. [3]) for the observation that additive attention beats unscaled dot-product attention at large $d_k$ [VASW17]; the scale is what closes the gap. One more thing the paper is careful about: its author notes credit Noam Shazeer with the scaled version and with multi-head attention; the dot product itself goes back to Luong [LUON15].

5. Multi-Head Attention

One head produces one $n\times n$ weighting and then averages the values under it. The paper's complaint (§2) is that a single head suffers "reduced effective resolution due to averaging attention-weighted positions": if "it" should look at "cat" for one reason and at "sat" for another, one softmax has to split the difference. Multi-head attention runs the formula several times in parallel, each time in a narrower slice with its own projections, so that each copy can hold a different opinion of who is relevant to whom. The paper's words: it lets the model "jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this." [VASW17]. Eight people reading the same sentence with different questions in mind, then a meeting where their notes are merged:

• Each head is Eq. 4 in a 64-wide slice, with its own three projection matrices.
• The eight 64-wide outputs are concatenated back to 512 and mixed by $W^O$; that is the meeting.
• What one head computes: an $n\times n$ weight matrix and an $n\times 64$ output. Eight heads give eight $n\times n$ pictures of the same sentence.

Let's split the toy layer of Section 2 into heads. Change h and compare the slices with the full-width head on the left:

heads h = 2^k

2

With h = 4 the slices disagree on purpose: one matches identical words, one sends "it" to "cat", one sends verbs and prepositions to nouns. The full-width head's dot products are the sums of the slices' dot products, so it can only show their average, which is the paper's complaint in one picture.

5.1 The Numbers

• Parameters do not grow with h. 8 heads × 3 projections × (512 × 64) = 786,432, exactly 3 × 512² for one full-width head; plus $W^O$ at 512² = 262,144. Total 1,048,576 = $4d_{model}^2$ (bias-free), for any $h$. The paper: "due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality" [VASW17].
• More heads is not monotonically better. Table 3 row (A) of the paper, dev BLEU on newstest2013: h = 1 → 24.9, 4 → 25.5, 8 → 25.8, 16 → 25.8, 32 → 25.4. "Single-head attention is 0.9 BLEU worse than the best setting, quality also drops off with too many heads."
• Narrow keys hurt. Row (B): $d_k$ = 16 → 25.1, 32 → 25.4, 64 → 25.8. "Reducing the attention key size $d_k$ hurts model quality."

NOTE: "Multi-head = more parameters" is the common misreading. Same count for any h; what changes is how many independent n × n weightings the layer can form at once.

6. What It Costs

Count the multiply-adds in Eq. 4 for $n$ queries and $m$ keys (Table 1 of the paper, first row):

• The pairs cost $n^{2}d$ in compute and $h{n}^2$ numbers in memory. The projections cost $4n{d}^2$ and do not depend on pairs at all.
• A ViT-sized input, n = 196 tokens, d = 768 (the number An Image Is Worth 16×16 Words — ViT, Patch by Patch inherits): $Q{K}^{\top }$ is 196² × 768 = 29.5 M multiply-adds, $AV$ another 29.5 M, and the four projections 4 × 196 × 768² = 462 M. At n = 196 the projections cost 8× more than the pairs. The weight matrix is 196² = 38,416 numbers per head, × 12 heads = 1.8 MB in fp32 per image per layer.
• At pixel resolution, n = 224² = 50,176, the same two lines give the 10 GB per head of the last read; Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks starts from that number.

Table. 1 The first three rows of Table 1 in [VASW17], with the per-layer term evaluated at two sequence lengths.

Self-attention is cheaper than a recurrent layer "when the sequence length n is smaller than the representation dimensionality d" (§4 of [VASW17]), and a convolution costs k times a recurrent layer again; that holds for sentences and for 14 × 14 patch grids. At 50,176 tokens the order flips completely.

NOTE: Attention has O(n²) activations and compute, not O(n²) parameters. The parameters are 4d², independent of n.

7. The Attention Map as a Picture

Every head's $\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(Q{K}^{\top }/\sqrt{d_k})$ is an $n\times n$ matrix of numbers between 0 and 1 whose rows sum to 1, so it can be drawn as a heat map, exactly like Bahdanau's alignment matrices. The paper's appendix (Figs. 3–5) draws several heads of layer 5 of the encoder. Heads attend from "making" to "more difficult", completing the phrase making … more difficult; two other heads resolve its to the noun it refers to. The toy figures above imitate this: the row for "it" lands on "cat".

NOTE: The map is the weights, not an explanation. The appendix captions say the heads are "apparently involved in" and "seem related to" those tasks. Keep the paper's hedge when you read one.

8. Where the Three Roles Come From

§3.2.3 of the paper uses the same formula in three places, and the difference is only where $Q$, $K$ and $V$ are taken from:

• Encoder self-attention. $Q,K,V$ all from the previous encoder layer: every token looks at every token.
• Decoder self-attention. Same, but a token may not look at tokens after it; how those keys are hidden is in Attention Is All You Need — The Transformer, Block by Block.
• Encoder-decoder attention. $Q$ from the decoder, $K$ and $V$ from the encoder output. That is Bahdanau's cross-attention of the last read, written in Q/K/V form.

So cross-attention between two sequences and self-attention within one, the two cases the last read kept apart, are the same formula with different inputs. The formula does not know the difference.

9. What Came Next

The paper did not stop at the formula. It built a whole network out of it: eight heads, a small feed-forward layer per token, a residual connection and a normalisation around each, stacked six times in an encoder and six in a decoder, plus a trick to tell the model where each word sits, since nothing in Eq. 4 knows about order. That network, its numbers and its training recipe are Attention Is All You Need — The Transformer, Block by Block.

10. References

1. [VASW17] Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, Polosukhin. Attention Is All You Need. NIPS 2017.
2. [BAHD14] Bahdanau, Cho, Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015.
3. [LUON15] Luong, Pham, Manning. Effective Approaches to Attention-based Neural Machine Translation. EMNLP 2015.
4. [WANG17] Wang, Girshick, Gupta, He. Non-local Neural Networks. CVPR 2018.