Why do we scale the dot product attention by √dₖ in the Transformer architecture?

In this question, we will learn why the Transformer architecture divides the raw query key dot product by $\sqrt{d_k}$ before applying softmax. We will show the statistical reason, a concrete numeric example, and how this ties into multi head design and initialization.

We will cover the following:

• The intuition
• How it actually works (variance argument plus a worked example)
• Relation to multi head attention and initialization
• Tradeoffs and failure modes
• Questions the interviewer might ask

Direct answer: We scale by $\sqrt{d_k}$ to keep the dot-product logits at a consistent variance so the softmax is not driven into saturation and gradients remain well behaved. This normalization makes attention heads stable across different $d_k$ and simplifies optimization.

The intuition (an analogy that makes it click)

Imagine you are comparing similarity scores from longer and shorter lists. If longer lists produce sums of many small contributions, their totals will typically be larger just because there are more terms. The dot product of two $d_k$ dimensional vectors is the sum of $d_k$ componentwise products. As $d_k$ grows the raw sum tends to grow in magnitude, so we rescale by $\sqrt{d_k}$ to restore a fair comparison. That makes the softmax see scores on a similar scale whether you have small or large head dimensions.

How it actually works (the real mechanics, with one concrete worked example)

Assume each component of query and key vectors is independent with mean zero and variance $\sigma^2$. The dot product $q\cdot k$ is a sum of $d_k$ such products. For zero mean and independent components the expected mean is zero and the variance grows roughly linearly with $d_k$.

A compact expression when components have variance $\sigma^2$ is:

$\mathrm{Var}(q\cdot k) = d_k \sigma^2.$

If we divide the dot product by $\sqrt{d_k}$ the variance becomes $\sigma^2$, independent of $d_k$. That keeps the distribution of logits stable as we change head size.

Worked numeric example. Suppose a head produces raw dot products (logits) for three keys as [16, 8, 4]. Those are already somewhat large relative to 1, so softmax will be very peaky. Now scale by $\sqrt{d_k}$, assume $\sqrt{d_k}=8$, to get [2, 1, 0.5]. Compare their softmax probabilities.

The unscaled logits produce an almost one-hot distribution that kills gradient flow to the lower-scoring keys. After scaling we get a smoother distribution that allows learning to adjust weights for multiple keys.

Relation to multi head attention and initialization

Multi head attention splits the model dimension into multiple heads, each with dimension $d_k$. If we did not scale, heads with larger $d_k$ would produce logits with larger variance. Scaling by $\sqrt{d_k}$ makes head outputs comparable and avoids having some heads be always dominant because of dimension alone.

Scaling also interacts with parameter initialization. Standard initializations aim to keep activations and gradients roughly constant with depth. The $\sqrt{d_k}$ factor is consistent with that goal: it normalizes the attention logits so the softmax output variance does not blow up as $d_k$ changes.

Tradeoffs and failure modes

Scaling by $\sqrt{d_k}$ is simple and effective, but it implicitly sets an attention temperature. Smaller effective temperature leads to sharper attention; larger temperature gives smoother attention. If you change $d_k$ or add explicit temperature parameters you must be careful to maintain the intended sharpness of attention.

Questions the interviewer might ask

Some follow-up questions you might get:

Why assume zero mean and independence for components? Those assumptions are a simplifying model that captures the main effect: sums of many small independent terms have variance that grows with count. Real activations are not independent, but the scaling still empirically helps.

Could we learn the temperature instead of fixing $\sqrt{d_k}$? Yes, some variants learn a scalar temperature or add a learned scale per head. That adds flexibility but also more parameters and potential instability without careful regularization.

Why use $\sqrt{d_k}$ not $d_k$ or another factor? The dot product variance grows linearly in $d_k$, so the standard deviation grows with $\sqrt{d_k}$. Dividing by $d_k$ would shrink too fast; dividing by $\sqrt{d_k}$ normalizes variance to a constant.

Does layer normalization change the need for scaling? Layer normalization helps stabilize activations but does not remove the variance growth of the dot product between independent vectors. Scaling remains useful even with normalization layers.

How does scaling affect attention heads when d_k is small? When $d_k$ is small scaling is less dramatic, but the factor still keeps behavior consistent across designs. With very small $d_k$ the division can increase noise sensitivity, so design choices should balance head count and head size.

Some things to note:

• The $\sqrt{d_k}$ factor is a variance normalization choice, not a magic constant. It follows from basic statistics of sums.
• Empirically this scaling improves training stability across many Transformer models.

What the interviewer is really testing

They want to know you understand why raw dot products change scale with vector dimension and how that affects the softmax and gradients. They also want to see you can connect a statistical argument to practical architectural choices like multi head design and initialization. A concise variance argument plus a numerical example shows you grasp both theory and practice.