Scaled Dot Product Attention — Why We Divide by sqrt(d_k)

The whole operation, with shapes attached

For one attention head over a sequence of N tokens, with head dimension d_k for queries and keys and d_v for values (almost always d_k = d_v), the operation is:

Q: [N, d_k]   K: [N, d_k]   V: [N, d_v]
S = Q K^T          → [N, N]     raw scores
S' = S / √d_k       → [N, N]     scaled
S'' = S' + M       → [N, N]     M is the additive mask
A = softmax(S'', axis=-1)  → [N, N]  rows sum to 1
O = A V            → [N, d_v]   output

Row i of S holds the affinity of query i against every key; row i of A is a probability distribution over the N positions; row i of O is that distribution’s weighted average of the value vectors. In a real implementation everything carries two leading axes, [B, h, N, d_k], and the two matmuls are batched over B·h independent problems. The softmax is strictly row-wise — keys never compete across queries — and that is why the whole thing parallelises so cleanly.

Where the √d_k comes from

Treat the components of a query and a key as independent, zero-mean, unit-variance draws — roughly true at initialisation. A score is a sum of d_k such products, and independent variances add:

s = Σ_r q_r k_r ,  r = 1..d_k
E[q_r k_r] = 0 ,  Var(q_r k_r) = 1
Var(s) = d_k    →   std(s) = √d_k

So the spread of the logits grows with head width, and the softmax sees a different input scale at d_k = 32 than at d_k = 128. Dividing by √d_k makes the score distribution width-invariant: one architectural knob stops leaking into another. Note why the exponent is one half and not one. Dividing by d_k would shrink logit spread as 1/√d_k, driving every attention distribution toward uniform as the model widens — the opposite failure. The exact factor that cancels what the sum introduced is √d_k, nothing else.

What saturation does to the gradient

The cost of unscaled logits shows up in the backward pass. Softmax has Jacobian ∂p/∂s = diag(p) − p p^T, so an upstream gradient g on the weights becomes

∂L/∂s_j = p_j (g_j − Σ_i p_i g_i)

Every term carries a factor of p_j. If one weight is 0.99999 and the rest share 10^-5, the small entries get gradients scaled by 10^-5 and the large one gets g_j minus an average it itself dominates — also near zero. The layer stops learning which key to attend to; it can only reinforce the key it already picked.

How close is that to reality? At d_k = 128 the raw logit spread is √128 ≈ 11.3, so a one-standard-deviation gap between the top two scores means a probability ratio of e^11.3 ≈ 8.2×10^4 — effectively one-hot from the first step. Scaled, the same gap becomes 1.0 and the ratio is e ≈ 2.72: a preference, not a verdict.

A worked example on four dimensions

Take d_k = 4, one query, three keys:

q  = [ 1, -1,  2, 0.5]
k_1 = [ 2,  0,  1, -1]   k_2 = [0, 1, -1, 2]   k_3 = [1, 1, 1, 1]

raw    s = [ 3.5, -2.0,  2.5]
scaled s/√4 = [1.75, -1.0, 1.25]

softmax(raw)    = [0.729, 0.003, 0.268]
softmax(scaled) = [0.599, 0.038, 0.363]

Both rows rank the keys identically — scaling is monotone, it never changes which key wins. What it changes is how much probability mass the losers keep, and therefore how much gradient they receive. The raw row gives key 2 a weight of 0.003; the scaled row gives it 0.038, twelve times more signal to work with. At d_k = 4 the effect is a nudge. Redo the same arithmetic with d_k = 128 and typical logits near ±11 and the unscaled row is numerically one-hot, with the losers' gradients rounded to zero rather than merely small.

Masking is arithmetic on scores, never on weights

Causal and padding masks are applied additively to the scaled scores, before the softmax — a lower-triangular M with 0 on allowed positions and a large negative constant elsewhere. Zeroing entries of A after the softmax is the classic wrong fix: the rows no longer sum to 1, so the output is a shrunken, unnormalised average whose magnitude depends on how many positions were masked.

Two traps follow from the choice of constant. In fp16 the largest finite value is 65504, so the popular -1e9 becomes -inf on cast — use finfo(dtype).min or a value like -1e4 that survives. And a row that is entirely masked (a fully padded query position) yields -inf − (-inf) = NaN under max-subtraction, which then poisons every downstream tensor. Guard those rows explicitly rather than hunting the NaN later.

Numerical stability: subtract the row max

Every correct softmax computes the shifted form, using the identity that a constant shift cancels between numerator and denominator:

m_i = max_j S''_ij
p_ij = exp(S''_ij − m_i) / Σ_j exp(S''_ij − m_i)

Because every exponent is now ≤ 0, exp can never overflow — it can only underflow to 0, which is harmless since the largest term is exactly 1 and the denominator is therefore at least 1. Without the shift the ceiling is real: exp overflows above ln(3.4×10^38) ≈ 88.7 in fp32 and above ln(65504) ≈ 11.1 in fp16 — a threshold raw d_k = 128 logits reach routinely. Max-subtraction, not scaling, is what prevents overflow; scaling is about gradients. This same two-pass structure is what FlashAttention makes single-pass with a running max and a rescaled running sum.

The cost model: FLOPs and bytes

Per head, the two matmuls dominate: QK^T costs 2N^2 d_k FLOPs and AV costs 2N^2 d_v; the scale, mask and softmax add only O(N^2) elementwise work. Summed over h heads with h·d_k = d_model, the attention core is 4N^2 d_model FLOPs per layer.

N=2048, d_model=768, h=12, d_k=64
QK^T, one head : 2 × 2048^2 × 64 = 5.4e8 FLOPs
core, 12 heads : ≈ 1.29e10 = 12.9 GFLOP / layer
scores in fp32 : 2048^2 × 4 B = 16.8 MB / head → 201 MB / layer

The four projections cost 8 N d_model^2 ≈ 9.7 GFLOP at the same settings, so the quadratic term overtakes the linear one at N = 2·d_model — 1536 tokens here. Below that, attention is not your bottleneck; above it, nothing else matters.

Implementation details that are free wins

Apply the scale to Q, not to S. Scaling the [N, d_k] query block costs N d_k multiplies against N^2 for the score matrix — a saving of N/d_k, or 32× at N = 2048, d_k = 64. Better still, fold 1/√d_k into W_Q once at load time and the scale costs nothing at all at inference.

Order matters: scale, then mask, then softmax. Masking first means the mask constant gets divided too, quietly weakening it. K is never physically transposed in good kernels — the matmul consumes it with swapped strides, so QK^T is one call, not a permute plus a call. And the [N, N] score matrix, not the weights, is the memory hog: it is allocated per head per layer, which is exactly the allocation fused kernels exist to avoid materialising.

Caveats, alternatives, and the CPU picture

Be honest about the derivation’s scope: it is an initialisation-time argument. After training, Q and K entries are neither unit-variance nor independent, and logit magnitudes drift — which is why large models sometimes add QK-norm (normalise queries and keys before the product, then apply a learned temperature) to enforce the property throughout training rather than assume it. muP goes further and scales attention logits by 1/d_k to make the whole training recipe transferable across widths. A fixed √d_k is the cheap default, not a theorem.

On CPU, the practical consequences are two. The N^2 score matrix blows past cache long before it blows past RAM, so blocked kernels that keep tiles resident beat naive full-matrix code by a wide margin. And never materialise a dense mask you could express as a loop bound: for causal attention, half the score matrix is arithmetic you compute and then discard.