Scaled dot-product attention adds exactly one modification to Dot-Product Attention — dividing scores by \(\sqrt{d_k}\) before the softmax — and this single change is what the Transformer architecture actually uses throughout.
Formula
\(d_k\) is the dimensionality of the query/key vectors. This scaling factor is the entire difference from the previous note's formula.
Why the Scaling Is Necessary
As \(d_k\) grows, the dot product between two random vectors tends to grow in magnitude too (each of the \(d_k\) individual products contributes to the sum). For large \(d_k\) (in real Transformers, often 64 or more per attention head), raw dot-product scores can become quite large in magnitude. Recall from Softmax Function and Sigmoid Function that very large input magnitudes push softmax (like sigmoid) into a saturating region — where the output becomes extremely close to one-hot (nearly all weight on the single largest score) and, critically, the gradient becomes vanishingly small. This is exactly the same saturation-driven vanishing gradient mechanism from Vanishing Gradient Problem, just showing up inside the attention computation specifically.
Why Dividing by \(\sqrt{d_k}\) Specifically
Assuming query and key components are roughly independent with unit variance, the dot product's variance grows proportionally to \(d_k\) — so its standard deviation grows proportionally to \(\sqrt{d_k}\). Dividing by \(\sqrt{d_k}\) exactly counteracts this growth, keeping the scores' typical magnitude roughly constant regardless of how large \(d_k\) is — keeping softmax operating in its well-behaved, non-saturating region, with healthy gradients, no matter the dimensionality chosen.
Numerical Example
Suppose \(d_k=64\) and two typical (unscaled) dot-product scores are \(45\) and \(2\) — a huge, saturation-inducing gap. Scaled: \(\frac{45}{\sqrt{64}}=\frac{45}{8}=5.625\) and \(\frac{2}{8}=0.25\) — still meaningfully different, preserving relative ranking, but on a much gentler scale that keeps softmax's gradient healthy.
Code
import torch
import torch.nn.functional as F
import math
def scaled_dot_product_attention(Q, K, V):
d_k = Q.shape[-1]
scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)
weights = F.softmax(scores, dim=-1)
return weights @ V
Q = torch.randn(1, 5, 64) # 5 queries, d_k=64
K = torch.randn(1, 5, 64)
V = torch.randn(1, 5, 64)
output = scaled_dot_product_attention(Q, K, V)
print(output.shape) # (1, 5, 64)Common Mistakes
- Forgetting the scaling factor when implementing attention from scratch — with small \(d_k\), this can be barely noticeable; with the larger \(d_k\) values common in real Transformer models, omitting it can measurably hurt training stability and convergence.
- Dividing by \(d_k\) instead of \(\sqrt{d_k}\) — this specific square root is what correctly counteracts the variance growth of a sum of \(d_k\) independent products; the wrong power either over- or under-corrects.
Interview Relevance
Q: "Why does the Transformer's attention mechanism divide scores by \(\sqrt{d_k}\) before applying softmax?" Without scaling, dot-product scores grow in typical magnitude as \(d_k\) increases (since the dot product sums \(d_k\) individual products), which pushes softmax into a saturating region — producing near one-hot output and, critically, very small gradients. Dividing by \(\sqrt{d_k}\) counteracts this variance growth, keeping softmax's input in a well-behaved range with healthy gradients regardless of the chosen key/query dimensionality.
Practice Question
If \(d_k=16\), what is \(\sqrt{d_k}\), and by how much would a raw score of 40 be scaled down?
