Ein theoretischer Rahmen für lastverteilung ohne Hilfsverluste bei spärlichen Mixture-of-Experts in großskaligen KI-Modellen
A Theoretical Framework for Auxiliary-Loss-Free Load Balancing of Sparse Mixture-of-Experts in Large-Scale AI Models
December 3, 2025
papers.authors: X. Y. Han, Yuan Zhong
cs.AI
papers.abstract
Bei groß angelegtem KI-Training ermöglichen Sparse Mixture-of-Experts (s-MoE)-Schichten die Skalierung, indem pro Token nur eine kleine Teilmenge von Experten aktiviert wird. Eine operationelle Herausforderung dieses Designs ist der Lastausgleich: die Zuweisung von Tokens, um die Anzahl inaktiver Experten zu minimieren, was für die effiziente Nutzung (kostspieliger) GPUs wichtig ist. Wir stellen einen theoretischen Rahmen zur Analyse des Auxiliary-Loss-Free Load Balancing (ALF-LB)-Verfahrens – vorgeschlagen von DeepSeek’s Wang et al. (2024) – bereit, indem wir es als Primal-Dual-Methode mit einem Schritt pro Iteration für ein Zuordnungsproblem darstellen. Zunächst liefert unser Rahmenwerk in einem stilisierten deterministischen Setting mehrere aufschlussreiche strukturelle Eigenschaften: (i) eine monotone Verbesserung eines Lagrange-Ziels, (ii) eine Präferenzregel, die Tokens von überlasteten zu unterausgelasteten Experten verlagert, und (iii) eine Garantie für approximativen Ausgleich. Anschließend integrieren wir die stochastische und dynamische Natur des KI-Trainings mittels einer verallgemeinerten Online-Optimierungsformulierung. Im Online-Setting leiten wir eine starke Konvexitätseigenschaft des Ziels her, die unter bestimmten Schrittweitenwahl zu einer logarithmischen Schranke für den erwarteten Regret führt. Zusätzlich präsentieren wir reale Experimente mit 1-Milliarde-Parameter DeepSeekMoE-Modellen, um unsere theoretischen Erkenntnisse zu ergänzen. Zusammengenommen bilden diese Ergebnisse einen prinzipienbasierten Rahmen zur Analyse des auxiliarverlustfreien Lastausgleichs von s-MoE in KI-Modellen.
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.