Scaling Dot Products for Stable Softmax

Introduction

In our last lesson, we constructed the full scaled dot-product attention mechanism, following the data from input Query, Key, and Value matrices to the final context-aware output vectors. We focused on the what and how of the computation, guided by the formula:

Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V

We noted the scaling factor, $\frac{1}{d _{k}}$ , as a crucial step but deferred its justification. Today, we will address that "why." This lesson delves into the statistical reasoning behind this specific design choice, explaining how it is essential for the stable and effective training of Transformer models. You will learn why unscaled dot products can be problematic and how this simple division prevents common deep learning pitfalls like softmax saturation and vanishing gradients.

The Problem with Unscaled Dot Products: Variance and Dimension

Let's start by considering what would happen if we omitted the scaling factor. The attention scores would simply be the raw dot products, $Q K^{T}$ . The core issue stems from a fundamental statistical property: as the dimension of the vectors being multiplied increases, so does the variance of their dot product.

To understand this, let's make a reasonable assumption. Due to mechanisms like Layer Normalization and specific weight initializations (like Xavier initialization), which are standard in the Transformer architecture, we can assume that the elements of our Query ( $Q$ ) and Key ( $K$ ) vectors are independent random variables with a mean of 0 and a variance of 1.

Given two such vectors, $q$ and $k$ , of dimension $d_{k}$ , their dot product is $s = \sum_{i = 1}^{d_{k}} q_{i} k_{i}$ . A key result from statistics, known as the Bienaymé formula, states that the variance of a sum of independent random variables is the sum of their variances. In this case, it can be shown that the variance of the dot product $s$ is approximately equal to the dimension $d_{k}$ .

Var (q \cdot k) \approx d_{k}

This means that as we increase the representational power of our model by using larger embedding dimensions (a larger $d_{k}$ ), the attention scores will become more spread out. A larger variance implies that the scores will have a wider range of values, with some becoming very large in magnitude.

The Consequences: Softmax Saturation and Vanishing Gradients

This growing variance poses a significant problem for the next step in our computation: the softmax function. The softmax function, $σ (z)_{i} = \frac{e ^{z_{i}}}{\sum _{j} e ^{z_{j}}}$ , is highly sensitive to the scale of its inputs.

When the input values have high variance, a few scores will likely be much larger than the rest. Due to the exponential nature of softmax, these large scores get amplified, resulting in a probability distribution that is extremely "peaky" or close to one-hot. For example, softmax([1, 5, 10]) results in [0.000, 0.007, 0.993], where almost all the probability mass is assigned to the largest score.

This phenomenon is called softmax saturation. It leads to two major issues for training:

Poor Context Aggregation: The attention mechanism is forced to attend almost exclusively to a single token, losing the ability to create rich contextual representations by blending information from multiple relevant tokens.
Vanishing Gradients: In regions where the softmax output is saturated (very close to 0 or 1), the gradient of the loss function with respect to the pre-softmax scores becomes vanishingly small. During backpropagation, this means that almost no learning signal flows back to the weights that produced the non-dominant scores. The model effectively stops learning from those connections, severely hindering the training process.

The following video provides an excellent intuitive and mathematical walkthrough of this entire problem.

Why Scaling by the Square Root of Dimensions Matters in Attention | Transformers in Deep Learning

The video "Why Scaling by the Square Root of Dimensions Matters" from the channel Learn With Jay explains this concept from first principles.

First, watch the segment that establishes the core statistical problem, showing how the variance of the dot product result is proportional to the dimension d_k. Focus on the explanation from this section. Next, watch the part that explains the consequences of this high variance when fed into a softmax function, from this discussion. Pay close attention to the explanation of how a "peaky" softmax output prevents most weights from being updated during training.

The Solution: Normalizing the Variance

Now that we've diagnosed the problem—the variance of the dot products scales with $d_{k}$ , causing softmax saturation—the solution becomes clear. We need to counteract this effect by normalizing the scores before they enter the softmax function. Our goal is to make the variance independent of $d_{k}$ , ideally keeping it constant (e.g., at 1).

This is where the scaling factor comes in. We can use another property of variance: for a random variable $X$ and a constant $c$ , $Var (c X) = c^{2} Var (X)$ .

Let our unscaled attention score be $s = q \cdot k$ , with $Var (s) \approx d_{k}$ . We want to find a scaling factor $c$ to apply to $s$ , creating a new score $s^{'} = c \cdot s$ , such that $Var (s^{'}) = 1$ .

Using the variance property:

Var (s^{'}) = Var (c \cdot s) = c^{2} Var (s) = c^{2} d_{k}

Setting this to our desired variance of 1:

c^{2} d_{k} = 1 ⟹ c^{2} = \frac{1}{d _{k}} ⟹ c = \frac{1}{d _{k}}

This elegant result shows that by dividing the dot product scores by the square root of the key dimension, we effectively normalize their variance to be approximately 1. This ensures that the inputs to the softmax function remain in a "healthy," non-saturated range, regardless of the chosen dimension $d_{k}$ .

A concrete example of scaling from the "Illustrated Transformer". The raw dot product score (q1・k1 = 112) is scaled by dividing by $\sqrt{d_k}$ (here, $\sqrt{64}=8$), yielding a value of 14, which is then fed to the softmax function.

This histogram shows the distribution of attention scores *after* scaling. They are centered around zero with a standard deviation of 1.0, an ideal input distribution for the softmax function, preventing saturation.

For a rigorous, text-based explanation of this derivation, the following resource is excellent.

Why use a "square root" in the scaled dot product

This Stack Exchange answer provides a concise and mathematically precise explanation for the scaling factor, directly referencing the original "Attention Is All You Need" paper.

Please read the first and most upvoted answer by user pi-tau, found in this section. It formalizes the argument we've just discussed. Also, review the other two short answers, which quote the original paper and offer a slightly different framing of the same concept.

To complete our understanding, let's return to the video from Learn With Jay, which beautifully demonstrates how this scaling factor provides the solution.

Why Scaling by the Square Root of Dimensions Matters in Attention | Transformers in Deep Learning

This final segment of the video demonstrates the mathematical fix.

Watch the concluding part of the explanation from this chapter. It uses the variance property \text{Var}(aX) = a^2 \text{Var}(X) to derive the \frac{1}{\sqrt{d_k}} factor, bringing the entire argument together.

Conclusion

In this lesson, we have uncovered the statistical necessity of the scaling factor in scaled dot-product attention. It is not an arbitrary hyperparameter but a principled solution to a problem that arises from the interaction between high-dimensional vector operations and the softmax function.

Here are the key takeaways:

The Problem: The variance of the dot product between query and key vectors is proportional to their dimension, $d_{k}$ . As $d_{k}$ grows, the attention scores spread out.
The Consequence: High-variance scores push the softmax function into saturation, where it produces nearly one-hot distributions. This leads to vanishing gradients, stalling the training process.
The Solution: By dividing the dot product scores by $d_{k}$ , we normalize their variance to be approximately 1. This keeps the inputs to the softmax in a stable range, ensuring that gradients can flow effectively and the model can learn to weigh information from multiple tokens.

You now understand a critical detail that makes the attention mechanism, and by extension the entire Transformer architecture, trainable at scale. In our next lesson, we will explore another crucial modification to the attention mechanism: causal masking, which is essential for building autoregressive language models like GPT that generate text one token at a time.

Recommended resources

Introduction

The Problem with Unscaled Dot Products: Variance and Dimension

The Consequences: Softmax Saturation and Vanishing Gradients

The Solution: Normalizing the Variance

Conclusion