Un Cadre Théorique pour l'Équilibrage de Charge Sans Fonction Auxiliaire des Mélanges d'Experts Creux dans les Modèles d'IA à Grande Échelle

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

Dans l'entraînement d'IA à grande échelle, les couches Sparse Mixture-of-Experts (s-MoE) permettent une mise à l'échelle en n'activant qu'un petit sous-ensemble d'experts par token. Un défi opérationnel dans cette conception est l'équilibrage de charge : router les tokens pour minimiser le nombre d'experts inactifs, ce qui est important pour une utilisation efficace des GPU (coûteux). Nous fournissons un cadre théorique pour analyser la procédure d'Équilibrage de Charge Sans Perte Auxiliaire (ALF-LB) – proposée par Wang et al. (2024) de DeepSeek – en la modélisant comme une méthode primale-duale à une étape par itération pour un problème d'affectation. Premièrement, dans un cadre déterministe stylisé, notre cadre produit plusieurs propriétés structurelles perspicaces : (i) une amélioration monotone d'un objectif Lagrangien, (ii) une règle de préférence qui déplace les tokens des experts surchargés vers les experts sous-chargés, et (iii) une garantie d'équilibrage approximatif. Ensuite, nous incorporons la nature stochastique et dynamique de l'entraînement d'IA en utilisant une formulation généralisée d'optimisation en ligne. Dans le cadre en ligne, nous déduisons une propriété de forte convexité de l'objectif qui conduit à une borne de regret attendu logarithmique sous certains choix de taille de pas. De plus, nous présentons des expériences réelles sur des modèles DeepSeekMoE à 1 milliard de paramètres pour compléter nos résultats théoriques. Ensemble, ces résultats construisent un cadre principié pour analyser l'Équilibrage de Charge Sans Perte Auxiliaire des modèles s-MoE dans l'IA.

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.