Теоретическая основа сбалансированной работы разреженных смесей экспертов в крупномасштабных моделях искусственного интеллекта без вспомогательных функций потерь
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
Аннотация
При обучении крупномасштабных моделей ИИ разреженные смеси экспертов (s-MoE) позволяют масштабировать модель, активируя для каждого токена лишь небольшое подмножество экспертов. Операционной проблемой такой архитектуры является балансировка нагрузки: маршрутизация токенов для минимизации числа простаивающих экспертов, что важно для эффективного использования (дорогостоящих) GPU. Мы предлагаем теоретическую основу для анализа процедуры балансировки нагрузки без вспомогательных потерь (ALF-LB), предложенной Wang et al. (2024) из DeepSeek, интерпретируя её как прямой-двойственный метод с одним шагом на итерацию для задачи назначения. Сначала, в детерминированной постановке, наша модель выявляет несколько содержательных структурных свойств: (i) монотонное улучшение лагранжевой целевой функции, (ii) правило предпочтения, перемещающее токены из перегруженных экспертов в недогруженные, и (iii) гарантию приближённой балансировки. Затем мы учитываем стохастическую и динамическую природу обучения ИИ с помощью обобщённой формулировки онлайн-оптимизации. В онлайн-режиме мы выводим свойство сильной выпуклости целевой функции, которое приводит к логарифмической границе ожидаемого регрета при определённых выборах размера шага. Кроме того, мы представляем реальные эксперименты на моделях DeepSeekMoE с 1 миллиардом параметров, дополняющие наши теоретические выводы. В совокупности эти результаты формируют принципиальную основу для анализа балансировки нагрузки без вспомогательных потерь в 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.