Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

Finding the right learning rate for a 155-billion-parameter MoE model trained on 10 trillion tokens is not a problem you can brute-force. A naive two-dimensional sweep across both model size and token budget at that scale is computationally ruinous. This paper decomposes the problem into two cheaper, sequential steps whose combined proxy cost is a small fraction of the full training run.

The first step adapts Maximal Update Parameterization (μP) to MoE architectures that use Multi-head Latent Attention and the Muon optimizer. The key design choice is that scaling happens along two coupled axes simultaneously—hidden dimension and total expert count—rather than width alone. Because increasing the number of experts while holding active experts fixed doesn't change the fan-in or fan-out of any individual expert beyond what width scaling already accounts for, the same μP scaling rule applies cleanly. The experiments confirm it: the optimal learning rate transfers consistently from a 0.6B proxy to a 30.7B model under μP, while under standard parameterization it simply drifts.

The second step handles the token dimension. Rather than running separate decay-phase experiments for each token budget—which would require one full training run per data point—the method trains a single proxy run in the stable phase of a Warmup-Stable-Decay scheduler and applies Exponential Moving Average to the weights at roughly 2-billion-token intervals. Each EMA checkpoint approximates what a decayed model would look like at that token count, yielding dozens of data points from one run. Fitting a second-order polynomial to validation loss versus log learning rate at each checkpoint gives an estimated optimal learning rate per token scale. A log-log linear regression over those estimates then extrapolates to 10 trillion tokens. The fit on the proxy model achieves an R² of 0.95, and held-out validation on unseen token budgets near 500 billion tokens shows predicted-to-actual ratios within a few percent—substantially tighter than prior comparable work.

The full 155B/17B-active model trained on 10 trillion tokens with the extrapolated learning rate shows no loss spikes and lands on the Pareto frontier of MMLU-Pro accuracy versus estimated training compute among comparable open-weight MoE models. That's the practical validation the method needed.

A few honest caveats appear in the paper itself. The framework was developed specifically for MLA plus Muon; whether it generalizes to other attention variants or optimizers is left open. The sparsity and width axes are scaled jointly throughout, so their individual contributions to transfer behavior can't be cleanly separated. And the paper acknowledges that definitively proving the extrapolated learning rate is truly optimal would require sweeps costing roughly half the full training run.

The expert routing appendix is a useful bonus: it shows that balanced marginal expert load and domain-specific routing specialization are not in tension. Deeper MoE layers exhibit strong domain-dependent routing even when aggregate load imbalance is low, with code routing diverging from the marginal early and multilingual routing specializing only in the final layers.