AUXILIARY-LOSS-FREE LOAD BALANCING STRATEGY FOR MIXTURE-OF-EXPERTS

文章标题：混合专家模型（MoE）的无辅助损失负载均衡策略
作者/机构：Lean Wang¹,², Huazuo Gao¹, Chenggang Zhao¹, Xu Sun²⋄, Damai Dai¹⋄
¹DeepSeek-AI
²北京大学计算机学院，多媒体信息处理国家重点实验室

A1 主要贡献

本文旨在解决训练混合专家模型（Mixture-of-Experts, MoE）时面临的核心挑战：专家负载不均衡问题。该问题可能导致路由坍塌（部分专家训练不足）或计算开销增加。现有方法普遍采用辅助损失函数来促进负载均衡，但这种方法会引入与主要训练目标（如语言建模）冲突的干扰梯度，从而损害模型最终性能。在实际应用中，研究者不得不在负载均衡和模型性能之间进行权衡。

为了在不产生干扰梯度的情况下有效控制负载均衡，本文提出了一种名为无损均衡（Loss-Free Balancing）的全新策略。其核心思想和贡献如下：

  1. 提出无辅助损失的均衡策略：该策略不引入任何辅助损失函数。它通过在MoE的top-K路由决策前，为每个专家的原始路由分数（routing scores）添加一个专家级别的偏置项（expert-wise bias）来实现负载均衡。
  2. 动态偏置更新机制：这些偏置项会根据近期训练数据上的专家负载情况进行动态更新。具体而言，负载较重的专家的偏置会被调低，而负载较轻的专家的偏置会被调高。通过这种迭代调整，系统能持续引导token流向负载较轻的专家，从而实现均衡的路由结果。由于该过程不产生干扰梯度，因此训练过程更纯净、友好。
  3. 提升模型性能上限：由于消除了辅助损失带来的干扰梯度，无损均衡策略能够提升通过MoE训练获得的模型性能上限，打破了传统方法中“负载均衡”与“模型性能”之间的两难困境。

图1：无损均衡策略在每个训练步骤中根据“带偏置的门控分数”选择专家，并在每个训练步骤后更新该专家级偏置。
图1：无损均衡策略在每个训练步骤中根据“带偏置的门控分数”选择专家，并在每个训练步骤后更新该专家级偏置。

实验验证在高达30亿参数的MoE模型上进行，训练数据量最高达到2000亿个token。结果表明，与传统的辅助损失控制方法相比，无损均衡策略不仅取得了更好的模型性能（以验证集损失衡量），还实现了更优的全局和批次级别的负载均衡，并且天然兼容大规模MoE模型训练中常用的专家并行技术。

A3 背景知识

2.1 混合专家模型（MIXTURE-OF-EXPERTS）

主流MoE架构：当前主流的MoE架构，如【索引4，Gshard: Scaling giant models with conditional computation and automatic sharding，Dmitry Lepikhin等人，2020】、【索引3，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，William Fedus等人，2021】和【索引1，Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models，Damai Dai等人，2024】，通常将标准Transformer中的MLP层替换为MoE层。在MoE层中，采用Top-K路由机制为每个token选择专家。对于一个包含N个专家的MoE层，第t个token的输入为$u_t$，其输出$h_t$的计算方式如下：

公式1
公式1

其中，$G$是一个非线性门控函数，$e_i$是第$i$个专家的质心（centroid）。

2.2 用于负载均衡的辅助损失

辅助损失的引入：不受控制的路由策略很容易导致负载不均衡，这带来两个显著的弊端。首先，存在路由坍塌（routing collapse）的风险，即模型持续选择少数几个专家，导致其他专家无法得到充分训练【索引8，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，Noam M. Shazeer等人，2017】。其次，当专家分布在多个设备上时，负载不均衡会加剧计算瓶颈。为了解决这些问题，通常会采用一个辅助损失函数来控制负载均衡【索引3，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，William Fedus等人，22021】、【索引4，Gshard: Scaling giant models with conditional computation and automatic sharding，Dmitry Lepikhin等人，2020】。对于一个长度为T的序列，辅助损失定义如下：

公式2
公式2

其中，$N$是专家总数，$K$是为每个token选择的专家数量，$s_{i,t}$是专家$i$对token $t$的路由分数，$f_i$表示路由到专家$i$的token比例，$P_i$表示专家$i$的平均门控分数，$\alpha$是控制辅助损失强度的超参数。

负载均衡与模型性能的两难困境：上述辅助损失虽然能促进负载均衡，但作为一个额外的正则化项，它会干扰语言建模的训练。如果辅助损失系数$\alpha$过小或缺失，会导致严重的负载不均衡；而$\alpha$过大则会损害训练效果，导致次优的模型性能。为了说明这种困境，图2展示了负载均衡与模型性能之间的关系。实验中调整$\alpha$的值（1e-2, 1e-3, 1e-4, 0），并展示了相应的MaxVio_global（一种衡量负载均衡程度的指标，计算细节见4.1节）。如图所示，小的$\alpha$值导致路由坍塌，影响模型效率，并可能使某些专家学习不足；而大的$\alpha$值虽然能控制负载均衡，但显著降低了模型性能。为了打破这一困境，本文提出了无损均衡策略，它直接控制专家负载均衡，但不会引入除语言建模损失之外的任何意料之外的梯度。

图2：辅助损失控制训练中负载均衡与模型性能的两难困境。小的辅助损失系数α导致负载均衡差，而大的α则损害模型性能。相比之下，我们的无损均衡方法打破了这一困境。
图2：辅助损失控制训练中负载均衡与模型性能的两难困境。小的辅助损失系数α导致负载均衡差，而大的α则损害模型性能。相比之下，我们的无损均衡方法打破了这一困境。

A2 方法细节

引入专家级偏置项：为了找到一种不直接干扰主训练目标梯度的负载均衡替代方案，本文提出了无损均衡（Loss-Free Balancing）策略，该策略根据每个专家的均衡状况直接调整其门控分数。如图1所示，我们为每个专家的门控分数$s_{i,t}$添加一个专家级偏置项$\{b_i\}_{i=1}^{N}$，并使用调整后的带偏置分数来确定top-K选择：

公式3
公式3

需要注意的是，专家偏置项$b_i$仅用于通过影响top-K选择来调整路由策略。在计算MoE层的最终输出时，它并不会被加到用于加权所选专家输出的$g_{i,t}$中。

偏置项的迭代更新机制：为了获得合适的偏置，我们根据以下原则对每个偏置$b_i$进行迭代调整：当对应专家的负载相对较重时，减小其偏置；反之亦然。具体来说，我们持续监控每个$b_i$对应专家在前一个批次上的负载。如果一个专家在前一个批次上负载较重，我们会减小其偏置。否则，我们会增加它。算法1详细描述了专家级偏置的更新算法。值得注意的是，我们基于历史的均衡状况来更新偏置，因为利用当前序列的负载信息会破坏语言建模的因果约束，导致未来token的信息泄露。通过对偏置的动态调整，我们可以在不引入像辅助损失控制方法那样的噪声梯度的情况下，实现良好的专家负载均衡。

算法1：训练期间调整每个专家的偏置$b_i$

算法1：训练期间调整每个专家的偏置bi
算法1：训练期间调整每个专家的偏置bi

与其他负载均衡方法的比较：为了展示无损均衡的理论优势，我们将其与另外两种主流的负载均衡方法进行了比较，即辅助损失控制方法【索引4，Gshard: Scaling giant models with conditional computation and automatic sharding，Dmitry Lepikhin等人，2020】、【索引3，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，William Fedus等人，2021】和专家选择（Expert Choice, EC）方法【索引10，Mixture-of-experts with expert choice routing，Yan-Quan Zhou等人，2022】。如2.2节所述，辅助损失控制方法面临着负载均衡和模型性能之间的两难困境，可能不存在完美的权衡点。至于EC方法，它会破坏语言建模的因果约束，因为每个token的目标专家是以后续token为条件的（在同一序列或批次中）。这会导致未来token的信息泄露，从而破坏模型的泛化能力。表1总结了不同负载均衡方法的特性。

表1：不同负载均衡方法之间的比较。绿色表示良好属性，红色表示不良属性。

表1：不同负载均衡方法之间的比较。绿色表示良好属性，红色表示不良属性。
表1：不同负载均衡方法之间的比较。绿色表示良好属性，红色表示不良属性。

A4 实验环境

A4 实验结果

主要结果

实验结果表明，无损均衡（Loss-Free Balancing）策略在模型性能和负载均衡两方面均优于传统的辅助损失控制方法。

偏置更新算法的实证研究

为验证主实验中配置的最优性，本文对偏置更新算法的更新率和不同变体进行了实证研究。

讨论

A5 结论

本文提出了一种名为无损均衡（Loss-Free Balancing）的新型MoE负载均衡控制方法，它无需引入辅助损失梯度。该方法有效解决了传统辅助损失控制方法在强制实现负载均衡时会引入额外梯度并可能损害模型性能的问题。在1B和3B参数的MoE模型上进行的实验（分别在100B和200B tokens上训练）表明，与传统的辅助损失训练方法相比，无损均衡策略实现了更好的模型性能和更优的负载均衡。

A6 附录

A 模型架构

架构细节：本文采用DeepSeekMoE【索引1，Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models，Damai Dai等人，2024】架构作为骨干网络。该架构引入了共享专家（shared experts）以减轻路由专家（routed experts）之间的知识冗余。其计算方式如下：

公式6
公式6

其中，$r$表示路由专家，$s$表示共享专家。DeepSeekMoE将除了输入嵌入层之后的全连接层（FFN）外的所有FFN层都替换为MoE层。

超参数：详细的架构超参数列于表5中。

表5：模型架构

表5：模型架构
表5：模型架构

B 训练设置

初始化与序列长度：遵循【索引1，Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models，Damai Dai等人，2024】的工作，所有可学习参数以0.006的标准差进行初始化，最大训练序列长度设置为2048。

1B模型训练参数：
* 学习率调度器：带有预热（warmup）的余弦调度器。
* 学习率：峰值学习率为1e-3，最小学习率为1e-4。
* 预热步数：1000步。
* 训练批次大小：1152。
* 总训练步数：40000步（对应1000亿tokens）。

3B模型训练参数：
* 学习率调度器：多步（multistep）调度器。
* 阶段步数：[45211, 50862, 56514]。
* 对应阶段学习率：[7.8e-4, 2.47e-4, 7.8e-5]。
* 预热步数：2000步。
* 训练批次大小：1728。
* 总训练步数：56514步（对应2000亿tokens）。

验证集：从训练语料库中预留约7100万tokens作为验证集。

C 使用Softmax门的实验

C.1 Sigmoid门与Softmax门基线比较
在1B规模的模型上，比较了不同辅助损失系数$\alpha$下Sigmoid门和Softmax门的基线性能。如图7所示，在相似的负载均衡条件下，Softmax门表现出更高的困惑度，并且其性能对负载不均衡更为敏感。

图7：Sigmoid门基线和Softmax门基线的比较。在相似的负载均衡条件下，Softmax门表现出更高的困惑度，并且对负载不均衡更敏感。
图7：Sigmoid门基线和Softmax门基线的比较。在相似的负载均衡条件下，Softmax门表现出更高的困惑度，并且对负载不均衡更敏感。

C.2 使用Softmax门的无损负载均衡
由于Softmax的归一化特性，调整其专家偏置更具挑战性。在此情况下，我们选择bi = bi + u * ei变体来维持负载均衡，其中$u$设为1e-3。基线选择$\alpha = 0.0003$，这是Softmax门能达到的最低困惑度配置。结果如表6所示，无损均衡实现了略低的困惑度，同时保持了显著更优的负载均衡。图8确认了无损均衡在大部分训练过程中都保持了优越的负载均衡。

表6：对于Softmax门，与辅助损失训练方法相比，无损均衡在达到显著更优的负载均衡的同时，实现了略低的困惑度。

表6：对于Softmax门，与辅助损失训练方法相比，无损均衡在达到显著更优的负载均衡的同时，实现了略低的困惑度。
表6：对于Softmax门，与辅助损失训练方法相比，无损均衡在达到显著更优的负载均衡的同时，实现了略低的困惑度。

图8：对于Softmax门，无损均衡在大部分训练过程中保持了优越的负载均衡。
图8：对于Softmax门，无损均衡在大部分训练过程中保持了优越的负载均衡。

D 专家选择（Expert Choice）中的未来Token泄露

D.1 理论泄露量证明
令$R = K/N$为MoE的稀疏度，其中$K$为每个token激活的平均专家数，$N$为专家总数。对于专家选择中的一个MoE层，最大信息泄露量$I$（单位：比特/token），即路由分配组合可以承载的信息量为：

公式7
公式7

对于一个稀疏度$R = 2/16 = 0.125$且包含9个MoE层的模型，总泄露信息量超过50比特/token。

D.2 实验证据
通过改变用于专家top-k选择的分块大小（从512 tokens到8192 tokens），研究了专家选择的潜在未来token泄露问题。我们在100B tokens上训练了一个2B的MoE模型。结果如图9所示，揭示了两个关键发现：
1. 小分块导致异常损失下降：使用512的小分块大小导致了异常的损失下降，这归因于显著的未来token泄露。更小的分块使得模型在训练中更容易利用分块内未来token的信息。
2. 打乱token可缓解泄露：在分块和选择之前，在批次内打乱token顺序可以缓解观察到的损失下降。这种打乱使得模型更难利用信息泄露，因为未来token不再处于其原始上下文中。这一发现支持了损失下降源于模型访问并利用未来token信息的假设。

图9：不同分块大小和加 shuffling 的专家选择方法比较。分块大小为512的专家选择方法相比于8192或2048的分块大小，表现出显著的损失下降。打乱（Shuffling）token消除了这种损失下降，表明存在未来token泄露。
图9：不同分块大小和加 shuffling 的专家选择方法比较。分块大小为512的专家选择方法相比于8192或2048的分块大小，表现出显著的损失下降。打乱（Shuffling）token消除了这种损失下降，表明存在未来token泄露。

引用文献汇总