Why Divide by √d_k

What is Scaled Dot Product Attention?

In the Transformer architecture, attention is calculated using:

$$\text{Attention}(Q, K, V) = \text{Softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right) V$$

Why do we explicitly divide by $\sqrt{d_k}$?

Why not use simple unscaled dot product attention $\text{Softmax}(Q K^T) V$?

The Mathematical Cause: Exploding Dot Product Variance

Assume Query vector $q$ and Key vector $k$ are $d_k$ dimensional vectors whose components are independent random variables with mean 0 and variance 1.

The dot product is the sum of $d_k$ individual products:

$$q \cdot k = \sum_{i=1}^{d_k} q_i k_i$$

1. Mean of $q \cdot k$ = $0$.
2. Variance of $q \cdot k$ = $d_k$.

As key dimension size $d_k$ grows larger (for example, $d_k = 64$ or $d_k = 128$), the variance of the dot product grows linearly to 64 or 128.

This means dot product values frequently reach large positive or negative numbers like $+30$ or $-40$.

The Softmax Derivative Trap

Look at what happens when Softmax receives large input values:

  Unscaled Inputs to Softmax:    [ +35.0,  -20.0,  +2.0 ]
  Softmax Probabilities:         [ 0.999999,  0.000000,  0.000000 ]  (Extremely Sharp Peak!)

When Softmax outputs extreme probabilities near 1.0 and 0.0:

Softmax local derivative approaches ZERO.

During backpropagation, multiplying upstream gradients by near zero derivatives causes Vanishing Gradients, stopping the model from updating early weights.

The Scaling Fix: Dividing by $\sqrt{d_k}$

To restore variance back to 1.0, we divide the dot product by $\sqrt{d_k}$:

$$\text{Var}\left( \frac{q \cdot k}{\sqrt{d_k}} \right) = \frac{1}{d_k} \text{Var}(q \cdot k) = \frac{d_k}{d_k} = 1.0$$

Scaling by $\frac{1}{\sqrt{d_k}}$ keeps dot product magnitude stable around 1.0 regardless of how large key dimension $d_k$ grows.

This keeps Softmax inputs in a smooth, well calibrated region where gradients remain strong and healthy during training.

Say this out loud

Scaled Dot Product Attention divides Query Key dot products by the square root of key dimension size d_k. As vector dimension d_k grows large, dot product variance increases linearly to d_k, pushing Softmax into extreme regions where gradients vanish. Scaling by 1 / sqrt(d_k) restores variance to 1.0, keeping Softmax gradients healthy during backpropagation.

Followups to expect

1. What is Additive Attention (Bahdanau Attention)? Computes similarity using a feedforward layer v_a^T tanh(W_q q + W_k k). It handles large dimensions without scaling issues but is much slower on GPUs than Scaled Dot Product matrix multiplication.
2. What happens if d_k is very small (like d_k = 4)? Scaling has minimal effect because dot product variance is already small. Scaling becomes critical as d_k scales to 64, 128, or 256.