Learning-rate transfer for mixture-of-experts models from 100M to 20B+ total parameters
Difficulty 37: transfer from 5.6B to 20.7B is now reported, but two orders of magnitude across expert granularity and top-k is untested and the parametrizations disagree; harder than the replication tasks (21 to 25). Application value 46: tuning on 100M proxies instead of multi-billion runs saves real money. Leverage 34: gates tuning practice for MoE pre-training.
The problem
Can the best learning rate found on a 100M-parameter expert model be reused unchanged on one 200 times larger?
- 1Tune learning rate on a tiny MoE
- 2μP rules rescale it with size
- 3Shown from 5.6B to 20.7B
- 4Test 100M to 20B+ across designs
The problem in one line. Training a large model requires picking a learning rate, the step size used when weights are nudged downhill along the derivative of the loss. Too small wastes compute; too large makes training unstable. Trying many learning rates on a 20-billion-parameter model would cost several full training runs, so practitioners want to tune on a cheap small model and reuse the answer.
Why a recipe is needed. In a naive setup the best learning rate drifts as the model gets wider, so the small model's answer is wrong for the big one. The maximal update parametrization (μP) is a set of scaling rules, derived from the limit of very wide networks, that rescales initial weights and per-layer learning rates with width so that the best setting stays put. It is well supported for ordinary dense transformers. Mixture-of-experts (MoE) models add more size knobs: the number of experts, the size of each, and how many experts each token uses (top-k). Several groups proposed MoE versions of μP and disagree about when transfer works.
Path to understanding
foundations → problemThe research entries this problem rests on, in the order to read them. Each links to its explainer.
- Start with
- Reverse-mode automatic differentiationAlgorithm
- Fast matrix multiplicationAlgorithm
- Step 2
- Adam, AdamW and the optimizer zooAlgorithm
- The transformer architectureAlgorithm
- Step 3
- Then
- Learning-rate transfer for mixture-of-experts models from 100M to 20B+ total parametersThe problem
State of the art
Jiang et al. (arXiv:2601.20205, ICML 2026) report transfer from 51M to about 2B total parameters at a fixed budget of about 1B tokens, and leave systematic transfer across training horizon to future work. Małaśnicki et al. (arXiv:2508.09752) test a Switch-style MoE μP (8 experts, top-1) at widths 128 to 2048. Vankadara et al. (arXiv:2605.14200, preprint) argue that μP does not reliably give learning-rate transfer or monotonic improvement for MoE and propose a Maximally Scale-Stable Parameterization, tested up to about 2.5B parameters. The one result at 20B+ is Kim et al. (arXiv:2608.20061, preprint): with μP adapted to MLA attention and the Muon optimizer, the optimal learning rate of a 5.6B-total / 1.8B-active proxy transferred to a 20.7B-total / 3.8B-active target at up to 100B tokens, and the recipe was used for a 155B-total model. That is a 4x jump from an already large proxy, for one architecture; transfer across two orders of magnitude, and across expert granularity and top-k, remains untested.
Realistic applications
A confirmed rule would let teams tune MoE hyperparameters on cheap proxies before committing to a multi-billion-dollar-scale run, and a failure would show where tuning at scale cannot be avoided.
What counts as solved
Released code, configs and loss curves for the proxy grid and at least two target-scale runs per configuration, showing the transfer criterion above for the proposed parametrization (or for each of the competing ones, giving a comparison), reproducible by a third party at the proxy scale.
How it gets verified
ReproductionA training run, experiment or computation rerun by someone else.
Dataset: one public pre-training corpus with fixed data order. Metric: validation loss at the end of training for each learning rate on a 2^0.5 grid (at least 7 points around the optimum). Threshold: target-scale optimal learning rate within 1.5x of the 100M-parameter optimum, with the loss gap at the transferred rate at most 0.5% relative. Budget: the grid at proxy scale with 3 seeds (of the order of 10^18 to 10^19 FLOPs per run), and for the target at least 5 learning rates at 20B+ total parameters, which is a cluster-scale budget (each run of the order of 10^22 FLOPs or more). The Adam or Muon variant must be stated.
Claims under verification
+ add a claimNo claimed results on the Verification board for this problem yet.
Where the work is
Theory exists in several competing forms; the decisive step is expensive large-scale training.
Bottlenecks
ranked by votes- #1ComputeCluster-scale runs for several learning rates
Checking transfer needs a learning-rate grid at 20B+ total parameters.
Community grades
Sign in to add your own grade.
Connected problems
No connections yet.