大規模AIモデルにおけるスパース混合エキスパートの補助損失不要ロードバランシングの理論的枠組み

A Theoretical Framework for Auxiliary-Loss-Free Load Balancing of Sparse Mixture-of-Experts in Large-Scale AI Models

December 3, 2025

著者: X. Y. Han, Yuan Zhong

cs.AI

要旨

大規模AIトレーニングにおいて、Sparse Mixture-of-Experts（s-MoE）層は、トークンごとに専門家の小さなサブセットのみを活性化することでスケーリングを可能にする。この設計における運用上の課題は、負荷分散である。つまり、（高価な）GPUの効率的な利用にとって重要な、アイドル状態の専門家の数を最小化するようにトークンをルーティングすることである。本論文は、DeepSeekのWang et al. (2024) によって提案されたAuxiliary-Loss-Free Load Balancing (ALF-LB) 手法を、割り当て問題に対する反復毎1ステップのprimal-dual法として定式化することにより、分析するための理論的枠組みを提供する。まず、様式化された決定論的設定において、我々の枠組みはいくつかの洞察に富む構造的特性を導き出す：(i) ラグランジュ目的関数の単調改善、(ii) 過負荷の専門家から負荷不足の専門家へトークンを移動させる選好規則、(iii) 近似的な負荷分散の保証である。次に、一般化されたオンライン最適化の定式化を用いて、AIトレーニングの確率的かつ動的な性質を組み込む。オンライン設定では、目的関数の強い凸性を導出し、それがあるステップサイズの選択の下で対数的な期待リグレット界をもたらすことを示す。さらに、理論的知見を補完するため、10億パラメータのDeepSeekMoEモデルを用いた実際の実験結果を提示する。これらの結果を総合することで、AIモデルにおけるs-MoEのAuxiliary-Loss-Free Load Balancingを分析するための原理に基づいた枠組みが構築される。

English

In large-scale AI training, Sparse Mixture-of-Experts (s-MoE) layers enable scaling by activating only a small subset of experts per token. An operational challenge in this design is load balancing: routing tokens to minimize the number of idle experts, which is important for the efficient utilization of (costly) GPUs. We provide a theoretical framework for analyzing the Auxiliary-Loss-Free Load Balancing (ALF-LB) procedure -- proposed by DeepSeek's Wang et al. (2024) -- by casting it as a one-step-per-iteration primal-dual method for an assignment problem. First, in a stylized deterministic setting, our framework yields several insightful structural properties: (i) a monotonic improvement of a Lagrangian objective, (ii) a preference rule that moves tokens from overloaded to underloaded experts, and (iii) an approximate-balancing guarantee. Then, we incorporate the stochastic and dynamic nature of AI training using a generalized online optimization formulation. In the online setting, we derive a strong convexity property of the objective that leads to a logarithmic expected regret bound under certain step-size choices. Additionally, we present real experiments on 1B-parameter DeepSeekMoE models to complement our theoretical findings. Together, these results build a principled framework for analyzing the Auxiliary-Loss-Free Load Balancing of s-MoE in AI models.