Exercise 3: Why Dividing by sqrt(d_k) Is Necessary — Possible Solution ==================================================================== WHAT THE RAW DOT-PRODUCT SCORES LOOK LIKE BEFORE SCALING ------------------------------ A dot product between two vectors is a sum of the products of their individual components. As the dimensionality of Q and K (d_k) grows, the number of terms being summed grows too — even if each individual component is modest in size, summing more and more of these products tends to produce dot products with a larger overall magnitude, simply because there are more terms contributing to the sum. WHY LARGE SCORE MAGNITUDES ARE A PROBLEM FOR SOFTMAX SPECIFICALLY ------------------------------ Softmax exponentiates its inputs before normalizing them. When the input scores are large in magnitude, the largest score's own exponential value can end up enormously bigger than the exponentials of the other scores, pushing softmax's own output toward an extremely peaked distribution — one weight very close to 1, and every other weight very close to 0. This isn't a graceful "confident" distribution; it's a saturation effect caused purely by the score magnitudes, not by genuine confidence in the underlying comparison. WHY THIS CAUSES VANISHING GRADIENTS DURING TRAINING ------------------------------ Once softmax's own output has saturated this way (one weight near 1, the rest near 0), the gradient of the softmax function with respect to its inputs becomes extremely small in the regions responsible for the near-0 outputs. During backpropagation, this means very little useful gradient signal flows back through the attention mechanism for those tokens, which slows or effectively halts learning for exactly the parts of the computation that produced the very large scores in the first place. WHY DIVIDING BY sqrt(d_k) FIXES THIS ------------------------------ Per this chapter, "dividing by sqrt(d_k) ... keeps the scores from growing too large as dimensionality increases." Since dot-product magnitude tends to scale with dimensionality in a roughly predictable way, dividing by the square root of that same dimensionality counteracts the growth, keeping the resulting scores in a range where softmax produces a well-behaved, non-saturated distribution regardless of how large d_k happens to be. WHY THIS MATTERS SPECIFICALLY FOR LARGER EMBEDDING DIMENSIONS ------------------------------ For small d_k, unscaled dot products might stay small enough that softmax behaves reasonably anyway — the problem is specifically that as models use larger embedding dimensions (as most real LLMs do, to represent richer information), unscaled dot-product magnitudes grow correspondingly larger, making the softmax-saturation problem worse exactly as the model scales up. The sqrt(d_k) scaling is what keeps attention numerically well-behaved across that entire range of model sizes, rather than only working for small, toy-sized examples. WHY THIS WORKS AS AN ANSWER ------------------------------ It explains why raw dot-product magnitude grows with dimensionality, connects that growth to softmax's own saturation behavior and the resulting vanishing-gradient problem during training, and explains specifically why the sqrt(d_k) scaling term counteracts this growth in a way that keeps attention stable across both small and large embedding dimensions.