mupblind

muP buys wide-model correctness by making attention start blind, and you can watch the blindness law in a single forward pass.

The maximal update parametrization (muP) lets you tune hyperparameters on a small model and transfer them to a large one. One of its rules scales attention logits by 1/d instead of the standard 1/sqrt(d), so that in the trained regime, where a query and its matching key are aligned and q . k is order d, the logits stay order 1 across width. This is about what that rule does at initialization, before any alignment exists.

The init law

At init the query and key maps are independent with muP hidden variance 1/m, so q = W_q x and k = W_k x have order-1 coordinates and are independent. The pre-softmax logit l = c * (q . k) over d coordinates has mean 0 and

Var(l) = c^2 * d * Var(q_i) * Var(k_i) = Theta(c^2 * d) .

• standard c = 1/sqrt(d): Var(l) = Theta(1), head-dim invariant.
• muP c = 1/d: Var(l) = Theta(1/d) -> 0.

For a softmax over n keys with small-variance logits, E[H] = log n - Theta(v), so the entropy gap G = log n - E[H] inherits the logit variance: muP's gap is Theta(1/d) and collapses toward uniform as head dimension grows, while the standard gap is Theta(1) and width-invariant. Measured across d from 16 to 512, the log-log slope of G is -1.001 under muP and -0.002 under standard, and muP's gap halves on every head-dim doubling.

So the coordinate-check invariance muP guarantees for MLP activation scale is violated in the attention subblock at init: muP starts attention maximally diffuse, and more so at larger width, because at init q . k is only order sqrt(d) rather than the order d the 1/d scale was calibrated for.

It is a transient, restored by training

This is an initialization effect, not a permanent defect. Training a single attention head to attend to a target key per query, the muP entropy gap starts near zero (0.004, maximally diffuse) and rises to 1.9 over a few hundred steps as query-key alignment develops and q . k grows from order sqrt(d) toward order d. Standard parametrization starts already peaked (0.45) and rises from there. muP's attention sharpens into place; the init blindness is the price of the width-invariance guarantee, paid back through training.

Layout

• entropy_law.py: the initialization entropy gap versus head dimension, the Theta(1/d) collapse under muP (slope -1) against the flat standard baseline.
• restoration.py: the entropy gap through training, showing muP sharpens from near-uniform as alignment develops.
• test_mupblind.py: the slopes, the collapse, and the restoration as tests.

Reproduce

python entropy_law.py
python restoration.py
python test_mupblind.py

Both experiments run in seconds; the init law needs only a forward pass.

Takeaway

Under muP a wider model starts with more diffuse attention, by a clean 1/d law, and reaches useful attention only after training aligns queries and keys. The 1/d attention scale is right for the trained regime and deliberately wrong for init, which is visible as an init entropy gap that collapses with width and then recovers.