# Why use a "square root" in the scaled dot product

In attention settings, typically when the both Query Q and Key K are of the same dimension d we can compute their attention score in the following manner:

$$\frac{Q^T K}{\sqrt{d}}$$

The justification being that the dot product tends to blow up in higher dimension, and $$\sqrt{d}$$ helps to "keep things reasonable". I also read that it serves to reduce the standard deviation to 1, however chatGPT contradicts that statement, and I can't find resources that support this claim.

Therefore what is the mathematical justification behind $$\sqrt{d}$$ and not another form such as $$\log d$$, $$d^5$$, $$e^{\pi d}$$ or simply $$|Q||V|$$ (a cosine similarity!) if we want to keep things "bounded"?

In the official paper the justification for this scaling is given in footnote $$[4]$$ on page 4.

Here is a more lengthy explanation of the same thing:
Let’s assume that the keys $$K$$ and queries $$Q$$ have zero mean and unit std. The attention score between some query $$i$$ and some key $$j$$ is: $$\alpha = q_i k_j^T = \sum_{n=1}^{d_k} q_{in}k_{jn}$$ We know that the variance of a sum of random variables equals the sum of the variances of those variables. Thus, for the variance of the attention score we get: $$\text{Var}(\alpha) = d_k,$$ since $$\text{Var}(q_{in}k_{jn}) = 1 \quad \forall n.$$

The standard deviation is: $$\text{std}(\alpha) = \sqrt{d_k}$$

After you compute the attention scores you need to apply softmax in order to obtain the attention probabilities. The problem is that the softmax function saturates very quickly. This means that if one of the attention scores is slightly bigger than the others, then the softmax function will put almost all of the weight on that element. See for example this:

x = torch.FloatTensor([1, 1, 1, 5])
y = torch.softmax(x, dim=-1)
# y: [0.0174, 0.0174, 0.0174, 0.9479]


Since the variance of the attention scores is $$d_k$$ then almost certainly one random score will be much higher than the others. The softmax will output attention probability of $$1.$$ for that element and zeros for the others and you will get no gradient propagation.

What you want to do to ensure stable gradient flow is normalize your attention scores. That is, make them with zero mean and unit std. We already assumed that the keys and queries have zero mean and unit std, so the only thing you need to do in order to achieve this is to scale the attention scores with $$\sqrt{d_k}$$.

But is it safe to assume that our keys and queries have unit variance?
Well, yes! The architecture of the Transformer includes LayerNorm layers that are used exactly for this reason: normalizing the inputs before applying the self-attention layer. In addition the weights of the self-attention layer are initialized using Xavier initialization and so the computed keys, queries and values will also have unit std. In my opinion this scaled dot-product works so well only when combined with LayerNorm and Xavier initialization. I have written a very detailed blog post about the Transformer architecture, if you want to read it here is a link:
https://pi-tau.github.io/posts/transformer

The answer by @pi-tau is upto mark. Here is a TL-DR version of it.

Assuming the keys and queries have mean of $$0$$ and standard deviation of $$1$$, the variance of the attention score (the sum of the dot products) is $$d$$, and the standard deviation is $$\sqrt{d}$$. To ensure stable gradient flow, you want to normalize your attention scores, i.e., make them have mean of $$0$$ and standard deviation of $$1$$. Since we already assumed that the keys and queries have $$0$$ mean and unit standard deviation, the only thing you need to do to achieve this is to scale the attention scores with $$\sqrt{d}$$.

This scaling helps to maintain stable gradients and prevents the softmax function from saturating and becoming unresponsive to small changes in its input. Hence, it's a crucial part of the self-attention mechanism in Transformer models.

From the attention is all you need paper (section 3.2.1):

Dot-product attention is identical to our algorithm, except for the scaling factor of $$\frac{1}{\sqrt{d_k}}$$. [...] While for small values of $$d_k$$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling [...]

Then

We suspect that for large values of $$d_k$$ the dot product grow large in magnitude, pushing the softmax function into region where it has extremely small gradients. To counteract this effect, we scale the dot products by $$\frac{1}{\sqrt{d_k}}$$.

As you can read the scaling is to fight vanishing gradients, in fact you can interpret it as a sort of temperature that smooths the logits $$QK^\top$$ making them more uniformly distributed (since a peaked distribution would cause gradient instabilities), therefore increasing the gradient magnitude.

• The motivation of this particular expression and not another one (like the product of the norm $\frac{1}{|Q| |K|}$ or even $1/d$) is not very well explained Commented Aug 24, 2023 at 16:34