If you've read the Transformer paper ("Attention Is All You Need"), you'll know the scaled dot-product attention is the core of the model. The authors mention they scale the attention scores by 1/√(d_k) to prevent the scores from becoming "extremely large," which would cause gradients to vanish during training. But why this specific value?
Imagine our queries (Q) and keys (K) are random vectors where each element has a mean of 0 and a variance of 1. This is a common and desirable starting point, meaning the values are generally centered around zero without being too spread out.
Now, the dot product between a query and a key is a sum of products. When you start adding up these products, the variance of the final dot product grows. Specifically, it grows proportionally to the dimension of the vectors, d_k.
Small d_k: The dot product values are relatively contained.
Large d_k: The dot product values can become huge in magnitude (both very positive and very negative).
This is a problem because these dot products are fed into a softmax function. The softmax function is very sensitive to extreme values. If one dot product is much larger than the others, the softmax will assign it a probability very close to 1, and all others will be nearly 0. This is called a "very sharp" distribution.
So, how do we stop the dot product from getting large as d_k increases? We need to scale it back down to a more manageable rang.
This is where 1/√(d_k) comes in. It's not a random choice; it's a mathematical correction. Since the variance of the dot product increases by a factor of d_k, we need to scale the values down by its square root, √(d_k), to bring the variance back to 1.
(The attached image provides the mathematical proof for this relationship.)
