Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Scaled Dot-Product Attention › The Scaling Problem and Why It Matters

Michael BrenndoerferMay 25, 20258 min read

Language AI HandbookAll chapters

Part 13 of 67Self-Attention

Chapter 92Scaled Dot-Product Attention › The Scaling Problem and Why It Matters

Reading section 3 of 15 · The Scaling Problem and Why It Matters · 8 minSection 3 of 15 · 8 min

Section 3 of 15 · 8 min read

Scaled Dot-Product Attention › The Scaling Problem and Why It MattersOne section · Sections in this chapter3 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Before applying softmax to the score matrix, we must address a subtle but critical issue: as $d_k$ grows larger, the dot products tend to grow very large in magnitude, and softmax applied to large values has an extremely undesirable property. This is not a hypothetical edge case: it is a systematic failure mode that the original transformer paper identified and explicitly corrected with the $\sqrt{d_k}$ scaling factor. Without this fix, training transformers with large key dimensions would be dramatically slower and potentially unstable.

To see why dot products grow with dimension, assume that the components of $\mathbf{q}$ and $\mathbf{k}$ are independent random variables with mean 0 and variance 1. This is a reasonable initialization assumption for learned projection vectors. Then each term $q_l k_l$ has mean 0 and variance 1, since the product of two independent mean-zero unit-variance variables has variance equal to the product of their variances. The dot product is a sum of $d_k$ such terms:

\mathbf{q} \cdot \mathbf{k} = \sum_{l=1}^{d_k} q_l k_l

By properties of variance, the variance of this sum is $d_k$ , making the standard deviation $\sqrt{d_k}$ . So as $d_k$ grows, dot products have a spread proportional to $\sqrt{d_k}$ . A $d_k$ of 64 yields dot products with standard deviation around 8. A $d_k$ of 512 yields dot products with standard deviation around 22.6. These are large values to feed into softmax.

Variance of a Sum

If $X_1, X_2, \ldots, X_n$ are independent random variables each with variance $\sigma^2$ , then $\text{Var}(X_1 + X_2 + \cdots + X_n) = n\sigma^2$ , and the standard deviation is $\sigma\sqrt{n}$ . This is why dot products in high dimensions have large magnitudes: each additional dimension adds one more product term to the sum, and the sum's variance grows linearly with the number of terms.

Now consider what softmax does to large inputs. The softmax converts a vector of real numbers into a probability distribution. Given an input vector $\mathbf{z}$ with $n$ elements, the softmax output for element $i$ is:

\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_j e^{z_j}}

where:

$z_i$ : the $i$ -th score (the raw dot product value for one query-key pair)
$e^{z_i}$ : the exponential of that score, always positive
$\sum_j e^{z_j}$ : the sum of exponentials over all positions, serving as a normalizing constant

When one value in $\mathbf{z}$ is substantially larger than the others, the exponential amplifies the gap dramatically. For example, if $\mathbf{z} = [0, 0, 10]$ , the softmax is approximately $[0.0000226, 0.0000226, 0.9999548]$ : a near-one-hot vector that places almost all attention on the third position. Gradients through softmax in this regime are nearly zero everywhere except at the maximum, which means parameters receive almost no learning signal. The exponential function's property of amplifying differences between values is exactly what creates this problem: a difference of 10 units in the input becomes a factor of $e^{10} \approx 22000$ in the numerators, completely overwhelming all but the largest value.

This is the saturation problem: large inputs cause softmax to produce extremely peaked distributions, which kills gradient flow and makes training difficult. Think of it like a light switch stuck in the on position: the model collapses to always attending to one position, gradients cannot push it to distribute attention more broadly, and the mechanism loses its capacity to capture dependencies across multiple positions.

In[5]:

Code

def softmax(x, axis=-1):
    """Numerically stable softmax."""
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / e_x.sum(axis=axis, keepdims=True)


# Demonstrate saturation effect
z_small = np.array([1.0, 2.0, 0.5, 1.8])
z_large = np.array([10.0, 20.0, 5.0, 18.0])  # scaled up by 10

weights_small = softmax(z_small)
weights_large = softmax(z_large)

# Gradient proxy: softmax output * (1 - softmax output),
# larger for flatter distributions
grad_small = weights_small * (1 - weights_small)
grad_large = weights_large * (1 - weights_large)

Out[6]:

Console

Softmax of small inputs: [0.1527 0.415  0.0926 0.3398]
Softmax of large inputs: [0.     0.8808 0.     0.1192]

Gradient signal (small inputs): [0.1294 0.2428 0.084  0.2243]
Gradient signal (large inputs): [4.00000e-05 1.05021e-01 0.00000e+00 1.04990e-01]

Max gradient, small inputs: 0.2428
Max gradient, large inputs: 0.105021

The gradient signal is nearly zero when inputs are large. When $d_k$ is 512 (common in transformers), unscaled dot products would have standard deviation around 22.6, pushing softmax firmly into the saturation regime. The attention mechanism would effectively be a hard argmax from the very beginning of training, with no gradient signal to guide the model toward more distributed attention patterns.

Out[7]:

Visualization

Bar chart showing relatively uniform softmax distribution at scale 1. — Softmax at scale 1. The distribution is relatively flat, with attention spread across all four positions. Gradients are healthy at this scale, meaning all positions receive meaningful learning signal during backpropagation.

Bar chart showing moderately peaked softmax distribution at scale 2. — Softmax at scale 2. Sharpening is visible: the dominant input (v2) takes more weight, but other positions still receive meaningful attention. Gradients are beginning to weaken for low-weight positions.

Bar chart showing highly peaked softmax distribution at scale 4. — Softmax at scale 4. Most probability mass concentrates on the maximum. Gradient signal for non-maximum positions is weakening significantly, reducing the model's ability to learn from them.

Bar chart showing near-one-hot softmax distribution at scale 8. — Softmax at scale 8. Nearly all attention weight sits on the single maximum input, approaching a one-hot distribution. Gradients for all other positions are near zero, making learning difficult. The dashed line shows the uniform (0.25) baseline.

The Scaling FixLink Copied

The original "Attention Is All You Need" paper introduced a simple fix: divide the dot products by $\sqrt{d_k}$ before applying softmax. The purpose of this step is to counteract the variance growth we identified above. The scaled score matrix becomes:

\tilde{S} = \frac{QK^T}{\sqrt{d_k}}

where:

$QK^T \in \mathbb{R}^{n \times n}$ : the raw score matrix from matrix multiplication
$\sqrt{d_k}$ : the scaling factor, the square root of the query/key dimension
$\tilde{S} \in \mathbb{R}^{n \times n}$ : the scaled scores, ready for softmax

Why $\sqrt{d_k}$ specifically? Dividing by $\sqrt{d_k}$ normalizes the variance. Recall that the raw dot products have standard deviation $\sqrt{d_k}$ under our initialization assumption. After dividing by $\sqrt{d_k}$ , the standard deviation becomes 1, just like unit-variance inputs. This keeps the softmax operating in its sensitive, gradient-rich regime regardless of how large $d_k$ is. You can think of it as a standardization step: instead of letting the scale of the inputs to softmax depend on an architectural hyperparameter, we normalize them so the effective scale is always around 1.

Why does this formula make sense? Notice that scaling by $1/\sqrt{d_k}$ does not change the ranking of scores: if key $j$ had the highest raw score for query $i$ , it still has the highest scaled score. What changes is the sharpness of the resulting softmax distribution. With variance-normalized inputs, softmax produces distributions in the range where gradients are non-trivial, and the model can learn to make them sharper or flatter as needed by adjusting the directions of its query and key vectors. The scaling constant is simply a normalization factor, not a learnable parameter.

In[8]:

Code

# Compare raw vs scaled score statistics
scale = np.sqrt(d_k)
S_scaled = S / scale

raw_std = S.std()
scaled_std = S_scaled.std()

Out[9]:

Console

d_k = 4, scale factor = sqrt(4) = 2.000
Raw score std:    0.7420
Scaled score std: 0.3710

After scaling, the scores have a standard deviation close to 1, placing softmax in its operating range. The scaling factor is a constant; it doesn't change what the model learns, only the numerical scale of the inputs to softmax. This is a good example of a design pattern common in deep learning: a simple normalization step that costs almost nothing computationally but enables significantly more stable and effective training.