How to Set the Learning Rate for Large-Scale Pre-training?

TL;DR

Formalizes two paradigms for choosing the learning rate (LR) at large pretraining scale — Fitting (extrapolate a scaling law for the optimal LR from cheap small-scale runs) and Transfer (use µP/µTransfer to compute the target LR from a proxy model) — and runs the first apples-to-apples comparison at industrially-relevant scale: a 12B-total / 1.3B-active MoE trained on 500B tokens under WSD. Fitting wins, both at 4B/500B and 12B/500B. The paper also argues that modern architectural choices (QK-Norm in particular) have eroded µP’s classical justification: under SP the attention-logit “blow up” that µP was designed to prevent simply doesn’t happen, so module-specific LRs no longer buy you measurable performance.

Key claims

The optimal LR under WSD’s constant-LR stable phase follows a clean power-law in model size N and data size D; fitting the loss-vs-LR curve as a quadratic at each (N,D) plus an L(D) power law cuts search complexity from O(N·D·LR) to O(N·LR) [§3.1, Eqs. 2–6].
At 4B-total / 530M-active MoE trained from scratch on 500B tokens, the Fitting-paradigm LR beats µTransfer’s LR throughout training on downstream evals [Fig. 3].
At 12B-total / 1.3B-active MoE trained on 500B tokens — a >10× scale jump beyond prior hyperparameter studies — Fitting again outperforms µTransfer at every checkpoint [Fig. 4].
After WSD decay (10% of stable LR, +100B high-quality tokens) the Fitting model beats the µTransfer model by +1.28% on MMLU and +2.23% on CMMLU [Fig. 11].
Module-level (Embeddings / LM Head / Router / Hidden) greedy LR search around the global optimum yields negligible loss improvement, undermining µP’s premise that modules need separate LRs to balance feature learning [§6.1, Figs. 5–6].
Under AdamW the update magnitude across modules stays near 0.2 throughout training and is uniform across layers, consistent with rotational-equilibrium results; this is offered as a direct mechanistic reason module-specific LRs are unnecessary [§6.2, Fig. 7].
Under standard parametrization the “attention-logit blow-up” µP was designed to suppress does not occur in modern Qwen3-MoE-style architectures; ablating QK-Norm restores the instability µP reported, so QK-Norm (not parametrization) is doing the stability work [§6.3, Figs. 8–9].
Training-data size affects internal-state stability more sharply than model size does — the LR scaling exponent on D is larger in magnitude than on N for this reason [§6.4, Eq. 6, Fig. 10].
µTransfer’s effective horizon is short; extending it requires explicit weight-decay scaling (held proportional to LR) and token-horizon-aware initialization rules, both folded into the µTransfer baseline here so the comparison is fair [§3.2, Table 2 / Table 10].

Method

The Fitting paradigm: train a grid of small Qwen3-MoE proxies (550M, 1B, 2B, 3B total params; constant width/depth aspect ratio) over a discrete LR grid for ~120B tokens each under WSD. For each (N, D) point, fit L(LR) as a quadratic to find LR*(N, D); fit L(D) as a power law to extrapolate from 120B → 500B; then fit LR*(N, D) itself as a function of N and D. The final scaling law is plugged in at the target 4B and 12B scales. Critically the fit is over the stable phase of WSD, not cosine, which is what production pretraining actually uses.

The Transfer paradigm: instantiate µP on a 2B MoE proxy (Complete-µP rules for depth scaling with residual factor 1/√(2L), weight decay held proportional to LR, token-horizon-aware initialization from Complete-µ𝒟P), search LR and initialization-std on the proxy, then apply the µP transfer rules to project hyperparameters to the 4B/12B targets.

Both paradigms are then run from-scratch at 4B-A0.5B and 12B-A1.3B on 500B tokens, with periodic MMLU/CMMLU evaluation and a final WSD-decay phase on a high-quality corpus.

The mechanistic analyses (§6) repeat µP’s coordinate-check style instrumentation on the resulting target models: track per-module update magnitudes (§6.2), attention-logit norms across scale (§6.3), and the same across data horizon at fixed model size (§6.4). The QK-Norm ablation in §6.3 is the load-bearing experiment for the “µP’s assumptions are stale” claim.

Results

4B-A0.5B / 500B tokens, from scratch: Fitting > µTransfer on intermediate MMLU/CMMLU throughout training [Fig. 3].
12B-A1.3B / 500B tokens, from scratch: Fitting > µTransfer at every measured checkpoint; the gap is consistent, not noise [Fig. 4].
Post-decay (after additional 100B high-quality tokens, LR → 10%): Fitting beats µTransfer by +1.28% MMLU and +2.23% CMMLU [Fig. 11].
Module-level fine-tuning: per-module greedy LR search around the global optimum produces virtually indistinguishable validation-loss curves on the 4B model trained for 120B tokens — gain is within noise [Fig. 6].
Coordinate-check under SP: attention logits stay bounded across the 550M→3B sweep, contradicting the µP motivating diagnostic; the instability does reappear when QK-Norm is removed [Figs. 8–9].
Scaling-law fit quality: R² > 0.99 reported for the LR* surface as a function of (N, D) [§3.1 + Appendix A.3.1].

Why it’s interesting

This is a direct, well-instrumented broadside at the default “use µTransfer at scale” recipe in current open-pretraining practice, and the QK-Norm ablation is the part that should travel beyond LLM pretraining. The argument generalizes: any module-wise stability heuristic (µP, layer-wise LR rules, custom inits) is implicitly betting that the architecture needs it, and modern normalization tricks (QK-Norm, RMSNorm-everywhere, the residual-stream constraints in mHC: Manifold-Constrained Hyper-Connections) keep eating that motivation. For Luma’s DiTs, the corresponding question is whether the parametrization tricks people quote from LLM pretraining still pay off once you have stable normalization and rotational-equilibrium-style update dynamics — the answer here is “test it; the prior is no.” The Fitting recipe itself (quadratic-in-log-LR + power-law-in-D, then a closed-form LR*(N,D)) is also cheap enough to run on diffusion pretraining sweeps with minor adaptation.