Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Scaled Dot-Product Attention › Numerical Stability of Scaling

Michael BrenndoerferMay 25, 20253 min read

Language AI HandbookAll chapters

Part 13 of 67Self-Attention

Chapter 92Scaled Dot-Product Attention › Numerical Stability of Scaling

Reading section 9 of 15 · Numerical Stability of Scaling · 3 minSection 9 of 15 · 3 min

Section 9 of 15 · 3 min read

Scaled Dot-Product Attention › Numerical Stability of ScalingOne section · Sections in this chapter9 / 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

The $\sqrt{d_k}$ scaling prevents saturation at the start of training. It is also essential for maintaining numerical stability throughout training as model weights adapt. This distinction is worth dwelling on, because it explains why the scaling is a permanent architectural feature rather than just an initialization trick.

Consider what can go wrong without scaling. As the model trains, projection matrices $W^Q$ and $W^K$ can grow in magnitude to produce peaked attention distributions. Gradient descent does not naturally constrain the scale of learned parameters: if slightly larger weights consistently reduce the loss, the optimizer will keep increasing them. Once queries and keys have large dot products, softmax saturates, gradients vanish, and parameters stop updating. The model gets stuck with sharp, nearly-fixed attention patterns that cannot adapt as training proceeds. The training curve flattens, and the model is trapped in a suboptimal configuration.

With $\sqrt{d_k}$ scaling, the effective inputs to softmax remain in a numerically stable range. The model can still learn to produce peaked attention (by making queries and keys point in similar directions), but the gradients remain non-zero throughout, enabling continued learning. Think of the scaling as maintaining an implicit contract: the softmax always receives inputs with unit-scale variance, regardless of how the learned weights have grown. This contract is what keeps training stable over thousands of gradient updates.

There is a secondary numerical consideration: the softmax computation itself. A naive implementation computes $e^{z_i}$ for each score $z_i$ . Even after scaling, some scores may be positive enough that $e^{z_i}$ overflows floating-point representation. The standard fix is the log-sum-exp trick: subtract the maximum score before exponentiating.

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

where:

$z_i$ : the $i$ -th score value
$z_{\max} = \max_j z_j$ : the maximum score value in the row
$e^{z_i - z_{\max}}$ : the shifted exponential, at most $e^0 = 1$ , preventing overflow
$\sum_j e^{z_j - z_{\max}}$ : the shifted normalization constant

Subtracting the max shifts the largest value to zero, preventing overflow while producing identical results mathematically (the max cancels in both numerator and denominator). Why does this formula make sense? Notice that $\frac{e^{z_i - z_{\max}}}{\sum_j e^{z_j - z_{\max}}} = \frac{e^{z_i} / e^{z_{\max}}}{\sum_j e^{z_j} / e^{z_{\max}}} = \frac{e^{z_i}}{\sum_j e^{z_j}}$ , which is exactly the original softmax. The shift is mathematically neutral but numerically essential.

In[28]:

Code

def stable_softmax(z, axis=-1):
    """Numerically stable softmax using log-sum-exp trick."""
    z_max = np.max(z, axis=axis, keepdims=True)
    e_z = np.exp(z - z_max)
    return e_z / e_z.sum(axis=axis, keepdims=True)


# Demonstrate: naive softmax can overflow; stable version doesn't
z_extreme = np.array([1000.0, 1001.0, 999.0])

naive_result = (
    np.exp(z_extreme) / np.exp(z_extreme).sum()
)  # Will overflow -> NaN
stable_result = stable_softmax(z_extreme)

Out[29]:

Console

Naive softmax (overflow): [nan nan nan]
Stable softmax (correct): [0.244728 0.665241 0.090031]

Our softmax function already uses the log-sum-exp trick. Production implementations in PyTorch and JAX also use this by default, so overflow is typically not a concern in practice, but understanding why it matters explains why the stable version is the standard. When building your own attention implementations, always use the numerically stable version.