Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts
Nayeon Kim, Hojin Lee, Yunju Bak, Jaesun Park, Boseop Kim
COLM 2026
cs.LG, cs.AI, cs.CL
2026-08-20
Kakao adapts μP for MLA MoE with Muon, then fits a log-log LR-vs-tokens law on a 10.8B proxy to predict 3.85e-4 for a 155B/17B-active model at 10T (R²=0.95).
The best learning rate for large MoE pretraining moves when width or token budget changes. A joint 2D sweep over both axes is unaffordable past 100B parameters. μP was built for dense width transfer. MoE adds a sparsity axis: grow the expert count, keep the active set fixed. Prior MoE μP studies mostly froze expert counts and scaled hidden size, which does not match fine-grained MoE with 100+ experts.
Kakao and Upstage split the job. First make the optimal learning rate transferable across width, including total experts. Then extrapolate it along tokens from a cheap proxy.
The stack is MLA MoE with Muon. μP labels tensors by how many of their dimensions grow with width. Embeddings, biases, and expert FC2 are vector-like (initialization only). FFN, attention, router, and expert FC1 are matrix-like (initialization and learning rate both scaled by faninbase/fanin). Depth stays fixed. Head dim stays fixed; head count grows with hidden size. When scaling up, active experts and expert intermediate dim stay put; total experts and hidden size grow together. Sparsity does not add extra fan-in or fan-out changes per expert, so the same μP rule still applies. Running the proxy at the target sparsity would also waste hardware: arithmetic intensity is too low.
They never sweep 10T tokens. Proxies run the stable phase of WSD without decay, merge weights with EMA (α=0.6) every 2B tokens, and score every 10B. At each budget, validation loss versus log learning rate is a quadratic; the vertex is η. Then log η versus log tokens is a linear fit. Batch size is left out of the transfer: it is a throughput knob, and papers disagree on how it should scale. They pin batch size for GPU efficiency and only extrapolate learning rate. The fit uses points after the batch-size increase, from 255B tokens onward.
Under standard parameterization, the best learning rate of a 0.6B-total / 0.3B-active base does not transfer to 2×, 4×, and 8× width (2.2B/0.7B, 8B/1.5B, 30.7B/3.6B). Under μP it does. The same pattern holds for dense MLA models.
A 5.6B/1.8B proxy and a 2× 20.7B/3.8B model, trained to 40B–100B tokens, share parabola vertices and curvature. The optimum drifts down slightly as tokens grow. The 10T extrapolation uses a 10.8B/3.3B proxy (one-quarter the width of the 155B/17B target) trained for 500B tokens. The log-log fit on 255B–502B has R²=0.95 and predicts 3.85×10⁻⁴ at 10T. A held-out check that fits 255B–350B and predicts near 500B is off by about 4.4% on average.
| Setup | Metric | Result |
| 10.8B proxy, 255B–502B | log-log R² | 0.95 |
| Extrapolated to 10T | Optimal LR | 3.85×10⁻⁴ |
| Held-out near 500B | Predicted vs actual | 4.4% mean gap |
| Width sweep vs proxy only | Extra compute | 240.3 ZFLOPs on top of 64.8 |
| 155B target vs proxy | Total compute | 98× |
The 155B-total / 17B-active model then trains for 10T tokens. Batch size goes from 8M to 32M after 200B. Loss never spikes. The mix starts at 45% English, 12.5% Math/STEM, 27.5% code, 15% multilingual, then 22.5/27.5/25/25 after 6T. On a shared harness, MMLU-Pro sits near 64%, above dots.llm1 and GLM-4.5-Air and below Hunyuan-A13B and DeepSeek-V4-Flash, on the Pareto front. Exact bar heights in Figure 7 are not tabulated.
Past 100B MoE, you cannot grid-search learning rate at target scale. This pipeline turns a 2D sweep into μP across width plus a 1D token extrapolation. Proxy search is about 1/98 of target pretraining compute. Teams already on MLA and Muon can reuse the tensor taxonomy and the EMA checkpoint trick. It is a practical recipe, not a new optimizer.
The recipe is tied to MLA and Muon. Sparsity and width grow together, so the two axes are not isolated. Per-expert learning rates are left for later, even though top-k routing gives experts uneven effective batches. There is no 10T sweep to prove the predicted 3.85×10⁻⁴ is optimal; stable loss and a competitive Pareto point only show that it works. Stage 1 data favors diversity over quality filters, so the benchmarks are not a final model card. Batch size was deliberately excluded; whether the token law still holds after a hardware-driven batch change is open.