Formula & theorem reference

If \hat y = \operatorname{softmax}(z) and L = -\sum_i y_i \log \hat y_i for a one-hot label vector y, then the gradient of the loss with respect to the logits z is \frac{\partial L}{\partial z} = \hat y - y.

Take query and key vectors q, k \in \mathbb{R}^{d_k} with independent, zero-mean, unit-variance entries. Their dot product q \cdot k then has variance d_k. Dividing the score by \sqrt{d_k}, \operatorname{score}(q, k) = \frac{q \cdot k}{\sqrt{d_k}}, restores unit variance regardless of d_k. This keeps the softmax over attention scores away from the saturated, near-zero-gradient region that a growing d_k would otherwise push it into.