대규모 AI 모델에서 희소 Mixture-of-Experts의 보조 손실 없이 부하 분산을 위한 이론적 프레임워크

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

December 3, 2025

저자: X. Y. Han, Yuan Zhong

cs.AI

초록

대규모 AI 학습에서 Sparse Mixture-of-Experts(s-MoE) 계층은 토큰당 소수의 전문가만 활성화하여 확장을 가능하게 합니다. 이 설계에서 발생하는 운영상의 과제는 부하 분산으로, (고비용인) GPU의 효율적 활용에 중요한 유휴 전문가 수를 최소화하도록 토큰을 라우팅하는 것입니다. 본 논문은 DeepSeek의 Wang 등(2024)이 제안한 Auxiliary-Loss-Free Load Balancing(ALF-LB) 절차를 할당 문제에 대한 1-스텝-퍼-반복 프라이멀-듀얼 방법으로 해석하여 분석하는 이론적 체계를 제공합니다. 먼저, 정형화된 결정론적 설정에서 우리의 체계는 몇 가지 통찰력 있는 구조적 특성을 도출합니다: (i) 라그랑지안 목적함수의 단조적 개선, (ii) 과부하 전문가에서 저부하 전문가로 토큰을 이동시키는 선호 규칙, (iii) 근사적 균형 보장. 이후 우리는 일반화된 온라인 최적화 공식을 사용하여 AI 학습의 확률적 및 동적 특성을 통합합니다. 온라인 설정에서 우리는 특정 스텝 사이즈 선택 하에서 로그적 기대 후회 한계로 이어지는 목적함수의 강한 볼록성 특성을 유도합니다. 또한 이론적 결과를 보완하기 위해 10억 파라미터 DeepSeekMoE 모델에 대한 실제 실험 결과를 제시합니다. 이러한 결과를 종합하면 AI 모델에서 s-MoE의 Auxiliary-Loss-Free Load Balancing을 분석하기 위한 원칙적인 체계를 구축합니다.

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.