Un Marco Teórico para el Balanceo de Carga sin Pérdidas Auxiliares en Mezclas de Expertos Dispersas para Modelos de IA a Gran Escala

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

December 3, 2025

Autores: X. Y. Han, Yuan Zhong

cs.AI

Resumen

En el entrenamiento de IA a gran escala, las capas de Mezcla de Expertos Dispersa (s-MoE) permiten la escalabilidad activando solo un pequeño subconjunto de expertos por token. Un desafío operativo en este diseño es el equilibrio de carga: enrutar los tokens para minimizar el número de expertos inactivos, lo cual es importante para la utilización eficiente de GPUs (costosas). Proporcionamos un marco teórico para analizar el procedimiento de Equilibrio de Carga Sin Pérdida Auxiliar (ALF-LB) —propuesto por Wang et al. (2024) de DeepSeek— reformulándolo como un método primal-dual de un paso por iteración para un problema de asignación. Primero, en un entorno determinista estilizado, nuestro marco produce varias propiedades estructurales reveladoras: (i) una mejora monótona de un objetivo Lagrangiano, (ii) una regla de preferencia que mueve tokens de expertos sobrecargados a expertos infrautilizados, y (iii) una garantía de equilibrio aproximado. Luego, incorporamos la naturaleza estocástica y dinámica del entrenamiento de IA utilizando una formulación generalizada de optimización en línea. En el entorno en línea, derivamos una propiedad de convexidad fuerte del objetivo que conduce a un límite de arrepentimiento esperado logarítmico bajo ciertas elecciones de tamaño de paso. Adicionalmente, presentamos experimentos reales en modelos DeepSeekMoE de 1B de parámetros para complementar nuestros hallazgos teóricos. En conjunto, estos resultados construyen un marco fundamentado para analizar el Equilibrio de Carga Sin Pérdida Auxiliar en modelos de IA con s-MoE.

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.