In the Transformer architecture introduced by Vaswani et al. (2017) in "Attention Is All You Need," the dot products between queries and keys are divided by the square root of the key dimension before the softmax is applied. What is the primary reason for this scaling step?
- It converts the attention weights into a probability distribution that sums to one across the sequence
- It reduces the number of parameters needed in the query, key, and value projection matrices
- It counteracts dot products growing large in magnitude for larger key dimensions, which would push the softmax into regions with extremely small gradients
- It allows the same attention weights to be reused across all heads in multi-head attention
Why C? And why not the others?
Correct answer: C. It counteracts dot products growing large in magnitude for larger key dimensions, which would push the softmax into regions with extremely small gradients
Vaswani et al. (2017) scale the raw dot products by the square root of the key dimension because, for larger key dimensions, the dot products between query and key vectors tend to grow large in magnitude; feeding these large values into the softmax pushes it into saturated regions where the gradient with respect to its inputs becomes extremely small, hurting learning. Dividing by the square root of the dimension keeps the values in a range where the softmax remains well-behaved. The option about producing a probability distribution describes what the softmax function itself does, not what the scaling step contributes before the softmax is applied. The option about reducing parameter counts is wrong because scaling is a fixed arithmetic operation on activations, not a change to the learned projection matrices. The option about reusing weights across heads is wrong because each attention head computes its own independent set of attention weights from its own query and key projections; scaling does not link heads together in any way.
Source: Vaswani et al., "Attention Is All You Need" (2017), arXiv:1706.03762, Section 3.2.1 (Scaled Dot-Product Attention)