ID Balancing: Stable Training of Extremely Sparse MoE via PID-Based Load Control

Peng Jin ,

Zihan Qiu ,

Zekun Wang ,

Bo Zheng ,

Yang Xu ,

Tian Xie ,

Xiao Li ,

Huaqing Zhang ,

Haoran Lian ,

Rui Men ,

Dayiheng Liu

LLM PTSAD MoE ML POA MEM

265

2026年09月30日

通过混合专家（Mixture-of-Experts, MoE）机制扩展大语言模型（LLMs），可在几乎保持每词元计算量不变的前提下，实现参数量的大幅增长。然而，若要进一步扩大参数规模，则需采用愈发稀疏的路由策略，而这将导致专家负载失衡问题日益加剧。此类失衡不仅降低参数的实际利用率与训练效率，还可能损害训练稳定性，从而成为制约模型可靠扩展的关键瓶颈。本文中，我们将两种具有代表性的、无需辅助损失（auxiliary loss）的方法统一纳入比例-积分-微分（Proportional-Integral-Derivative, PID）控制器的理论框架下，视其为不完整的PID控制器：DeepSeek提出的无损方法等价于一个步长固定的积分控制器；而Kimi K3提出的分位数均衡（Quantile Balancing）则对应一种广义的比例控制器。基于这一控制论视角，我们进一步提出ID均衡（ID Balancing）方法——一种仅含积分项与微分项的PID控制器。该方法使积分项随负载误差线性缩放，并仅在失衡程度加剧时才激活微分项，从而对较大或持续恶化的失衡施加更强的校正力度，而在接近均衡状态时则采取更微小的更新步长。我们在Top-10、Top-5和Top-3三种路由策略下，于包含768个专家的MoE架构上对ID均衡进行了全面评估。结果表明，在Top-3设置中，相较于当前最优基线方法，ID均衡可将骨干网络最差情况下的最大违反率（backbone MaxVio）降低逾50%，同时将训练全程平均的最小违反率（backbone MinVio）降低12%。当模型总参数量从18.9B提升至69.9B（即在768个专家中选取Top-10路由）时，ID均衡下的骨干网络最差情况MaxVio基本保持稳定，且较依赖辅助损失的基线方法降低了约89.6%。此外，ID均衡在语言建模能力及各类下游任务性能上亦保持竞争力。尤为值得注意的是，随着路由稀疏度不断提高，ID均衡的优势愈发显著，因而有望成为支撑更大规模、更高稀疏度MoE模型可靠扩展的一项有力技术方案。

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.

许愿