Vol. IV · Attention Plate 3 of 9scaled dot-product attention, softmax saturation, square root of d, head dimension, gradientsAL

Scaled Dot-Product Attention.

Divide by the square root of the dimension, and attention behaves the same at 16 dimensions as at 256.

If the entries of a query and a key are independent with variance one, their dot product is a sum of d products and has variance d: its typical size grows like √d. At d = 4 the scores are a few units apart; at d = 256 they are tens apart, and a softmax over scores that far apart puts almost all the weight on the largest one. A saturated softmax is bad for learning: its gradients with respect to the scores are close to zero, so the queries and keys stop being trained.

The Transformer paper’s fix is one line: divide the scores by √d before the softmax. The scaled scores have variance one whatever the width, so attention starts training in the same soft, learnable regime at any head size. That is the whole of scaled dot-product attention: softmax(QKᵀ/√d) V.

Below: a random query compared with 32 random keys. Left: the entropy of the attention weights as the dimension d grows, with and without dividing by √d (averaged over five draws). Right: the weights themselves at the chosen d.

What you will learn

Show that dot products of random vectors grow like √d.
Explain softmax saturation and why it stops learning.
Apply the 1/√d scale.
Write scaled dot-product attention in matrix form.

Workshop

Left: entropy of the attention weights over 32 keys (5 bits would be perfectly even) against the dimension d, scaled by 1/√d (blue) and unscaled (red). Right: the 32 weights for one draw at the chosen d.

One square root

At d = 128 without scaling, the 32 scores spread by about 9.6 units and the attention weights have an average entropy of 0.29 bits — nearly all the weight on one or two keys, a softmax so saturated that its gradients barely move. Divide by √128 and the spread is about one, the entropy 4.27 bits, close to the 5 bits of perfectly even weights: the starting point from which a model can learn to sharpen attention where it helps. At d = 4 the two hardly differ (unscaled 3.28 bits), which is why the problem only appeared when attention was used with wide vectors.

In a Transformer, d here is the head dimension — the size of each head’s queries and keys, usually 64 or 128 — not the model’s full width (plate IV.5). The scale is not the only knob: trained models learn to make some heads sharp on purpose, and later architectures add normalisation of queries and keys (QK-norm, plate VII.7) to keep the scores from growing during training as well as at the start. The formula softmax(QKᵀ/√d)V, with Q, K and V holding every position’s queries, keys and values as rows, computes all of a sequence’s attention in two matrix products.

The Transformer paper explained the √d with exactly this variance argument, in a footnote; additive attention, the 2014 alternative, did not need it but was slower on matrix hardware.

Terms

Scaled dot-product attention: softmax(QKᵀ/√d) V: scores divided by the square root of the key dimension.
Softmax saturation: Scores so far apart that the softmax is nearly one-hot and its gradients nearly zero.
Head dimension: The size of each attention head’s queries and keys; the d in √d.
Entropy of attention: How spread the attention weights are, in bits: high when even, zero when one-hot.

Random dot products grow like √d; dividing by √d keeps attention soft and trainable at any head size.