Queries, keys, values and scaled dot-product attention

1. Dot-product self-attention through scaling

   Understanding Deep Learning - Transformers

   Sections 12.2 to 12.3.2 (pp. 208 to 214): value, query and key computation in equations 12.2 to 12.4, attention weights in 12.5, matrix form in 12.7 and 12.8, and the scaling argument with equation 12.9

   Done
2. Notebook 12.1: self-attention

   Understanding Deep Learning - Transformers

   The complete notebook, referenced in the margin of the matrix-form derivation

   Done
3. Queries, keys and values

   D2L - attention and transformer implementation route

   Section 11.1: the attention-pooling equation, the nonnegative, convex-combination, single-selection and average-pooling cases, and show_heatmaps

   Done
4. Dot-product attention and scaling

   D2L - attention and transformer implementation route

   Section 11.3.1 including the variance-\(d\) derivation of the scaled scoring function, then the Scaled Dot Product Attention class with the (2, 1, 2) query toy and heatmap check

   Depth · Skip Additive Attention; it is not a unit topic. Read the batch matrix multiplication part of Section 11.3.2 for shapes only.

   Done
5. Equivalent-implementation box

   PyTorch - scaled dot product attention

   The Efficient implementation equivalent code box: default scale \(1/\sqrt{E}\), score-times-scale order, and softmax-then-dropout order

   Done

Focus on

• Queries, keys and values are three learned linear maps of the same input. D2L builds them bias-free by default while UDL equations 12.4 and 12.6 carry bias vectors, so the two parameter counts differ by construction.
• Under independent zero-mean unit-variance components, \(q^\top k\) has variance \(d\). The \(1/\sqrt{d}\) factor restores unit variance before the softmax; without it large \(d\) saturates the softmax and gradients vanish.
• Stable softmax subtracts the row max first. Naive fp32 softmax on the UDL Problem 12.4 logits overflows at \(\exp(100)\): that entry becomes inf/inf, which is NaN, and the rest become 0. The shifted form returns the correct weights.
• UDL lays inputs out as \(D \times N\) columns with softmax per column; D2L uses batch by queries by keys with softmax on the last axis. In both, exactly one axis sums to 1: check which one before reading any heatmap.

Reproduce closed-book

Traps

• Computing softmax as \(\exp(z)/\sum(\exp(z))\) in fp32. \(\exp(100)\) overflows to inf, so the sum is inf, the large entry becomes NaN and every other entry 0. Always subtract the row max, and recheck any code that skips it.
• Softmaxing the wrong axis after a transpose. UDL and D2L sum over different axes by layout; a heatmap that sums to 1 along queries instead of keys inverts the routing story.
• Quoting D2L parameter counts for UDL equations or vice versa. The bias terms add \(3D\) parameters; carry them only when the equations you implement include them.

Practice

Self-test

1. Query and key components are independent with zero mean and unit variance, in dimension \(d\). What are the mean and variance of their dot product, and what scale restores unit variance?

   Show answer Got it Missed

2. Naive fp32 softmax meets UDL Problem 12.4 logits [-3, 1, 100, 5, -1]. What comes out, what does the max-subtracted form give, and what do the 25 derivatives say?

   Show answer Got it Missed

3. A self-attention block takes \(N = 5\) inputs of width \(D = 8\), with queries, keys and values also width 8. How many QKV weights and biases, how many attention weights, and how many weights and biases in a fully connected \(DN\) to \(DN\) map?

   Show answer Got it Missed

4. Why must attention weights be nonnegative and sum to one, and where does each output lie?

   Show answer Got it Missed

Leads to

• 7.2 Causal masks, padding masks and fully masked rows Feed masked scores into the same row-wise softmax.
• 7.3 Self-attention, cross-attention and multi-head shapes Reuse single-head attention inside per-head projections.

Resources for this unit

• Understanding Deep Learning - Transformers
  Simon J. D. Princeprimarytheory
  Use Chapter 12.1-12.11 and the matching notebooks for dot-product self-attention, extension to transformers, transformer layers and vision transformers.
• D2L - attention and transformer implementation route
  D2L authorsrequiredimplementation
  Use queries-keys-values, attention-scoring-functions, multihead-attention, self-attention-and-positional-encoding, transformer and vision-transformer. Implement masked single-head attention, multi-head attention and one pre-normalized block before comparing with PyTorch.
• PyTorch - scaled dot product attention
  PyTorchreferencereference
  Use the official function documentation as a differential oracle for shapes, masks, dropout and backend behavior after the manual version passes intervention tests.