「DeepSeek-V3 技术解析」：无辅助损失函数的负载均衡
编者按： 在混合专家模型（MoE）的实践中，负载不均衡俨然已成为制约模型性能提升的关键瓶颈之一。传统的均衡策略往往需要引入复杂的辅助损失函数，不仅增加了训练的复杂度，还可能干扰模型的核心学习目标。工程师们在提升模型效率的道路上，一直苦苦追寻着一个优雅而高效的平衡解决方案。
DeepSeek 团队的这项研究，为这一长期困扰业界的技术难题提供了令人耳目一新的解决思路：通过在门控分数中直接添加专家层面的偏置项，在绝大部分不引入额外损失函数的情况下，实现了模型训练过程中的自适应负载均衡。更令人惊叹的是，这一方法不仅保持了模型的因果关系，还显著提升了训练的稳定性和最终性能。
作者 | Shirley Li
编译 | 岳扬
这是 DeepSeek-V3 系列文章的第三篇，我们将探讨 DeepSeek 模型[1, 2, 3]中与混合专家模型（MoE）相关的另一项关键架构突破：无辅助损失函数的负载均衡（Auxiliary-Loss-Free Load Balancing）[5]。
在本文，我们将深入解析 DeepSeek 如何解决 MoE 的隐藏瓶颈------负载均衡，同时还通过消除梯度干扰和严格遵循因果关系约束，提升了训练和推理效率，为后续专家模型的优化方向提供了标杆。
本文目录：
- 技术背景：介绍混合专家模型（MoE）的基本原理，解释负载均衡的重要性，回顾之前的工作，包括辅助损失函数（auxiliary loss）方法和专家选择（Expert Choice）策略。
- DeepSeek 的无辅助损失函数的负载均衡：解析其运作原理
- 性能评估：讨论无辅助损失函数的负载均衡的性能表现
- 总结
- 参考文献
「DeepSeek-V3 技术解析」系列其他文章：
01 技术背景
1.1 MoE (Mixture-of-Experts) in Transformers
MoE（Mixture-of-Experts，混合专家模型）在 Transformer 中的实现方式通常为：每隔若干 Transformer 层，将其中的 FFN（前馈网络）替换为多个 FFN（每个 FFN 充当一个"专家"）。当 input token 进入该层时，通过门控操作（Gating）选择 Top-K 个专家，并将 input token 只路由至这些被选中的 FFN，从而仅激活对应的专家网络。
下图展示了这一过程：左侧标准 Transformer 层中的 FFN 子层被替换为右侧的 MoE 层。
如需更详细的 MoE 技术解析，可参考「DeepSeek-V3 技术解析」：DeepSeekMoE（文中通过餐厅类比直观解释了 MoE 原理）。
图 1. Transformer 中的 MoE 层（红框内）示意图。图片改编自文献 [6]。
1.2 负载均衡及其重要性
在本系列上一篇文章进行的餐厅类比中，我们通过一个能够提供多种菜系菜品的餐厅解释了 MoE 的概念：餐厅的每位厨师都是专家，主厨（Head Chef）的工作类似于门控操作，将每道菜品分配给具备对应技能的特定厨师。
为确保该系统高效运作，需满足以下条件：
- 每位专业厨师必须精通自身菜系所需技能（例如饺子厨师必须会包饺子），同时所有厨师需能共同处理所有菜品。
- 主厨需充分了解所有专业厨师的专长，并能高效分配订单。
在 MoE 中，前者对应 expert specialization 与 knowledge sharing 的权衡（我们已在介绍 DeepSeekMoE 的这篇文章中进行了详细讨论），后者则体现了负载均衡的重要性 ------ 这也是本文的核心主题。
为何负载均衡如此重要？
原因在于，当负载不均衡发生时，MoE 无法有效运作。 最常见的问题是路由崩溃（Route Collapse）：少数专家接收大部分 input token，而其他专家的利用率极低。
因此，大部分计算由超负荷工作的专家承担，而这些专家通常分布在多个 GPU 核心上，因此会导致硬件资源浪费。
由于梯度冲突（gradient conflict），路由奔溃也会导致训练不稳定。超负荷工作的专家接收更多 input token，他们积累的梯度也会更大，学习速度也会比工作负荷不足的专家快得多；因此，两者的梯度在幅值和方向上均可能发生偏离，导致训练难以收敛。
最后，MoE 中的负载不均衡也会导致性能低下和泛化效果不佳，因为工作负荷不足的专家会因为训练 tokens 不足，而难以学习有效知识。
由于负载均衡技术对 MoE 至关重要，因此研究者们针对这一问题提出了多种解决方案。其中，最常用的策略是为负载均衡添加辅助损失函数（Auxiliary Loss）和专家选择（Expert Choice）。
1.3 带辅助损失函数的负载均衡
一种常见的改善负载均衡的策略是在模型训练的目标函数基础上引入辅助损失函数。
图 2. 用于强化负载均衡的辅助损失函数示例。图片编辑自文献[5]。
上图展示了一个辅助损失函数的示例，其中
- N 是专家数量，T 是 token 数量，K 是每个 input token 激活的专家数量。
- s_{i, t} 是门控机制的输出，通过 Softmax 归一化到 [0, 1] 区间，表示第 t 个 token 选择第 i 个专家的概率。向量 u_t 是第 t 个 token 的输入隐藏状态，而 e_i 是第 i 个专家的"质心"，可以看作历史上路由到第 i 个专家的 token 嵌入平均值。因此，s_{i, t} 度量的是当前输入与第 i 位专家历史接收 token 的平均值的接近程度。
- 因此，P_i 可视为整个输入序列选择第 i 个专家的平均概率。
- f_i 表示被路由到第 i 个专家的 token 比例。
需要注意的是，f_i 是不可微分的，因此最小化上述损失函数实际上转化为了最小化 s_{i, t}。同时由于 f_i 依赖于 s_{i, t}，对 s_{i, t} 的调整也会影响 f_i，从而实现对各专家负载分配的调节。
然而，用这种辅助损失函数来均衡负载需要付出一定代价，因为其梯度可能会干扰语言建模目标（language modeling objective）的梯度，导致模型性能下降，在极端不平衡情况下（工作负荷过大的专家的 f_i 和 P_i 都变得极大时）尤其明显。
因此，采用这种方法进行负载均衡需要谨慎设置辅助损失函数的权重。为更清晰地说明这一点，文献[5]的作者进行了一个实验，用不同 alpha 值训练模型，结果如下图所示，其中纵轴表示困惑度指标下的模型性能，横轴表示 MaxVio（衡量负载不平衡程度的指标，MaxVio 值越高表示负载越不平衡，i 表示第 i 个专家）：
图 3. 辅助损失函数控制训练中负载均衡与模型性能的权衡困境。图片引自文献[5]。
如图所示，当 alpha 过小时（alpha=0），MaxVio 保持高位，说明辅助损失函数未能有效实现负载均衡目标。另一方面，当 alpha 过大时（alpha=0.01），模型最终会产生更高的困惑度。
综上，辅助损失函数控制的负载均衡是把双刃剑：若 alpha 未经仔细调校，可能损害模型性能。实际 LLM 训练中，由于资源限制，alpha 的调校过程充满挑战，这进一步增加了优化难度。
上图同时展示了本文提出的无损失函数方法在相同 Perplexity-MaxVio 坐标系下的表现，该方法同时实现了低困惑度和低 MaxVio，证明了无损失函数方法的有效性。
1.4 专家选择（Expert Choice）策略
在此需提及的另一项前期工作是 Expert Choice [7]，它提出了一种简单高效的负载均衡方法，将路由策略从"token choice"切换为"expert choice"。
具体而言，MoE 路由中的门控分数通常通过对 affinity matrix（译者注：二维矩阵，用于量化 input token 与各个专家之间的匹配程度。）应用 Softmax 计算得出，如图 2 所示。传统路由方法从 token 维度应用 Softmax 为每个 token 选择专家，因此这些方法被称为"token choice"。问题在于，该机制下我们无法控制每个专家接收的 token 数量，最终导致负载不均衡问题。
而 Expert Choice 方法则从专家维度应用 Softmax，为每个专家选择被路由的 token。通过这种设计，每个专家接收的 token 数量能够天然达到完美均衡，因此无需依赖辅助损失函数实现负载均衡。在文献[7]中，这种方法同时展现出更优的模型性能和更快的训练速度。
然而，Expert Choice 这种方法也存在局限 ------ 未来 token 泄露问题。由于每个专家需要先查看所有 token 的路由分数才能决定处理哪些 token，这违反了因果关系（causality），在文本生成、机器翻译等自回归任务中可能引发严重问题。
02 DeepSeek 的无辅助损失函数的负载均衡
为在不引入 gradient inference（译者注：此处或为作者笔误？应当为"gradient interference（梯度干扰）"，多个损失函数（或多个优化目标）的梯度方向发生冲突。） 的情况下解决负载均衡问题，DeepSeek 提出了一种名为" Loss-Free Balancing"的新技术，通过直接调整门控分数 s_{i,t} 实现。
如前文所述，当我们最小化图 2 所示的辅助损失函数时，最终会通过调整 s_{i,t} 来实现最小化 P_i。
因此，若能直接调整 s_{i,t}，理论上应能达到与施加辅助损失函数相似的效果。
为此，我们在每个专家的门控分数上添加了专家层面的偏置项，如下图所示。需注意 b_i 并不用于最终门控分数的计算（后文也将说明该偏置项也是不可微分的），而是用于 TopK 选择专家时：
图 4. 在门控分数中引入偏置项 b_i。图片引自文献[5]
该偏置项 b_i 的计算方式非常直观，如下图所示：首先获取分配给各专家的 token 数量平均值及所有专家的理论全局均值，然后计算给各专家分配的 token 数量与理论全局均值的差值，偏置项由该差值（或误差）的符号乘以固定更新率（fixed update rate）决定（该更新率为可调超参数）。后续章节我们将对该超参数的影响进行更多实验分析。
图 5. DeepSeek 无损失函数的负载均衡算法。图片引自文献[5]
现可通过下表总结不同负载均衡方法的优势与局限：
图 6. 不同负载均衡方法对比。图片引自文献[5]
图 3 已展示该方法在模型性能与负载均衡间取得了更好的权衡，但仍有多方面需要验证。下一章节我们将深入分析实验结果。
03 Evaluation
有三个关键问题需要回答：
- DeepSeek 所提出的方法能否在性能和负载均衡之间实现更好的权衡？
- 图 5 中更新率 u 有什么影响？
- 我们能否进一步优化偏置项更新规则（鉴于其如此之简单）？
3.1 性能 vs. 负载均衡
为回答第一个问题，作者在 1B 和 3B 模型上进行了实验，比较 loss-controlled 负载均衡和 loss-free 负载均衡的困惑度（Perplexity）和 MaxVio，结果如下图所示：
图 7. loss-controlled 负载均衡和 loss-free 负载均衡的对比。图片来自文献[5]。
上述结果与我们在图 3 中看到的结果类似：所提方法同时实现了更低的困惑度和更低的MaxVio。
除了评估最终的 checkpoint 外，作者还展示了训练过程中的 MaxVio 曲线，以便更全面地理解该方法在整个训练过程中的表现，如下图所示：
图 8. 训练过程中的 MaxVio 曲线。图片来自文献[5]。
如图中所示，在 1B 和 3B 的模型配置下，loss-free 方法在整个训练过程中都展现出更优的负载均衡能力，体现了该方法的稳定性。
3.2 超参数的影响（更新率）
如图 5 所示，所提方法引入了一个新的超参数 u（称为更新率（update rate）），该超参数如何影响 loss-free 方法的有效性？具体而言，我们需要理解 loss-free 方法对超参数 u 的取值是敏感还是不敏感，以及如何选择一个最优值来最大化该方法的效果。
如前文所述，在门控分数中添加偏置项的概念类似于绕过损失函数的反向传播，直接对门控分数进行梯度更新。在这种情况下，更新率 u 的作用与梯度更新中的步长（step size）类似。由此可推测其影响也相似：过小的更新率可能会导致收敛速度缓慢。过大的更新率可能导致不稳定和引发波动。
在原文中，作者对更新率进行了实验测试（取值从 1e-4 到 1e-2），结果如下图所示：
图 9. 更新率（update rate）对训练负载均衡的影响。图片来自文献[5]。
与预期一致，当 u 过小时（如 1e-4），MaxVio 下降速度较慢；而过大的 u（如 1e-2）则因波动性增强，导致训练过程中 MaxVio 持续偏高。
3.3 其他偏置项更新规则
为回答第三个问题，研究者尝试了多种备选策略，并将它们与 DeepSeek 提出的版本进行对比：
- 策略变体 1：使用 e_i 的数值（而不仅仅是符号）计算偏置项，即从 b_i = b_i +u∗sign(e_i) 改为b_i = b_i +u∗e_i。
- 策略变体 2：使用乘法偏置项而非加法偏置项。
其中策略变体 2 可以更正式地描述如下：
实验表明，策略变体 1 能带来略优的负载均衡效果，但未提升模型性能：
图 10. 策略变体 1 的性能表现。图片来自文献[5]。
而策略变体 2 甚至显示出略差的模型性能：
图 11. 策略变体 2 的性能表现。图片来自文献[5]。
以上所有结果均表明，最简单的策略反而是最佳选择。
04 Summary
在本文中，我们解释了 DeepSeek 模型的核心架构创新之一 ------ DeepSeekMoE 中使用的无辅助损失函数的负载均衡方法。
本文首先介绍了混合专家模型（MoE）的基本原理，强调了负载均衡的重要性，并回顾了先前的解决方案（包括 auxiliary loss 方法和 Expert Choice 机制）。接着，本文阐释了 DeepSeek 的无损失函数的负载均衡方法及其性能表现。
DeepSeek 的无损失函数方法在保持因果关系的同时避免了引入梯度干扰，其有效性已通过原论文的实证结果得到验证。
感谢您花时间阅读本文！
参考文献
[1] DeepSeek（https://www.deepseek.com/）
[2] DeepSeek-V3 Technical Report（https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf）
[3] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model（https://arxiv.org/abs/2405.04434）
[4] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models（https://arxiv.org/abs/2401.06066）
[5] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts（https://arxiv.org/abs/2408.15664）
[6] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding（https://arxiv.org/abs/2006.16668）
[7] Mixture-of-Experts with Expert Choice Routing（https://arxiv.org/abs/2202.09368）
Thanks for reading!
Hope you have enjoyed and learned new things from this blog!
About the author
Shirley Li
I am a Machine Learning Engineer working on building multi-modality models to solve real-world problems.
END
本期互动内容 🍻
❓在实际工程中，您认为负载均衡对模型性能的影响有多大？除了本文提到的技术路径，您还了解哪些有效的负载均衡方案？
原文链接：
https://ai.gopubby.com/deepseek-v3-explained-3-auxiliary-loss-free-load-balancing-4beeb734ab1f
