Learning Rate Transfer

Why the optimal learning rate moves with scale

Train the same architecture at several widths and plot loss against learning rate. Under the parametrization almost every codebase ships by default — call it standard parametrization (SP) — the U-shaped curves do not line up. As width grows, the whole curve slides toward smaller learning rates: the value that was optimal at width 256 is too hot at width 4096 and diverges or plateaus.

The reason is dimensional. A pre-activation is a sum over fan_in terms, h_j = Σ_i W_ji x_i. When a gradient step nudges every W_ji, the change in h_j is itself a sum over fan_in contributions, so its magnitude grows with width unless something scales it down. A fixed learning rate that produced an O(1) update at small width produces an over-large one at large width, so you must keep dialing the rate back down. That is the drift.

Zero-shot hyperparameter transfer, stated plainly

Zero-shot hyperparameter transfer means you fix the hyperparameters on a small model and reuse them, unchanged, on a large one — no tuning at the target scale at all. It only works if the quantity you tuned has the same optimum at both scales, so the real engineering task is not clever search; it is reparametrizing the network so that the optimum stops moving. Pick a small proxy, sweep the learning rate to the bottom of its loss-vs-LR curve, and apply that value to the full-size model. If the curves for every width share a common minimum, the proxy’s answer is the target’s answer — and a sweep that costs a rounding error on a 40M model replaces one that would cost more than the entire training run on a 7B model.

muP: the maximal update parametrization

muP is a specific set of width-dependent choices for three things per layer: the initialization variance, a fixed multiplier on the layer’s output, and the learning rate. The design goal is a stability condition: as width increases, (1) every activation stays O(1), (2) the change in each activation after one optimizer step stays O(1), and (3) the logits stay O(1) — nothing blows up, nothing decays to zero.

‘Maximal’ refers to condition (2): muP is the parametrization in which the hidden features actually learn — move by an O(1) amount — in the wide limit, rather than freezing near their initialization the way a lazy, kernel-like regime would. That makes a wide model behave like a scaled-up narrow one, so the LR that steers the narrow one steers the wide one too.

The width scalings that pin the optimum

muP treats layers by their role: input/embedding, hidden, and output (readout) layers each get their own scaling in the network width n, because a pre-activation formed from a width-n input behaves differently from one formed from a fixed-size token embedding.

For the case that matters most in practice — hidden weights trained with Adam — the muP learning rate scales like 1 / fan_in, inversely with width; the output-layer multiplier is scaled by 1 / fan_in and its init reduced, while input/embedding layers keep an O(1) learning rate. These factors are folded into the parametrization, so the single number you tune — the base learning rate — is width-independent: the knob you sweep on the proxy is the knob you set on the target. Get any one of the three scalings wrong and transfer quietly breaks.

The muTransfer recipe

Put together, the workflow is short. Express the model in muP against a chosen base width; build a proxy by shrinking only the width (keep depth, data, and the rest fixed); sweep the learning rate on the proxy and read off the optimum; then instantiate the full-size model in the same muP and apply the proxy’s values directly.

The base width is a bookkeeping anchor — the width at which the muP scalings equal one, so the numbers you sweep stay interpretable. The proxy must be wide enough to be in the asymptotic regime: 128 or 256 is usually fine, 16 is not. Hold everything you do not intend to transfer identical between proxy and target, so width is the only thing that changed.

A worked example

Suppose you plan to train a 6.7B model with hidden width d = 4096, and a full sweep at that size costs twenty runs you cannot afford. Choose a base width of 256 and a proxy at that width — roughly a 40M-parameter model. Sweep the muP learning rate on the proxy over, say, {5e-4, 1e-3, 2e-3, 4e-3, 8e-3} and find the loss minimum at η* = 2e-3.

proxy:   width 256   -> sweep LR -> best muP LR = 2e-3
target:  width 4096  -> set  muP LR = 2e-3   (no sweep)

internally, muP scales the hidden-layer step by 1/fan_in:
  effective hidden LR(256)  ∝ 2e-3 / 256
  effective hidden LR(4096) ∝ 2e-3 / 4096
the base number 2e-3 is unchanged; muP absorbs the 16× ratio.

You set 2e-3 on the 6.7B model and train once. The twenty exploratory runs happened on a model that trains in minutes, not one that costs a cluster-week — search cheap, commit expensive.

What transfers, and what does not

Width is the axis muP is built for, and LR transfer along width is the robust, well-tested case. muTransfer also carries other hyperparameters across width in practice: initialization scale, the layer/output multipliers, and to a good degree Adam’s β terms and the LR schedule shape.

The important caveats are about the other axes. Transfer across depth is far more fragile: plain muP does not make the optimum depth-invariant, and doing so needs extra residual-branch scaling (each block by roughly 1/√depth, the idea behind depth-aware variants). Transfer across batch size and training length is only approximate, since the optimal LR has its own mild dependence on both. The safe reading: transfer across width with confidence, keep depth and the data recipe fixed between proxy and target, and re-check if you change those.

Implications for small and CPU-trained models

Transfer is most valuable exactly when compute is scarce, and it always points the cheap way: tune on the small model. For someone training a small language model on a CPU or a single modest GPU, that is the whole game — run a real learning-rate sweep on a tiny proxy, then commit your budget to one well-chosen run instead of gambling on a guessed LR.

There is a second, quieter benefit: muP tends to make training more stable even at a single scale, because the O(1)-activation condition keeps pre-activations and logits from silently growing — fewer mystery divergences and less loss-spike firefighting when you cannot babysit a run.

A crude fallback when muP is too much

Implementing full muP is a real code change, and sometimes you just want a rough transfer without it. The blunt heuristic is to scale the learning rate as η ∝ 1 / √width when you widen the model: it captures part of the drift SP leaves in, and it beats holding the LR fixed. It is not principled the way muP is — it does not touch initialization or multipliers, so the transfer degrades as the width gap grows — but for a modest jump it often lands close enough to save a sweep.

Treat it as a stopgap. If you are scaling by a large factor, or if training stability matters, pay the one-time cost of real muP: the square-root rule is a patch on the symptom; muP removes the cause.

Common pitfalls

The failure modes are almost always incomplete or inconsistent muP rather than a flaw in the idea. The frequent ones: scaling the learning rate but forgetting the init and multiplier changes, so activations still fan out with width and transfer silently misses. Choosing a proxy that is too narrow to be in the asymptotic regime, so its optimum is a small-width artifact. Changing depth or the data mixture between proxy and target and expecting the width theory to cover it.

Two subtler ones: weight decay interacts with the LR scaling — decoupled decay behaves differently under muP than coupled L2 decay, so pick one and keep it consistent across scales. And confusing the base width definition between sweep and target rescales every number by a constant and quietly detunes the transfer. The discipline that avoids all of it: define muP once and change only width between proxy and target.