Softmax Scaling & Vanishing Gradients

Introduction

In our last session, we constructed the full formula for scaled dot-product attention and walked through a concrete computation. We identified the formula as:

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

We briefly mentioned that the scaling factor, $\frac{1}{\sqrt{d_k}}$, is important for stabilizing the training process. Today, we will explore precisely why this is the case.

Our learning goal for this lesson is to explain the role of the scaling factor in preventing vanishing gradients through softmax. We will answer three key questions:

1. Why do the dot products in $Q{K}^T$ tend to get larger as the dimension of the key vectors, $d_k$, increases?
2. Why are large input values a problem for the softmax function?
3. How does dividing by exactly $\sqrt{d_k}$ solve this problem?

1. The Problem: Variance Grows with Dimensionality

Let's first establish the core problem. The dot product between two vectors is calculated by multiplying their corresponding elements and summing the results.

For two vectors $q$ and $k$ of dimension $d_k$, their dot product is:
$q\cdot k={\sum }_{i=1}^{d_k}{q}_i{k}_i$

Intuitively, as the dimension $d_k$ grows, you are summing up more and more products. This suggests that the resulting dot product values will have a larger spread, or variance.

This isn't just an intuition; it's a statistical property. If the components of our query and key vectors are drawn from a random distribution with a mean of 0 and a variance of 1, then the dot product $q\cdot k$ will have a mean of 0 and a variance of $d_k$.

This means: The variance of the dot products in $Q{K}^T$ scales linearly with the dimension $d_k$.

To build a strong intuition for this, the following video provides an excellent visualization. It uses histograms to show how the distribution of dot product values widens dramatically as the vector dimension increases.

2. The Consequence: Softmax and Vanishing Gradients

Now we know that for large models with high embedding dimensions (e.g., $d_k=512$ or higher), the attention scores in the $Q{K}^T$ matrix will have a large variance. Why is this a problem?

The issue arises when these high-variance scores are fed into the softmax function. The softmax function involves an exponential term, $e^x$. This has the effect of exaggerating differences: it pushes large input values towards 1 and small input values towards 0.

Consider two sets of scores:

• Low variance scores: [1, 2, 3] -> softmax -> [0.09, 0.24, 0.67] (a gentle distribution)
• High variance scores: [1, 5, 10] -> softmax -> [0.0001, 0.0067, 0.9932] (a highly "peaked" or "spiky" distribution)

When the softmax output is extremely close to 0 or 1, it has entered a saturation region. In these regions, the gradient of the softmax function with respect to its input is nearly zero. This is the infamous vanishing gradient problem.

If the gradient is zero, no learning signal can flow backward during backpropagation. This means the model stops updating the weights that produced the smaller attention scores, effectively ignoring parts of the input sequence and severely hindering its ability to learn.

The next segment of the video you just watched uses a fantastic classroom analogy to explain this exact problem.

3. The Solution: Scaling to Normalize Variance

So, the problem is clear:
High dimension (d_k) -> High variance in QK^T -> Saturated softmax -> Vanishing gradients

The solution is to "tame" the variance of the $Q{K}^T$ matrix before it enters the softmax function. We want to scale it down so that its variance remains constant (ideally, at 1), no matter how large $d_k$ gets.

But what should the scaling factor be? This is where a key property of variance comes into play:

For a random variable $X$ and a constant $c$, the variance of their product is: $\text{Var}(cX)=c^2\text{Var}(X)$

We established that $\text{Var}(Q{K}^T)\approx d_k$. Our goal is to find a scaling constant $c$ such that the variance of the scaled scores is 1.

$\text{Var}(c\cdot Q{K}^T)=1$
$c^2\cdot \text{Var}(Q{K}^T)=1$
$c^2\cdot d_k=1$
$c^2=\frac{1}{d_k}$
$c=\frac{1}{\sqrt{d_k}}$

This elegant derivation shows that dividing by the square root of the dimension $d_k$ is precisely the operation needed to normalize the variance of the attention scores to 1, keeping the inputs to softmax in a healthy, non-saturating range.

The following resources provide a final, clear walkthrough of this mathematical justification. You can choose whichever you prefer, as they both cover the core derivation effectively.

Or, for a text-based version with the same derivation:

By performing this scaling, we ensure that the attention mechanism can be trained stably, even in very deep networks with high-dimensional embeddings.

Conclusion

In this lesson, we took a deep dive into a single, critical component of the attention formula. You have moved from knowing what the scaling factor is to understanding why it is essential.

Key Takeaways:

• The variance of the dot product attention scores ($Q{K}^T$) increases linearly with the dimension of the key vectors, $d_k$.
• Feeding high-variance scores into the softmax function pushes the output probabilities to extremes (0 or 1), causing the function to saturate.
• A saturated softmax function has near-zero gradients, which prevents the model from learning effectively (the vanishing gradient problem).
• By dividing the attention scores by $\sqrt{d_k}$ before the softmax, we normalize the variance to 1. This keeps the inputs in a stable range and ensures that gradients can flow properly during training.

Preview of the Next Lesson:
We have now mastered the mechanics of a single "attention head." However, a single head might only learn to focus on one type of relationship (e.g., subject-verb agreement). To make the model more powerful, Transformers run multiple attention heads in parallel, allowing the model to simultaneously attend to different aspects of the input sequence. In our next lesson, we will explore this concept of multi-head attention.