Un Quadro Teorico per il Bilanciamento del Carico Senza Perdite Ausiliarie nei Modelli Sparse Mixture-of-Experts su Larga Scala
A Theoretical Framework for Auxiliary-Loss-Free Load Balancing of Sparse Mixture-of-Experts in Large-Scale AI Models
December 3, 2025
Autori: X. Y. Han, Yuan Zhong
cs.AI
Abstract
Nell'addestramento di AI su larga scala, i livelli Sparse Mixture-of-Experts (s-MoE) consentono la scalabilità attivando solo un piccolo sottoinsieme di esperti per token. Una sfida operativa in questo design è il bilanciamento del carico: instradare i token per minimizzare il numero di esperti inattivi, aspetto cruciale per un utilizzo efficiente delle (costose) GPU. Forniamo un quadro teorico per analizzare la procedura di Bilanciamento del Carico Senza Perdita Ausiliaria (Auxiliary-Loss-Free Load Balancing, ALF-LB) – proposta da Wang et al. (2024) di DeepSeek – inquadrandola come un metodo primale-duale a un passo per iterazione per un problema di assegnazione. In primo luogo, in un contesto deterministico stilizzato, il nostro quadro produce diverse proprietà strutturali significative: (i) un miglioramento monotono di un obiettivo Lagrangiano, (ii) una regola di preferenza che sposta i token da esperti sovraccarichi a esperti sottocaricati, e (iii) una garanzia di bilanciamento approssimato. Successivamente, incorporiamo la natura stocastica e dinamica dell'addestramento di AI utilizzando una formulazione generalizzata di ottimizzazione online. Nell'ambito online, deriviamo una proprietà di convessità forte dell'obiettivo che conduce a un limite di regret atteso logaritmico sotto certe scelte del passo di apprendimento. Inoltre, presentiamo esperimenti reali su modelli DeepSeekMoE da 1 miliardo di parametri per integrare i nostri risultati teorici. Nel complesso, questi risultati costruiscono un quadro principiato per analizzare il Bilanciamento del Carico Senza Perdita Ausiliaria negli s-MoE nei modelli di AI.
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.