muP and Hyperparameter Transfer

Tuning a learning rate at the scale you actually intend to train is often more expensive than the training run itself. The standard workaround, tune on a small proxy model and use the same value at full scale, has an inconvenient failure mode: under conventional network parametrisation, the learning rate that works best genuinely changes as you widen a network, so the proxy's answer is frequently wrong for the target. Maximal Update Parametrisation, muP, is a specific way of scaling initialisation and learning rate with width that is designed to make that transfer actually hold.

Why standard parametrisation breaks with width

Under the conventional way of initialising and updating weights, several quantities that determine how a training step behaves, the scale of activations, the effective size of a parameter update relative to the activations it feeds, drift systematically as the number of hidden units per layer grows. A learning rate tuned for a narrow network can produce updates that are comparatively too large or too small once the network is widened, and the direction of the drift is not something you can casually eyeball. In practice this shows up as needing to re-tune the learning rate, sometimes substantially, every time you meaningfully change model width, which defeats the point of using a cheap proxy in the first place.

What muP changes

Yang et al.'s Tensor Programs V paper derives a specific set of per-layer scaling rules for initialisation variance and learning rate, as a function of width, chosen so that in the infinite-width limit, the scale of each layer's update relative to its activations stays stable regardless of how wide the network is. The practical payoff: under muP, the optimal learning rate (and several other hyperparameters, including relative learning rates for the embedding layer versus the rest of the network) converges to a value that does not depend on width. Sweep those hyperparameters cheaply on a small proxy trained under muP, and the same values transfer to a much larger model in the same architecture family, built on the underlying feature-learning theory laid out in Tensor Programs IV.

The coordinate check

Implementing muP correctly across every layer type in a real architecture, attention projections, embeddings, output head, is enough surface area that subtle mistakes are easy to make even with the maths right on paper. The standard diagnostic, released alongside the muP work, is the coordinate check: train several models of different widths under a candidate parametrisation for a few steps, and plot activation or logit scale against width. A correct implementation produces flat curves as width grows; a curve that visibly drifts with width is evidence something in the implementation, not the theory, is wrong. It functions as a cheap unit test for a parametrisation before committing a full tuning budget to it.