Een Theoretisch Kader voor Hulpverliesvrije Taakverdeling bij Sparse Mixture-of-Experts in Grootschalige AI-modellen

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

December 3, 2025

Auteurs: X. Y. Han, Yuan Zhong

cs.AI

Samenvatting

Bij grootschalige AI-training maken Sparse Mixture-of-Experts (s-MoE)-lagen schaalvergroting mogelijk door per token slechts een kleine subset van experts te activeren. Een operationele uitdaging in dit ontwerp is *load balancing*: het routeren van tokens om het aantal inactieve experts te minimaliseren, wat belangrijk is voor een efficiënt gebruik van (kostbare) GPU's. Wij bieden een theoretisch kader voor het analyseren van de Auxiliary-Loss-Free Load Balancing (ALF-LB)-procedure – voorgesteld door DeepSeek's Wang et al. (2024) – door deze te modelleren als een primaal-duale methode met één stap per iteratie voor een toewijzingsprobleem. Ten eerste levert ons kader in een gestileerde deterministische setting verschillende inzichtelijke structurele eigenschappen op: (i) een monotone verbetering van een Lagrangiaan doel, (ii) een voorkeursregel die tokens van overbelaste naar onderbelaste experts verplaatst, en (iii) een garantie voor een benaderd evenwicht. Vervolgens incorporeren we het stochastische en dynamische karakter van AI-training met behulp van een gegeneraliseerde online-optimalisatieformulering. In de online-setting leiden we een sterk convexiteitseigenschap van het doel af die leidt tot een logaritmische verwachte regret-bound onder bepaalde keuzes van de stapgrootte. Daarnaast presenteren we echte experimenten op DeepSeekMoE-modellen met 1B parameters om onze theoretische bevindingen aan te vullen. Gezamenlijk vormen deze resultaten een principieel kader voor het analyseren van Auxiliary-Loss-Free Load Balancing in s-MoE AI-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.