Optimal MoE Learning Rates Extrapolate to 10 Trillion Tokens From Small Proxy Runs at R-squared 0.95

arXiv 2608.20061·medium signalXBlueskyLinkedInCopy link

Sweeping learning rates for Mixture-of-Experts models at extreme model and token scale is computationally prohibitive, so the authors build a two-step transfer framework: a Maximal Update Parameterization adaptation for MoE with Multi-head Latent Attention and the Muon optimizer, shown to transfer optimal learning rates consistently across width-scaled models. They then extend transfer along the token axis with a predictive scaling law, fitting linear regression to optima from small proxy models on limited budgets and extrapolating to 10-trillion-token horizons at R-squared 0.95. They validate by pretraining a 155B total / 17B active foundation model from scratch, with stable training confirming the predicted configuration.