ICML2026

Rethinking Convergence in MoE Training: The Role of Routing Sparsity

Weihao Zhu, Long Shi, Kang Wei, Zhe Wang, Yipeng Zhou, Haixia Zhang

摘要

In Mixture-of-Experts (MoE) training, sparse routing, i.e., activating only the top-KK experts per token, is essential for balancing convergence speed and computational cost. However, existing works typically choose KK empirically, without theoretical guidance. To address this gap, we characterize the convergence behavior of MoE training using stochastic optimization theory. Specifically, we derive a convergence upper bound of O(1+M/KT)\mathcal{O}\left(\frac{1+M/K}{\sqrt{T}}\right), where TT is the number of training iterations and MM is the total number of experts per MoE layer. This result guarantees convergence and shows that increasing KK can accelerate training. By further fixing the total computational budget RR (in FLOPs), we obtain a refined bound of O(KR+MKR)\mathcal{O}\left(\sqrt{\frac{K}{R}} + \frac{M}{\sqrt{K R}}\right), which is convex in KK and implies the existence of an optimal K[1,M]K^{*}\in[1,M] that achieves the best convergence performance. Extensive experiments validate our theoretical analysis under diverse settings.