Scaled dot-product attention & the √d_k

01 · Concept

Concept

Lesson 4.2 left one part of the attention formula unexplained: the mysterious division by $\sqrt{d_k}$. The concrete problem it solves appears the moment head dimension changes. Compare two hypothetical heads scoring the same token pair, one 4-dimensional and one 256-dimensional. Even if each individual feature carries the same amount of signal, the wide head’s dot product sums 64 times more terms, so its raw scores swing over a far larger range. Something downstream is about to be very sensitive to that range.

Write one score as

Under the intuition used in the original Transformer paper — components roughly independent, mean zero, variance one — each product $q_i k_i$ has variance near one, so the sum has $\operatorname{Var}(s) \approx d_k$ and standard deviation near $\sqrt{d_k}$. For Qwen3.8-27B’s attention heads, $d_k = 256$: typical raw scores are on the order of $\sqrt{256} = 16$ times larger than a single feature’s contribution, purely because of width.

Why does that matter? Softmax exponentiates its inputs, so scale is behavior. The naive first approach — feed raw scores straight into softmax — works acceptably for tiny heads and fails quietly for wide ones. Walk through it with three keys. Suppose the “true” preference pattern is mild: one key twice as compatible as the next. At small scale the raw scores might be $[0.5, 0.25, 0]$, and softmax spreads meaningful weight across all three. Now let the 256-wide head inflate everything by its natural factor of 16: the same pattern becomes $[8, 4, 0]$. Softmax of $[8, 4, 0]$ puts roughly 98% of the weight on the first key — near winner-take-all certainty the model never earned, and in the saturated region most gradients are nearly zero, so learning stalls exactly where it should be adjusting.

The correction is to divide by the standard deviation the width introduced:

With Qwen’s numbers, the inflated scores $[8, 4, 0]$ divided by $\sqrt{256} = 16$ come back to $[0.5, 0.25, 0]$ — the trainable regime. The divisor can also be read as a fixed softmax temperature of $T = \sqrt{d_k} = 16$: a scale chosen once from geometry rather than exposed as a knob. The model can still learn sharper or flatter attention by growing or shrinking query and key norms; the scaling only removes the automatic sharpening that width alone would cause.

Numerical stability is a separate, complementary fix. Practical softmax subtracts the row maximum before exponentiating, which prevents overflow without changing the resulting distribution. Max-subtraction solves representability; $\sqrt{d_k}$ solves saturation. You need both.

Now the second reason this lesson matters: the size of the thing being normalized. The score matrix has one entry per query–key pair. Qwen3.8-27B’s native context length is 262,144 tokens. At full context, one head of one layer scores

query–key pairs — about 69 billion score entries, before multiplying by the model’s 24 query heads, and before noting that 16 of its 64 layers run this computation. Nothing about the √256 division changes that count; scaling fixes the statistics of the scores, not their quantity. This is the quadratic bill from lesson 4.1 made concrete, and it is the standing motivation for track 7: FlashAttention-style kernels that avoid materializing the matrix, the KV cache whose arithmetic lands in lesson 7.2, and the grouped-query design from lesson 7.6 that Qwen already bakes into its 24-query/4-KV head layout. Qwen’s deeper answer — replacing attention entirely in 48 of 64 layers — arrives in lesson 4.15.

The lesson generalizes beyond one denominator. Whenever a sum collects many roughly independent contributions, its typical magnitude grows with the number of terms, and passing it into a sensitive nonlinearity changes behavior merely because width changed. The Transformer compensates with a scale derived from variance — a two-character fix in the formula, load-bearing at every width from a toy head to 256.