ID Balancing: Stable Training of Extremely Sparse MoE via PID-Based Load Control
ID Balancing uses an integral-derivative controller to stabilize extremely sparse MoE expert load during training.
ID Balancing treats sparse MoE load control as an integral-derivative controller: the integral term scales with load error and the derivative term activates only when imbalance worsens. Across Top-10, Top-5, and Top-3 routing over 768 experts, it reduces worst-case backbone MaxVio by over 50% and training-average MinVio by over 12% versus the best baselines in the Top-3 setting. Scaling from 18.9B to 69.9B parameters leaves worst-case MaxVio nearly unchanged and about 89.6% lower than an auxiliary-loss baseline, with competitive language-modeling performance.
- Casts DeepSeek loss-free balancing and Kimi K3 quantile balancing as partial PID controllers.
- Derivative term activates only when expert load imbalance is worsening.
- Top-3-of-768 routing cuts worst-case MaxVio by over 50% versus the best baselines.
- At 69.9B parameters, MaxVio is about 89.6% below the auxiliary-loss baseline.
Full article223 words · extracted from arxiv.org · click to collapse
Scaling Large Language Models (LLMs) via Mixture-of-Experts (MoE) enables massive parameter growth with nearly constant per-token computation. However, further scaling the parameter count requires increasingly sparse routing, where expert load imbalance becomes more severe. This imbalance reduces parameter utilization and training efficiency, and can undermine training stability, becoming a bottleneck to reliable scaling. In this work, we unify two representative auxiliary-loss-free methods as incomplete Proportional-Integral-Derivative (PID) controllers: DeepSeek's loss-free method acts as a fixed-step integral controller, while Kimi K3's Quantile Balancing functions as a generalized proportional controller. Building on this control perspective, we propose ID Balancing, an Integral-Derivative controller. It scales its integral term with load error and activates its derivative term only when imbalance worsens, enabling stronger corrections for large or worsening errors and smaller updates near balance. Evaluated across Top-$10$, Top-$5$, and Top-$3$ routing over $768$ experts, ID Balancing reduces worst-case backbone MaxVio and training-average backbone MinVio by over $50\%$ and $12\%$, respectively, relative to the best baselines in the Top-$3$ setting. When the total parameter count increases from $18.9$B to $69.9$B (Top-$10$-of-$768$), ID Balancing's worst-case backbone MaxVio remains nearly unchanged and is approximately $89.6\%$ lower than that of the auxiliary-loss baseline. ID Balancing also maintains competitive language-modeling and downstream performance. The advantages of ID Balancing grow as sparsity increases, making it a promising solution for scaling larger, sparser MoE models.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.39137