Elastic routing improves sparse mixture-of-experts performance and stability
Smoothing the Top-k Exposure Boundary for Sparse Mixture-of-Experts
Computation and LanguageMachine Learning
Summary
Sparse Mixture-of-Experts models pick a fixed number of experts to handle each token, but this can cause problems when experts have similar scores—some get full attention and others none. The authors propose Elastic Expert Routing, which randomly varies the number of active experts around that fixed number during training, smoothing the sharp cutoff. This new method keeps the same computational cost but improves model learning and performance in different training scenarios.
What this means in practice
- •For machine learning engineers: Improve model training stability and downstream accuracy for large language models using sparse mixture-of-experts architectures.
- •For ai infrastructure teams: Maintain fixed inference compute cost while improving efficiency and robustness of expert routing in production models.
Authors
Yunkai Chai, Tong Zhu, Xiaoye Qu, Xuyang Hu, Guanjie Chen, Qipeng Guo, Yu Cheng
Abstract
Sparse Mixture-of-Experts models scale parameter capacity efficiently while maintaining a fixed compute budget per token. However, traditional training paradigms enforce a static choice of top-$k$ experts, which converts a continuous routing distribution into a rigid step function. This constraint introduces a brittle boundary where highly competitive experts are arbitrarily separated into full-supervision and zero-feedback zones based on minor score fluctuations. To address this issue, we propose Elastic Expert Routing, which stochastically samples the active expert budget from a localized discrete distribution centered at $k$. Over multiple training iterations, this mechanism softens the sharp threshold into a gradual probability distribution. Because the sampling neighborhood remains symmetric, this approach matches the expected computational cost of deterministic training, while preserving the inference budget. Extensive experiments demonstrate the efficacy of our method on both supervised fine-tuning and from-scratch pretraining settings. During supervised fine-tuning, elastic routing improves downstream macro-averages on OLMoE-1B-7B and Qwen3-30B-A3B by $+0.84$ and $+2.02$ points, respectively. In addition, in from-scratch pretraining, it outperforms the static top-$k$ baseline by $1.6$ points on average across downstream tasks.