Vol. XI · Optimisation at Scale Plate 8 of 10muP, maximal update parametrisation, hyperparameter transfer, width, learning rate, Tensor ProgramsLC

Transferring Hyperparameters.

Tune on a model a hundred times smaller — if the parametrisation makes its best learning rate the same as the large one’s.

A sweep of learning rates on a 400-billion-parameter model would cost several training runs of that model. Labs tune on small proxies instead, and the question is whether the best settings carry over. In the standard parametrisation (SP) they do not: as a model widens, a step of the same learning rate changes each layer’s output more, because more weights contribute, and the best rate drifts downwards — roughly as 1/width for Adam-trained layers.

The maximal-update parametrisation (μP) of Yang and colleagues (2022) rescales the initialisation, the output multiplier and the per-layer learning rates with width so that one step changes each layer’s activations by an amount that does not depend on width. The best learning rate then stays put, and settings tuned on a small model transfer; they reported tuning a 40-million-parameter proxy and transferring to a 6.7-billion-parameter GPT-3 configuration. The workshop measures how much one Adam step changes the output of a two-layer network as it widens, in each parametrisation.

Below: left, the change in a network’s output from one Adam-like step at a fixed learning rate, against width, in standard and μP-style parametrisation. Right, the learning rate that keeps that change equal to the narrowest network’s.

What you will learn

Explain why the best learning rate drifts with width.
State what μP changes.
Measure update size against width in both parametrisations.
Describe hyperparameter transfer from a proxy.

Workshop

A two-layer ReLU network of width n on one 16-dimensional input; one sign step (Adam’s first step) of learning rate 0.01 on every weight. Standard parametrisation: readout weights ~ N(0, 1/n), the same learning rate everywhere. μP-style: readout weights ~ N(0, 1) with an output multiplier 1/n and readout learning rate scaled by 1/n. Left: |Δ output| against width, log–log. Right: the learning rate that would give the same |Δ output| as width 16.

One learning rate for every width

In the standard parametrisation the same step changes the output 24.0 times more at width 1 024 than at 16, a log–log slope of 0.74: each of n readout weights moves by about the learning rate and their effects add. To keep the change constant the learning rate must fall with width, so a rate tuned on a narrow model is far too large for a wide one. In the μP-style parametrisation the output multiplier and the scaled readout rate cancel that growth: the change varies by only 1.23 times across the same range, and the learning rate that suits the small network suits the large one.

That is the whole value of μP: it makes small models faithful proxies for large ones, so hyperparameters can be tuned where tuning is cheap. Yang and colleagues transferred settings across a factor of about 168 in parameters. The idea has spread in several forms — some labs use μP itself, others fit per-width scaling rules empirically — and later work found that with careful per-layer choices even standard parametrisation can be made to transfer. What does not change is the principle: decide how every hyperparameter should scale with width before trusting a small run to choose it.

μP addresses width. Depth, batch size and training length change the best settings too, and extensions for depth (scaling residual branches by the depth) follow the same logic of keeping each update’s effect constant.

Terms

Standard parametrisation (SP): Initialisation scaled by fan-in and one learning rate for all layers.
μP: The maximal-update parametrisation: width-dependent initialisation, multipliers and per-layer learning rates.
Hyperparameter transfer: Using settings tuned on a small model for a large one.
Proxy model: A small model trained with the same recipe to tune or predict a large one.

In μP a step changes the output by the same amount at every width, so a learning rate tuned on a small proxy transfers to the large model.

Sources

Yang et al., “Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer”, arXiv:2203.03466, 2022.
Everett et al., “Scaling Exponents Across Parameterizations and Optimizers”, ICML 2024.