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.

What it is and is not

muP's demonstrated wins are specifically about width transfer: hyperparameters tuned on a narrower model in the same family transfer near-exactly to a much wider one, cutting the cost of tuning at frontier scale substantially. It is not a general solution to hyperparameter transfer across every axis of scale. Transfer across depth is a separate and considerably less settled question, and transfer across a genuinely different data mix or architecture family is outside what the theory guarantees. Cross-linking to the practical schedule and warmup choices it still leaves open: learning-rate-schedules and warmup-and-why-it-helps still need their own tuning even once the peak learning rate itself transfers cleanly.

When it falls down

• The guarantees are asymptotic. muP's theory is proved in the infinite-width limit; at the finite widths anyone actually trains, transfer is good in practice but not exact, and some quantities transfer more cleanly than others.
• Not every hyperparameter transfers equally well. Learning rate transfers best in published results; weight decay, warmup length, and batch size are reported as transferring less cleanly and often still need some adjustment at the target scale.
• It is a real engineering commitment, not a flag. Implementing muP correctly means touching the initialisation and learning-rate handling of every layer type in the architecture consistently; it is not a drop-in option applied after the fact.
• It says nothing about depth or data scale. A muP-tuned learning rate still needs revalidation if the target model is also much deeper, or trained on a substantially different data distribution, than the proxy it was tuned on.

Further reading

• Yang et al., 2022, Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, arXiv:2203.03466 - the muP parametrisation and its width-transfer guarantees.
• Yang and Hu, 2020, Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks, arXiv:2011.14522 - the underlying feature-learning theory muP builds on.