Um Referencial Teórico para Balanceamento de Carga sem Perdas Auxiliares em Mistura Especialista Esparsa em Modelos de IA de Grande 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
Resumo
No treinamento de IA em larga escala, as camadas Sparse Mixture-of-Experts (s-MoE) permitem a escalabilidade ativando apenas um pequeno subconjunto de especialistas por token. Um desafio operacional nesse projeto é o balanceamento de carga: o roteamento de tokens para minimizar o número de especialistas ociosos, o que é importante para a utilização eficiente de GPUs (custosas). Fornecemos uma estrutura teórica para analisar o procedimento Auxiliary-Loss-Free Load Balancing (ALF-LB) – proposto por Wang et al. (2024) da DeepSeek – enquadrando-o como um método primal-dual de uma etapa por iteração para um problema de atribuição. Primeiro, em um cenário determinístico estilizado, nossa estrutura produz várias propriedades estruturais reveladoras: (i) uma melhoria monotônica de um objetivo Lagrangiano, (ii) uma regra de preferência que move tokens de especialistas sobrecarregados para subutilizados, e (iii) uma garantia de balanceamento aproximado. Em seguida, incorporamos a natureza estocástica e dinâmica do treinamento de IA usando uma formulação generalizada de otimização online. No contexto online, derivamos uma propriedade de convexidade forte do objetivo que leva a um limite de arrependimento esperado logarítmico sob certas escolhas de tamanho de passo. Adicionalmente, apresentamos experimentos reais em modelos DeepSeekMoE com 1 bilhão de parâmetros para complementar nossas descobertas teóricas. Juntos, esses resultados constroem uma estrutura fundamentada para analisar o Balanceamento de Carga Livre de Perda Auxiliar em modelos de IA com 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.