Hyperparameter rules adapt to model sparsity for mixture-of-experts
Hyperparameter Scaling Laws Across MoE Sparsity
Machine LearningArtificial IntelligenceComputation and Language
Summary
Training large AI models with many parts activated only some of the time (called sparsity) is tricky because the usual settings for learning speed and batch size don't work well. The authors ran many experiments and found that how often parts of the model activate changes the best training settings in predictable ways. They created formulas that help pick good learning speeds and batch sizes no matter how sparse the model is. This helps make training these big sparse models more reliable and efficient.
What this means in practice
- •For machine learning engineers: Set learning rates and batch sizes for training sparse Mixture-of-Experts models using the derived scaling laws for better efficiency.
- •For cloud ai service providers: Optimize training configurations across different sparsity levels in large AI models to reduce computational costs and improve training throughput.
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.