Hyperparameter Scaling Transfer

The tuning tax outgrows the model

Hyperparameter search cost scales with the thing you are searching over. A 20-point learning-rate sweep on a 6.7B model trained for 300B tokens is twenty full pretraining runs — nobody does that. So in practice large models get tuned by extrapolation: plot the optimal learning rate found at 125M, 350M and 1.3B, fit something, and hope the fit holds two orders of magnitude further out.

That hope is usually misplaced, and the reason is not noise. Under the parametrization almost every codebase ships with, the optimum genuinely moves as width grows, and it keeps moving. You are extrapolating a curve whose shape is an artifact of a scaling convention rather than a fact about the problem. The fix is not a better fit. The fix is to change the convention so the curve is flat — so that “the best learning rate” is a width-independent constant you can measure once, on a model small enough to sweep exhaustively.

Why the optimal learning rate drifts

Take one hidden weight W: [n, n] mapping activations x: [n], with the standard parametrization (SP): initialize W_ij with variance 1/fan_in = 1/n, and use a single global Adam learning rate η. At initialization this is fine. After one step it is not.

at init:   W_ij = Θ(1/√n), independent of x
           (Wx)_i = Σ_j W_ij x_j — n terms, random signs → Θ(√n · 1/√n) = Θ(1)   OK

after one Adam step:   ΔW_ij = Θ(η)   (Adam normalizes, so every entry moves ~η)
           ΔW is built FROM x, so the terms are correlated, not random
           Δ(Wx)_i = Σ_j ΔW_ij x_j → Θ(n · η)   ← grows with width

want Δ(Wx)_i = Θ(1)  ⇒  η = Θ(1/n)

The cancellation that saves the forward pass at initialization does not survive training, because the update is a function of the very activations it multiplies. So a fixed η makes the per-step change in every hidden preactivation grow linearly in width. Widen the model and the same learning rate becomes proportionally more violent — which is exactly why the empirical optimum slides toward zero roughly like 1/n.

Three weight classes, one invariant

The remedy generalizes the above into a design goal: hold every width-dependent quantity at Θ(1), both at initialization and after updates. Concretely, three conditions must hold simultaneously as n → ∞: activations are Θ(1); logits are Θ(1); and one optimizer step changes each preactivation by Θ(1) — not Θ(n) (blow-up) and not Θ(1/n) (a wider model that learns less per step).

Satisfying all three needs different rules for different weights, because the weights differ in which of their dimensions grows. Input weights and biases have fixed fan_in (vocabulary, feature count) and growing fan_out. Hidden weights grow in both. Output/readout weights have growing fan_in and fixed fan_out (the vocabulary). SP treats all three identically. That is the entire bug: one rule cannot satisfy three different constraints.

The μP rules

Write m_d = n / n_base for the width multiplier between your target model and the proxy you tuned on, and let σ, η be the base values found by sweeping at n_base. μP then prescribes, per class:

Three of these surprise people. The readout is initialized m_d times smaller than usual, so logits vanish at initialization — deliberately, since after training the readout becomes correlated with the final hidden state and would otherwise produce Θ(n) logits. The embedding keeps a constant learning rate while hidden layers shrink theirs. And attention uses 1/d_head scaling, because trained q and k are correlated and their dot product grows like d_head, not √d_head.

Why the optimum stops moving

The payoff is best stated as a limit. Under SP, the infinite-width limit is a kernel limit: features are effectively frozen at their initial values and the network only fits their linear readout, so finite-width models are all approaching a degenerate object at different rates — and their optima chase that approach. Under μP, the width limit is a feature-learning limit: hidden representations still move by Θ(1) no matter how wide the model is.

Because every width-dependent quantity has been normalized to Θ(1), the training dynamics a width-1024 model sees at learning rate η are, to leading order, the same dynamics a width-8192 model sees at the same η. Plot loss against log η for a family of widths and the curves under SP fan out, each minimum sitting left of the last. Under μP the minima stack on top of one another; the curves shift down as width grows, which is the behaviour you actually wanted, but they do not shift sideways.

The recipe, step by step

The procedure (usually called μTransfer) is short, and the discipline is in holding everything except width fixed.

1. Instrument the model in μP. Tag every parameter with its class, define m_d against a base width, and apply the table above — including the per-class optimizer parameter groups and the output multiplier.
2. Verify with a coordinate check. Train a handful of steps at several widths and plot the mean absolute entry of each activation tensor against width. Correct μP gives flat lines; anything that fans out or collapses is a misclassified tensor or a missing multiplier.
3. Sweep on the proxy. Shrink only width. Keep depth, data, batch size, sequence length and step count as close to the target as you can afford, then random-search η, σ, α_out, α_emb and warmup.
4. Scale and run once. Set m_d to the real ratio, keep the winning base values verbatim, and launch the target run.

Worked example: width 256 → 4096

Say the proxy is a 12-layer model at n_base = 256 and the target is the same 12 layers at n = 4096. Then m_d = 4096/256 = 16. Suppose the proxy sweep’s best point is η = 2×10⁻³, σ = 0.02, α_out = 1.0. Applying the table:

hidden LR     = 2e-3 / 16      = 1.25e-4
readout LR    = 2e-3 / 16      = 1.25e-4
embedding LR  = 2e-3           (unchanged)
hidden init   σ = 0.02/√16    = 0.005
readout init  σ = 0.02/16     = 0.00125
logit scale   = 1.0/16         = 0.0625

Note what you did not do: no sweep at width 4096, no trend fit, no “we usually halve the LR when we double the model.” The target run’s hyperparameters are a deterministic function of the proxy’s and one integer. If the target had also doubled d_head, the 1/d_head attention scale would follow automatically from the same rule — there is no second decision to make.

What transfers, and what only sort of

Transfer is strong and well-supported across width, and it covers the optimization hyperparameters: learning rate, momentum, initialization scale, the multipliers, and warmup fraction. Everything else carries a caveat.

Depth transfers only partially. Plain μP says nothing about the number of layers; residual-branch magnitudes still accumulate with L. Scaling each residual branch by 1/√L (“depth-μP”) makes it behave far better in pre-LN transformers, but treat depth transfer as empirical rather than guaranteed. Batch size, sequence length and step count usually transfer in practice but are not covered by the width argument — verify at two points. Regularization does not transfer at all: dropout and weight decay trade off against dataset size and overfitting pressure, which a proxy run does not reproduce. With AdamW, note also that fixing λ while shrinking η shrinks the effective decay ηλ with width; holding ηλ constant instead is a common and reasonable correction.

The cost arithmetic

The reason to care is budget. Use the standard C ≈ 6·N·D meter for training FLOPs, and price the two options.

target run:  N = 6.7e9, D = 3e11   →  C = 6 · 6.7e9 · 3e11  ≈ 1.2e22 FLOPs
proxy run:   N = 4.0e7, D = 4e9    →  C = 6 · 4.0e7 · 4e9   ≈ 9.6e17 FLOPs

200-sample proxy search  ≈ 1.9e20 FLOPs  ≈ 1.6% of ONE target run

Two hundred configurations — far more than anyone would attempt at scale — cost a rounding error against a single target run, because the proxy is both smaller and trained on fewer tokens and the savings multiply. The μTransfer paper reports a comparable figure end to end, tuning a 6.7B GPT-3 for roughly 7% of its pretraining compute including the proxy setup work. The precise percentage depends on your sweep budget; the structural point is that once the optimum is width-invariant, search happens where compute is cheap and the expensive run is executed exactly once.