Motivation

In transformer attention, the logits are $q^\top k$, where $q, k \in \mathbb{R}^d$ are query and key vectors. So if μP says to scale attention by $1/d$ rather than the usual $1/\sqrt{d}$, the natural question is: why?

If $q$ and $k$ were independent, the usual $1/\sqrt{d}$ scaling would be exactly the right guess. Each term $q_i k_i$ is mean-zero with bounded variance, so the sum $q^\top k = \sum_i q_i k_i$ lives in the familiar Central Limit Theorem world and fluctuates at scale $\sqrt{d}$.

But during training, $q$ and $k$ pick up correlations. Then the summands $q_i k_i$ can reinforce one another, so the dot product starts looking less like a CLT object and more like a Law of Large Numbers object. In the large-width regime that μP cares about, this matters because we want the size of $q^\top k$ to stay $O(1)$ as $d$ grows.

So the real issue is whether $q^\top k$ behaves like a $\sqrt{d}$ object or a $d$ object. This note gives a simple Gaussian calculation that makes that distinction explicit. The object we will compute is the $L^2$ norm

$$\|q^\top k\|_2 := \left(\mathbb{E}[(q^\top k)^2]\right)^{1/2}$$

for jointly Gaussian $q$ and $k$. The answer splits neatly into a baseline term, which is there even when $q \perp k$, and two alignment terms controlled by the cross-covariance $C = \mathbb{E}[qk^\top]$. That decomposition makes the $\sqrt{d} \to d$ transition completely explicit.

---

Setup

Of course, real query and key vectors in a trained network are not literally Gaussian. But in the spirit of μP, it is useful to start with a simplified probabilistic model and see what it already explains.

Before specializing to $S = q^\top k$, let me fix the standard probability notation: if $X$ is a random variable on a probability space and $\mathbb{E}[|X|^p] < \infty$, then its $L^p$ norm is $$ \|X\|_p := \left(\mathbb{E}[|X|^p]\right)^{1/p}. $$ See Vershynin, High-Dimensional Probability, Chapter 1.

---

Main Result

---

Interpretation

The formula is easy to read once you look at the three terms separately:

• Independent ($C = 0$): $\|q^\top k\|_2 = \sqrt{d}$ — CLT scaling.
• Isotropic correlation ($C = \rho I$): $\|q^\top k\|_2 = \sqrt{d + \rho^2 d + \rho^2 d^2}$ — any nonzero $\rho$ makes the $d^2$ term dominate.
• Perfect alignment ($k = q$): $\|q^\top k\|_2 \approx d$ — LLN scaling.

So the whole story is really carried by $C$. If $C$ is negligible, you stay in the $\sqrt{d}$ world. If $C$ has coherent structure, the $d$-scale behavior starts to take over.

---

Proof of Theorem 1

This clean version of the proof using Stein's formula was found with the help of ChatGPT. My own proof computed $\mathbb{E}[(q^\top k)^2]$ by expanding everything into fourth-order Gaussian moments and I realized that I needed to Isserlis's theorem to continue the proof. That works, but in coordinates it quickly turns into quite a mess.

After a few rounds of discussion with ChatGPT, it suggested using Stein's lemma , which was a suprisingly elegant alternative. We package $q^\top k$ as a quadratic form $w^\top A w$ in the joint vector $w = (q, k)$, and then use Stein's lemma to get the needed quadratic-form identity almost for free.