muP and Hyperparameter Transfer

Sweep learning rate on a 40-million-parameter transformer, find the optimum, scale the model to 4 billion parameters, and the optimum has moved. It moves in a consistent direction — wider models want smaller learning rates — which is why practitioners developed folklore rules of thumb for shrinking it. The folklore exists because the underlying parametrisation is broken in a specific, fixable way.

The diagnosis is about what happens to activations as width grows. Under standard parametrisation, the size of the update a layer's pre-activations receive per optimiser step depends on the layer's fan-in. Make the model wider and updates to hidden pre-activations grow with width, so a learning rate that was well-scaled at one width overshoots at another. The optimum shifts not because "big models are different" but because the parametrisation lets a per-step quantity scale with a number you are changing.

Fixing the scaling, layer by layer

The maximal update parametrisation rescales three things per layer group so that, in the infinite-width limit, every layer's activations and the updates to them stay \(\Theta(1)\) as width grows (Yang & Hu et al., 2022, Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, arXiv:2203.03466).

The three knobs are initialisation variance, per-layer learning rate, and an output multiplier, and they are set differently for input/embedding layers, hidden layers, and the output layer, with the hidden and output scalings depending on fan-in. Two consequences are visible in any muP implementation and are good diagnostics that you are looking at one: the output logits are divided by fan-in rather than left unscaled, and attention uses a \(1/d_{\text{head}}\) scale rather than \(1/\sqrt{d_{\text{head}}}\) — a change that looks like a typo against the original transformer and is deliberate, because under muP the query-key dot product is already correctly scaled by the initialisation.

muTransfer: tune small, train big

Once the optimum no longer depends on width, the tuning procedure inverts. Parametrise the target model in muP, sweep hyperparameters on a much smaller proxy with the same shape, and transfer them zero-shot to the full-size run. Yang and Hu report transferring from a 40M-parameter proxy to a 6.7B GPT-3-scale model, outperforming the published numbers for that model, with total tuning cost around 7% of pretraining compute.

What that buys is not primarily money. It is the ability to sweep at all. A large pretraining run is typically hyperparameter-tuned once, badly, because nobody can afford a grid at that scale; muTransfer makes a real sweep affordable by moving it to a model that trains in hours.

When it breaks

• Transfer is across width, not everything. Depth, batch size, sequence length, and total training tokens are separate axes, and transferring across them is either an active research question or requires additional corrections. A proxy that differs from the target in depth as well as width is not a muTransfer setup.
• Implementation is fiddly and silently wrong when partial. muP requires touching initialisation, per-parameter-group learning rates, the attention scale, and the output multiplier consistently. Get one wrong — a fused layer that keeps its default init, a weight-tied embedding treated as hidden — and you have neither standard parametrisation nor muP, just an untested third thing.
• It fixes the optimum's location, not the optimiser. muP does not make training stable by itself; loss spikes, gradient clipping, warmup, and normalisation choices remain your problem (see loss spikes and divergence).
• Verification requires the coordinate check, not the loss. The right test is that per-layer activation and update magnitudes stay flat as you widen the proxy across a few widths. A model that trains fine at one width can still be misparametrised, and you find out only at scale.
• Interaction with newer optimisers is not automatic. muP's per-layer scalings were derived for SGD and Adam; optimisers with different update geometry need their own derivation rather than the same constants.