Summary
Mixture-of-Experts (MoE) models use many small expert sub-models to build very large AI systems without needing a lot more computing power. But when only a small number of these experts are active during training, it becomes tricky to choose the best training settings like learning rate and batch size. The authors ran many experiments to find patterns in how these settings change with sparsity and model size. They discovered new scaling rules that predict good training choices across different sparsity levels, helping make training efficient for these large sparse models. This work also clarifies previous conflicting results and shows their rules apply even in very sparse cases with many experts.
Mixture-of-Experts (MoE)sparsityhyperparameterslearning ratebatch sizescaling lawsmodel capacityactivation ratiotraining computepower-law relationships
Authors
Changxin Tian, Kunlong Chen, Jia Liu, Ziqi Liu, Zhiqiang Zhang, Jun Zhou
Abstract
Mixture-of-Experts (MoE) models expand model capacity without a proportional increase in training compute, but increasing sparsity makes reliable hyperparameter transfer challenging. In this work, we show that conventional hyperparameter scaling laws are insufficient for ultra-sparse MoEs: the optimal learning rate and batch size vary with activation ratio, and these shifts cannot be explained by either total or activated parameter count alone. To characterize this dependence, we conduct 1,800 pre-training runs spanning six activated-parameter scales and models with up to 6B total non-embedding parameters, processing approximately 20 trillion tokens at a cost of 200,000 equivalent H800 GPU-hours. Our results reconcile conflicting findings in prior work by revealing two scaling regimes. At fixed sparsity, the optimal batch size follows a power-law relationship with training tokens $D$, whereas the optimal learning rate scales with training compute $C$ and remains robust to the allocation between model size and data. Across sparsity levels, the activation ratio $A$ enters both relationships as an additional multiplicative power-law factor. These observations lead to unified hyperparameter scaling laws that transfer across MoE sparsity levels. Large-scale evaluation shows that the scaling form outperforms alternative functional forms. On a held-out ultra-sparse MoE with 12B total parameters and only 1/64 of its experts activated, the predicted hyperparameters remain close to the observed optima, supporting joint extrapolation across model scale and sparsity. Further experiments demonstrate transfer across expert granularities and isolate the effect of activation ratio from that of total expert count.