Defend an ML derivation when its assumptions change

Derive an attention-score variance from stated random-variable assumptions and test whether the reasoning survives a change in dependence.

Problem

Context. The fixture below defines four-dimensional query and key vectors using independent random signs, then changes their within-vector dependence.

Task. Find the unscaled and scaled dot-product variances in both cases, showing the assumptions needed at each step. Explain what the result does and does not establish about attention scaling in a learned model.

Constraints. Use the specified zero-mean sign distributions rather than assuming arbitrary learned coordinates are independent. Check covariance terms when independence changes; a familiar formula is not a substitute for the derivation.

Let $d = 4$ and $z = q^{T} k$ . All eight coordinates of $q$ and $k$ independently take $+ 1$ or $- 1$ , each with probability $1/2$ . Find the variance of $z$ and of $z / d$ . Then decide whether the same conclusion holds when every query coordinate shares one sign and every key coordinate shares another.

Picture

Each product $X_{i} = q_{i} k_{i}$ is another independent sign: mean zero, variance one. Four independent products add variance; halving their sum multiplies variance by one quarter.

An exact four-coordinate distribution

Square-root scaling stabilizes variance only under the stated dependence assumptions

Derivation

Read it this way: follow the random variable and its variance together. Scaling the variable by one half scales variance by one quarter, not one half.

Check: Why does dividing by 2 yield variance 1 for independent terms but variance 4 for shared terms?

Read it this way: With independent, unit-variance coordinate products, a 64-term dot product has variance 64. Dividing by 8 returns variance 1; dependence can change this result.
Why does dividing the dot product by eight divide its variance by sixty-four?

References: Vaswani et al. and PyTorch scaled-dot-product attention.

Worked example: attention scaling

The decisive identity retains covariance until the assumptions remove it:

Var (z) = i \sum Var (X_{i}) + 2 i < j \sum Cov (X_{i}, X_{j}) .

Here the covariance terms vanish, so $Var (z) = 4$ and $Var (z /2) = (1/2)^{2} \times 4 = 1$ . The $2^{8} = 256$ equally likely sign assignments give an exact check:

$z$	$- 4$	$- 2$	$0$	$2$	$4$
Assignment count	$16$	$64$	$96$	$64$	$16$

Check also $d = 1$ : no scale change is needed. Dividing variance by $d$ would incorrectly give two at $d = 4$ ; variance squares the multiplier.

Now set every $q_{i} = Q$ and every $k_{i} = K$ , with independent signs $Q, K$ . All products equal $Q K$ , so $z = 4 Q K$ has variance 16. Six pair covariances contribute $2 \times 6 = 12$ beyond the four individual variances. Dividing by two leaves variance four.

Change the assumptions or object

1. Change coordinate variances

New task: independent zero-mean query coordinates have variance four; independent key coordinates have variance nine.

Picture and change: retain independence of vectors and products, but each product now has variance $4 \times 9 = 36$ . For $d = 4$ , the sum has variance 144.

Check: dividing by two leaves variance 36. Dividing by $σ_{q} σ_{k} d = 12$ restores one in this stylized distribution.

Use when: testing the unit-variance premise, not claiming exact whitening of trained heads.

2. Average correlated worker errors

New task: average four unit-variance errors with pairwise correlation $0.5$ .

Picture and change: apply the same covariance identity to their average.

Var (\overset{e}{ˉ}) = \frac{4 + 2 \times 6 \times 0.5}{16} = 0.625.

Check: independence gives $0.25$ ; perfect correlation gives one.

Use when: shared worker errors invalidate independent-noise averaging.

Transfer: recover a missing identity

You forget the importance-sampling identity, and the proposal misses some target support. What can you derive safely?

Check your answer

Start from the target integral and insert $q / q$ only where $q > 0$ . Identify any missing contribution; support mismatch can block equality. Large weights or an effective-sample-size number cannot recover omitted support.

Use the same four-part structure

1. Setup

Define variables, dimensions, probability direction and assumptions. For gradients, state the differentiated variable and row/column convention.

2. Key identity

Expose the decisive step: chain rule, log-ratio expansion, log-sum-exp derivative, Jensen’s inequality or a variance identity. Justify it rather than hiding it behind a name.

3. Sanity checks

Use at least two: dimensions, sign, equality case, scalar case, limit, symmetry or numerical scale.

4. Interpretation

Explain the model or system consequence, then change an assumption and identify what survives.

Coordinate independence is sufficient, not necessary, for the attention result. For independent zero-mean vectors with covariance matrices $Σ_{q}, Σ_{k}$ :

Var (q^{T} k) = tr (Σ_{q} Σ_{k}) .

This conditional result does not establish universal softmax saturation or training stability. Learned projections and normalization alter empirical head distributions.

A remembered formula is not a proof

“It prevents large values” does not justify $1/ d$ or identify when the variance calculation changes.

Eight proof checkpoints

These are separate derivations connected by a presentation protocol, not one common mathematical proof.

Prompt	Decisive step and check
Softmax cross-entropy	Log-sum-exp differentiation gives $p - y$ for normalized targets; gradients sum to zero; common logit shifts preserve probabilities.
KL direction	$D (P ∥ Q) = E_{P} lo g (P / Q)$ ; check support, equality and asymmetry.
ELBO	$lo g p (x) = ELBO (q) + D (q (z) ∥ p (z ∣ x))$ , when well-defined; equality uses the exact posterior.
L1 versus L2	$
Geometric attempts	Independent identical success trials counted from one give $E [T] = 1/ p$ ; $p = 1$ gives one attempt.
Averaging	Independent errors give $σ^{2} / n$ ; equal pair correlation gives $σ^{2} [1 + (n - 1) ρ] / n$ .
Importance sampling	Insert $q / q$ in the target integral only where the proposal supports the contributing integrand.
Attention	Keep covariance terms until assumptions justify removing them; square any scale factor in variance.

For Bernoulli parameters $0.8$ and $0.5$ , the KL is about $0.192745$ nats; reversing direction gives $0.223144$ . The unequal values check asymmetry, not an error in the expectation direction. At geometric success probability zero, success never occurs and the expectation is not finite. Importance weights can have infinite variance despite a valid support identity. Effective sample size is a diagnostic, not a finite-variance or support certificate.

Which conclusion survives?

Independence, convexity, asymptotics, unbiasedness and support overlap are different premises. State which identity uses each one. The ELBO gap concerns posterior approximation, not an unspecified reconstruction-plus-regularization slogan.

Evidence in a derivation

Defined notation, justified identities, two checks and explicit error recovery make the reasoning inspectable. No format guarantees a hiring outcome.

Errors to expose

Undefined notation, missing support and unjustified algebra remain errors even when the final formula looks familiar. Stop at the missing step rather than inventing one.

Common follow-up

“You cannot remember the identity needed for the next step. What do you do?”

State what is known, derive the missing piece from a definition or simpler case, and identify the exact unresolved step.

Use the timed math oral with an observer and a changed-assumption follow-up.

Defend an ML derivation when its assumptions change

Practice before you read