Why does attention scale the dot product before softmax?

To score how well a query matches a key, the transformer takes their dot product, divides by the square root of the key’s dimension, and applies a softmax to turn the scores into weights that sum to one. The division is not a detail. The paper explains that “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,” so the scaling keeps learning stable.^[1]

The paper spells out exactly why the dot products get large. Treat a query’s and a key’s individual components as independent random values with mean 0 and variance 1: their dot product is a sum of d_k such terms, and a sum of d_k independent terms has variance equal to d_k itself.^[1] So as the key dimension grows, the typical size of an unscaled score grows right along with it, with no ceiling. Dividing by the square root of d_k cancels that growth exactly, pulling the variance back to roughly 1 no matter the dimension, which is why the fix is specifically a square root and not some fixed constant: an 8-head model with d_k of 64 needs a different scale than a model built with a different head count, and dividing by the square root of d_k adjusts automatically for either.

The cost of skipping this step is not cosmetic. Once a handful of scores run far larger than the rest, softmax pushes almost all of the probability mass onto them and assigns the remaining scores a gradient close to zero. A model in that state stops getting a useful learning signal for telling those low-scoring positions apart, in exactly the high-dimension regime every modern model operates in.