Math Foundations for ML › Probability & Statistics · Note 1

Expectation and Variance

Xinheng Ding

Oct 18, 2026

Most "why is there a constant here?" questions in deep learning come down to the same four facts about means and variances. These include initialization schemes, normalization layers, and the \(1/\sqrt{d_k}\) in attention. This note collects the definitions and those four rules, with proofs. It then uses them, step by step, to explain why scaled dot-product attention divides by \(\sqrt{d_k}\).

Expectation

A random variable \(X\) is a quantity whose value is decided by chance. Its expectation (mean) is the probability-weighted average of its values:

For a fair die, \(\mathbb{E}[X] = \tfrac{1}{6}(1+2+3+4+5+6) = 3.5\). The expectation of a function of \(X\) is computed the same way: \(\mathbb{E}[g(X)] = \sum_x g(x)\,P(X=x)\). We write \(\mu = \mathbb{E}[X]\).

Variance

The variance measures how far \(X\) typically lands from its mean. It is the expected squared deviation:

Expanding the square gives the form used in almost every calculation:

The second step uses linearity, which is Rule 1 below. Rearranged, this gives a fact we will need later:

The standard deviation \(\sigma = \sqrt{\operatorname{Var}(X)}\) has the same units as \(X\). It is the natural answer to "how big is a typical fluctuation?" For a fair die, \(\mathbb{E}[X^2] = 91/6\), so \(\operatorname{Var}(X) = 91/6 - 3.5^2 \approx 2.92\) and \(\sigma \approx 1.71\).

Four rules

Rule 1 · Linearity of expectation

Proof (discrete). Sum over the joint distribution: \(\mathbb{E}[aX+bY] = \sum_{x,y}(ax+by)P(x,y) = a\sum_x x\sum_y P(x,y) + b\sum_y y\sum_x P(x,y) = a\,\mathbb{E}[X] + b\,\mathbb{E}[Y]\), because summing the joint probability over one variable leaves the marginal of the other. Constants pass through since \(\sum P = 1\). ∎

Rule 2 · Product of independent variables

Proof. Independence means \(P(x,y) = P(x)P(y)\), so the double sum factors: \(\mathbb{E}[XY] = \sum_{x,y} xy\,P(x)P(y) = \big(\sum_x x P(x)\big)\big(\sum_y y P(y)\big)\). ∎

Rule 3 · Variance of a sum

Proof. Apply \(\operatorname{Var}(Z) = \mathbb{E}[Z^2] - \mathbb{E}[Z]^2\) to \(Z = X+Y\):

Under independence, Rule 2 makes the covariance zero. For \(n\) terms, expand the square of the sum. Every cross term is a covariance between two different variables, and each one is zero. ∎

Rule 4 · Scaling by a constant

Proof. The mean follows from Rule 1. For the variance, the shift \(b\) cancels against the mean: \(\operatorname{Var}(aX+b) = \mathbb{E}\big[(aX + b - a\mu - b)^2\big] = \mathbb{E}\big[a^2(X-\mu)^2\big] = a^2\operatorname{Var}(X)\). ∎

Variance is quadratic, so constants come out squared. A useful special case is standardization: \(Z = (X-\mu)/\sigma\) always has mean 0 and variance 1. That is exactly what normalization layers do to activations.

Application: why attention divides by \(\sqrt{d_k}\)

Scaled dot-product attention (Vaswani et al., 2017) computes \(\operatorname{softmax}(QK^\top/\sqrt{d_k})\,V\). The Annotated Transformer implements it in one line:

scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

The paper justifies the constant in a short paragraph. Paraphrased, it makes three claims. First, without scaling, dot-product attention does worse than additive attention once \(d_k\) is large, although the two are similar for small \(d_k\). Second, the likely cause is that large dot products push the softmax

"into regions where it has extremely small gradients" Vaswani et al., 2017, §3.2.1

Third, to see why dot products grow, assume the components of \(q\) and \(k\) are independent with mean 0 and variance 1. Then \(q\cdot k\) has mean 0 and variance \(d_k\), so the authors scale it by \(1/\sqrt{d_k}\). The rest of this note derives each step using the four rules.

Why assume mean 0 and variance 1?

The assumption is not a particular distribution. A standard normal, a uniform on \([-\sqrt3,\sqrt3]\) and a fair \(\pm1\) coin all qualify. The derivation uses only means, variances and independence, so the shape does not matter. The scale is chosen because it is the natural reference point, and it is roughly what the network sees at initialization: the inputs to attention are layer-normalized, and the projection weights are initialized (e.g. Xavier) to preserve variance. The question being asked is therefore: if every component has an ordinary size, how big is their dot product?

Step 1 · One term \(q_i k_i\)

By Rule 2, using the independence of \(q_i\) and \(k_i\):

For the variance, use \(\operatorname{Var} = \mathbb{E}[\,\cdot^2] - \mathbb{E}[\,\cdot\,]^2\). Then apply Rule 2 to \(q_i^2\) and \(k_i^2\), which are also independent:

where \(\mathbb{E}[q_i^2] = \operatorname{Var}(q_i) + \mathbb{E}[q_i]^2 = 1\) comes from the identity in the variance section.

Step 2 · Sum of \(d_k\) terms

The dot product is \(q\cdot k = \sum_{i=1}^{d_k} q_i k_i\). Different terms use different components, so the \(d_k\) terms are independent. Linearity gives the mean, and Rule 3 gives the variance:

The typical size of a dot product is therefore \(\sigma = \sqrt{d_k}\). For the paper's \(d_k = 64\), that is 8. Each term is a random-sign number of size about 1. Summing \(d_k\) of them, the signs partly cancel, so the total grows like \(\sqrt{d_k}\), not \(d_k\). This is the same law as a random walk.

Step 3 · Divide by \(\sqrt{d_k}\)

Apply Rule 4 with \(a = 1/\sqrt{d_k}\):

The variance of a score is now 1 regardless of \(d_k\). In simulation, it measures 1.000, 0.996 and 1.008 for \(d_k\) = 4, 64 and 512. The constant is \(\sqrt{d_k}\) and not \(d_k\) because Rule 4 squares it. Dividing by \(d_k\) would give variance \(1/d_k\), which squeezes all scores toward 0 and makes every attention row nearly uniform.

Step 4 · Why large scores starve the gradient

Strictly speaking, the problem is not \(d_k\) itself but the size of the scores, whose standard deviation is \(\sqrt{d_k}\). The paper's claim is that large scores push the softmax into a region where its gradient nearly vanishes. To see what that region is, start with the smallest case.

Two keys. Softmax only depends on differences between scores, because adding a constant to every \(z_j\) cancels in the ratio. With two keys and gap \(\Delta = z_1 - z_2\),

So a two-key softmax is a sigmoid of the gap. A sigmoid is steep in the middle and flat at both ends. The "region with extremely small gradients" is those flat ends. There, \(p_1\) is already close to 0 or 1, and moving the scores barely changes the output.

Going from a gap of 0 to a gap of 8 shrinks the slope about 750 times. For large \(\Delta\), \(p_1(1-p_1) \approx e^{-\Delta}\), so the decay is exponential in the gap.

Many keys. With \(m\) keys, the derivative is a Jacobian matrix. Write \(p_i = e^{z_i}/S\) with \(S = \sum_j e^{z_j}\), and note that \(\partial S/\partial z_j = e^{z_j}\). The quotient rule gives

When one score is far above the rest, \(p\) is nearly one-hot. Every entry of the Jacobian then contains a factor that is close to 0. On the diagonal, the winner has \(1 - p_i \approx 0\) and every loser has \(p_i \approx 0\). Off the diagonal, at least one of \(p_i, p_j\) is close to 0. The whole matrix is close to zero: this is the many-key version of the flat tails.

A temperature view. Dividing every score by a constant \(T\) is the same as running softmax at temperature \(T\). Low temperature (scores effectively multiplied up) makes softmax sharp and close to argmax. High temperature makes it smooth and close to uniform. Without scaling, the effective temperature of attention falls as \(1/\sqrt{d_k}\) when \(d_k\) grows. Dividing by \(\sqrt{d_k}\) fixes the temperature at a value that does not depend on \(d_k\).

What it does to training. The loss reaches the query and key projections only through the softmax:

and \(\partial L/\partial W_Q\), \(\partial L/\partial W_K\) are computed from \(\partial L/\partial z\). If \(J \approx 0\), those gradients are close to 0 too. At initialization, the random projections make each query attend almost entirely to one arbitrary key. If that key is wrong, the loss has almost no way to say "look elsewhere", because the signal is crushed as it passes back through the softmax. This is a local saturation of one function, not the layer-by-layer shrinking usually called vanishing gradients, but the effect is similar: attention learns slowly or not at all.

This is the whole argument. Rules 2 and 3 show that the scores grow like \(\sqrt{d_k}\). Rule 4 shows that dividing by \(\sqrt{d_k}\) undoes this exactly. The softmax Jacobian shows why it matters: without the constant, attention would start out nearly one-hot, with almost no gradient to move it.

Summary

Appendix: the simulation

Pure Python, no dependencies. It produces every number quoted above.

"""Simulate q·k for q, k with i.i.d. N(0, 1) components, before and after dividing by sqrt(d_k)."""
import math
import random
import statistics as st

random.seed(0)
N = 20000


def dot(d):
    return sum(random.gauss(0, 1) * random.gauss(0, 1) for _ in range(d))


print(f"{'d_k':>5} {'mean':>8} {'var':>9} {'std':>7} {'scaled var':>11}")
for d in (4, 64, 512):
    xs = [dot(d) for _ in range(N)]
    scaled = [x / math.sqrt(d) for x in xs]
    print(f"{d:>5} {st.mean(xs):>+8.3f} {st.pvariance(xs):>9.2f} {st.pstdev(xs):>7.2f} {st.pvariance(scaled):>11.3f}")


def softmax(z):
    m = max(z)
    e = [math.exp(v - m) for v in z]
    s = sum(e)
    return [v / s for v in e]


random.seed(1)
d = 64
z = [dot(d) for _ in range(5)]
for name, scores in (("raw", z), ("scaled", [v / math.sqrt(d) for v in z])):
    p = softmax(scores)
    print(f"\n{name:>6} scores : {[round(v, 2) for v in scores]}")
    print(f"{name:>6} softmax: {[round(v, 3) for v in p]}")
    print(f"{name:>6} max p(1-p) = {max(v * (1 - v) for v in p):.4f}")

  d_k     mean       var     std  scaled var
    4   +0.008      4.00    2.00       1.000
   64   -0.006     63.73    7.98       0.996
  512   +0.018    516.07   22.72       1.008

   raw scores : [-1.04, -6.83, -7.02, 3.25, 8.74]
   raw softmax: [0.0, 0.0, 0.0, 0.004, 0.996]
   raw max p(1-p) = 0.0042

scaled scores : [-0.13, -0.85, -0.88, 0.41, 1.09]
scaled softmax: [0.142, 0.069, 0.067, 0.242, 0.481]
scaled max p(1-p) = 0.2496

Related: Attention Masks in the Transformer.