文章标题：混合专家模型（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 实验环境

  • 模型架构：

    • 主干网络采用DeepSeekMoE【索引1，Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models，Damai Dai等人，2024】架构，该架构通过引入共享专家来减少路由专家间的知识冗余，性能优于GShard等传统MoE架构。
    • 门控函数（Gating Function）在主实验中选用Sigmoid，因其基线性能优于Softmax。
    • 实验涵盖两种模型规模：10亿（1B）参数和30亿（3B）参数。偏置更新率等超参数仅在1B模型上进行调优，3B模型直接沿用最佳配置。
    • 详细架构参数见附录A中的表5。

  • 数据集与训练设置：

    • 数据集：使用DeepSeek-AI创建的多语言训练语料库，来源广泛，包括网页文本、数学资料、代码脚本和已发表文献。
    • 分词器：使用HuggingFace Tokenizer训练了一个词汇量为32K的字节对编码（BPE）分词器【索引6，Neural machine translation of rare words with subword units，Rico Sennrich等人，2015】。
    • 训练规模：1B模型在1000亿（100B）tokens上训练，3B模型在2000亿（200B）tokens上训练，以确保充分收敛。
    • 学习率调度器：1B模型使用余弦学习率调度器【索引5，Sgdr: Stochastic gradient descent with warm restarts，Ilya Loshchilov和Frank Hutter，2016】，3B模型使用多步学习率调度器【索引1，Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models，Damai Dai等人，2024】。
    • 验证集：从训练语料库中预留约7000万（70M）tokens作为验证集，用于评估模型性能和负载均衡情况。

  • 软件与硬件配置：

    • 论文未提供具体的硬件配置信息（如GPU型号、数量、网络互联等）。
    • 软件实现基于主流深度学习框架，具体细节未在文中详述。

  • 基线与评估指标：

    • 基线方法：传统的辅助损失控制方法，其中辅助损失系数$\alpha$设置为0.001，这是一个在模型性能和负载均衡之间取得合理权衡的值（如图2所示）。
    • 评估指标：
      • 模型性能：使用困惑度（Perplexity）衡量。
      • 负载均衡：引入最大违规度（Maximal Violation, MaxVio）来量化MoE层的负载均衡程度，其计算公式为：
        公式4
        公式4

        其中，$Load_i$是分配给第$i$个专家的token数量，$Load_{\text{ideal}}$是理想均衡状态下的专家负载。
      • MaxVio的两种变体：
        • MaxVio_global：在整个验证集上计算，反映了专家的长期利用均衡度和批次大小趋于极限时的效率上限。
        • MaxVio_batch：在每个训练批次上计算，与训练效率更直接相关。

A4 实验结果

主要结果

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

  • 模型性能与全局负载均衡：如表2所示，对于1B和3B参数规模的模型，无损均衡方法都取得了更低的验证集困惑度（Perplexity）和显著更优的全局负载均衡（MaxVio_global）。这证明该方法成功打破了性能与均衡之间的两难困境。

    表2：无损均衡在1B和3B模型上均实现了更低的困惑度和更好的负载均衡。这些指标使用验证集计算得出（详见附录B）。

    表2：无损均衡在1B和3B模型上均实现了更低的困惑度和更好的负载均衡。
    表2：无损均衡在1B和3B模型上均实现了更低的困惑度和更好的负载均衡。

  • 训练过程中的负载均衡：图3展示了训练过程中批次级负载均衡（MaxVio_batch）的变化曲线。结果显示，无损均衡方法在绝大部分训练时间内都保持了比辅助损失方法更优的负载均衡状态，证明了其持续有效的控制能力。

    图3：无损均衡在大部分训练时间内保持了更好的负载均衡。为便于观察，此处的MaxVio_batch是每100个相邻步骤的平均值。
    图3：无损均衡在大部分训练时间内保持了更好的负载均衡。为便于观察，此处的MaxVio_batch是每100个相邻步骤的平均值。

偏置更新算法的实证研究

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

  • 更新率（Update rate）：更新率$u$控制专家偏置收敛到“合适值”的速度。如图4所示，过低的更新率（u=0.0001）导致收敛缓慢，在训练早期负载均衡较差；而过高的更新率（u=0.01）则可能在训练后期引起偏置的不必要波动，恶化负载均衡。实验证明，u=0.001是一个合适的选择，它在训练均衡性和验证困惑度上都表现良好。

    图4：更新率对训练负载均衡的影响。低更新率在训练早期表现出较差的负载均衡，而高更新率则在后期恶化了负载均衡。Validation PPL表示验证集困惑度。
    图4：更新率对训练负载均衡的影响。低更新率在训练早期表现出较差的负载均衡，而高更新率则在后期恶化了负载均衡。Validation PPL表示验证集困惑度。

  • 更新规则（Update rule）：研究了将更新规则从bi = bi + u * sign(ei)更改为bi = bi + u * ei的变体。后者使得违规误差大的专家偏置变化更快。如表3所示，虽然这种变体略微改善了负载均衡，但并未带来模型性能的提升。因此，保留了原始的sign版本。

    表3：变体bi = bi + u * ei略微改善了负载均衡，但模型性能未见提升。

    表3：变体bi = bi + u * ei略微改善了负载均衡，但模型性能未见提升。
    表3：变体bi = bi + u * ei略微改善了负载均衡，但模型性能未见提升。

  • 乘性偏置（Multiplicative bias）：除了加性偏置，还探索了使用乘性偏置的变体，即s'_{i,t} = b_i * s_{i,t}。如表4所示，与加性偏置相比，使用乘性偏置导致模型性能略差，且在负载均衡方面没有显著改善。基于这些发现，结论是加性偏置是更合适的选择。

    公式5
    公式5

    表4：与加性偏置相比，乘性偏置在负载均衡上表现相似，但性能略差。

    表4：与加性偏置相比，乘性偏置在负载均衡上表现相似，但性能略差。
    表4：与加性偏置相比，乘性偏置在负载均衡上表现相似，但性能略差。

讨论

  • 与专家并行的兼容性：在超大规模MoE模型训练中，通常采用专家并行技术将专家分布在不同设备上。在这种情况下，单个计算步骤内的负载均衡对效率至关重要。如图5所示，随着计算批次大小的增加，无损均衡策略的负载均衡度持续改善，而辅助损失方法的均衡度在批次较大时趋于稳定。由于专家并行会显著增大计算批次，因此无损均衡策略天然兼容大规模MoE训练，并且其在负载均衡上的优势会随着专家并行规模的扩大而进一步增强。

    图5：随着计算批次大小的增加，无损均衡比辅助损失训练实现了更好的均衡性，显示了其在使用中等大小计算批次时的优越性。
    图5：随着计算批次大小的增加，无损均衡比辅助损失训练实现了更好的均衡性，显示了其在使用中等大小计算批次时的优越性。

  • 负载均衡与未来Token泄露问题：对于自回归语言模型，负载均衡方法必须遵守因果约束以避免未来token信息泄露。传统的辅助损失方法和本文的无损均衡都遵守此约束，但专家选择（Expert Choice, EC）方法【索引10，Mixture-of-experts with expert choice routing，Yan-Quan Zhou等人，2022】违反了这一约束。EC通过为每个专家分配完全相同数量的token来确保完美的负载均衡，但这种做法内在地导致了严重的未来token信息泄露问题。如图6所示，未来token可以影响前面token的专家分配，从而传递信息。理论上，这种信息泄露量是巨大的，足以让一个token推断出其后继者的身份。实验通过减小分块大小和打乱token顺序证实了EC方法存在严重的泄露问题，这会破坏模型的泛化能力，使其评估不可靠。因此，与EC相比，使用无损均衡策略扩展MoE模型更为安全。

    图6：EC中未来token泄露的一个例子。未来token可以影响前面token的专家分配。这种分配可以帮助前面的token推断其后继者的身份。
    图6：EC中未来token泄露的一个例子。未来token可以影响前面token的专家分配。这种分配可以帮助前面的token推断其后继者的身份。

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泄露。

引用文献汇总

  • [1] Dai, Damai, et al. "Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models." ArXiv, abs/2401.06066, 2024.
  • [2] DeepSeek-AI, et al. "Deepseek-coder-v2: Breaking the barrier of closed-source models in code intelligence." ArXiv, abs/2406.11931, 2024.
  • [3] Fedus, William, Barret Zoph, and Noam M. Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." J. Mach. Learn. Res., 23:120:1–120:39, 2021.
  • [4] Lepikhin, Dmitry, et al. "Gshard: Scaling giant models with conditional computation and automatic sharding." ArXiv, abs/2006.16668, 2020.
  • [5] Loshchilov, Ilya, and Frank Hutter. "Sgdr: Stochastic gradient descent with warm restarts." arXiv: Learning, 2016.
  • [6] Sennrich, Rico, Barry Haddow, and Alexandra Birch. "Neural machine translation of rare words with subword units." ArXiv, abs/1508.07909, 2015.
  • [7] Shao, Zhihong, et al. "Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model." ArXiv, abs/2405.04434, 2024.
  • [8] Shazeer, Noam M., et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." ArXiv, abs/1701.06538, 2017.
  • [9] Vaswani, Ashish, et al. "Attention is all you need." In Neural Information Processing Systems, 2017.
  • [10] Zhou, Yan-Quan, et al. "Mixture-of-experts with expert choice routing." ArXiv, abs/2202.09368, 2022.