翻译:DeepSeek-V3 Explained 3: Auxiliary-Loss-Free Load Balancing

Posted by lili on June 18, 2025

本文翻译DeepSeek-V3 Explained 3: Auxiliary-Loss-Free Load Balancing


这是我们 DeepSeek-V3 系列的第三篇文章,我们将探讨 DeepSeek [1, 2, 3] 模型中与专家混合(MoE)相关的另一项关键架构突破:无辅助损失负载均衡(Auxiliary-Loss-Free Load Balancing)[5]。

目录

背景

Transformer 中的 MoE(专家混合)

MoE 代表专家混合(Mixture-of-Experts),在 Transformer 模型的语境下,这通常是指用多个 FFN(前馈网络)替换每隔几层 Transformer 层中的 FFN,每个 FFN 都充当一个专家(Expert)。因此,当一个输入 token 到来时,一个门控(Gating)操作将选择 Top-K 个专家并将该输入 token 路由到仅选定的 FFN,从而只有这些被选定的专家才会被激活。

如下图所示,左侧标准 Transformer 层中的 FFN 子层被 MoE 层取代。

如需更详细的 MoE 介绍,请查阅DeepSeekMoE,我们在其中用餐厅类比解释了 MoE。

图 1:Transformer 中 MoE 层(红框突出显示)的示意图。图片编辑自 [6]。

负载均衡及其重要性

在我们之前的餐厅类比中,我们用一家提供多种菜肴的餐厅来解释 MoE 的概念,其中每位厨师都充当专家,而主厨则充当门控操作,根据厨师的技能将每道菜分配给特定的厨师。

为了确保这样的系统有效运行,我们需要:

  1. 每位专业厨师都必须掌握自己菜系所需的技能(例如,饺子师傅必须知道如何包饺子),同时他们能够共同处理所有菜肴。
  2. 主厨对所有专业厨师的专长有充分了解,并且能够有效地分配订单。

在 MoE 中,前者对应于专家专业化与知识共享之间的权衡,我们已在《DeepSeekMoE》中详细讨论过。后者则反映了负载均衡的重要性,这也是本文的主要主题。

那么,为什么负载均衡如此重要呢?

原因是当负载不平衡发生时,MoE 无法很好地发挥作用,最常见的问题被称为“路由崩溃(route collapse)”,即只有一小部分专家接收到大部分输入 token,而其他专家则未得到充分利用。

结果,大部分计算由过载的专家进行,导致硬件利用率出现瓶颈,因为专家通常分布在多个 GPU 核心上。

路由崩溃还可能由于梯度冲突而导致训练不稳定。由于过载的专家接收到更多的输入 token,它们也会积累更大的梯度并比欠载的专家学习得更快。结果,过载专家和欠载专家的梯度可能在大小和方向上都发生分歧,使得训练过程更难收敛。

最后,MoE 中的负载不平衡还会导致性能不佳和泛化能力差,因为欠载的专家没有获得足够的训练 token 来学习有意义的知识。

由于负载均衡在 MoE 中如此关键,因此已经提出了各种技术来解决这个问题。在这些现有工作中,最常用的策略是添加用于负载均衡的辅助损失(auxiliary loss for load-balancing)专家选择(Expert Choice)

通过辅助损失进行负载均衡

一种常见的改善负载均衡的策略是在模型训练的原始目标函数之上引入辅助损失函数。

图 2:用于强制实现负载均衡的辅助损失示例。图片编辑自 [5]。

上图展示了一个辅助损失函数的示例,其中:

  • $N$ 是专家数量,$T$ 是 token 数量,$K$ 是每个输入 token 激活的专家数量。
  • $s_{i, t}$ 是来自门控(Gating)的输出,它已通过 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$,从而使分配给每个专家的负载得到调整。

【 译注: 我们可以通过一个例子来理解这个loss。假设只有两个专家,那么最理想的情况是:$f_1=0.5, P_1=0.5, f_2=0.5, P_2=0.5$,此时的loss是$f_1 \times P_1 + f_2 \times P_2 = 0.5$。如果不均衡,$f_1=0.6, P_1=0.6, f_2 = 0.4, P_2=0.4$,则loss为$0.6 \times 0.6 + 0.4 \times 0.4 = 0.52$。因此负载平衡时loss比较小。

有的读者可能会问,如果不均衡,为什么不能是$f_1 = 0.6, P_1=0.4, f_2=0.4, P_2=0.6$,这样的loss不是0.48吗?这比最理想的0.5还要低啊。

但是注意:这种状态时不稳定的,如果$P_1=0.4$,那么说明目前这个专家期望处理的token是比较小的。如果数据正常分布,它的$f_1$也应该接近这个期望。除非这个辅助loss的权重非常大,这样我们的模型宁愿其它的loss很大(专家不匹配)而也愿意减低这个辅助loss。这就是超参数设置的不合理了。 】

然而,通过此类辅助损失来平衡负载是有代价的,因为其梯度可能与语言建模目标的梯度发生冲突,导致模型性能下降,尤其是在 $f_i$ 和 $P_i$ 对于过载专家都变得非常大的极端不平衡情况下。

因此,使用这种方法平衡负载需要仔细权衡辅助损失的权重。为了更清楚地看到这一点,[5] 的作者通过训练具有不同 alpha 值的模型进行了一项实验,结果如下图所示,其中 y 轴显示了以困惑度表示的模型性能,x 轴显示了 MaxVio,这是一个指示负载不平衡程度的指标,MaxVio 越高表示负载不平衡越严重(i 代表第 $i$ 个专家):

图 3:辅助损失控制训练中负载均衡与模型性能之间的两难困境。图片来自 [5]。

如上图所示,当 alpha 过小(alpha = 0)时,MaxVio 保持较高水平,这意味着辅助损失不足以实现负载均衡目标。另一方面,当 alpha 变得过大(alpha = 0.01)时,模型最终会产生更高的困惑度。

总而言之,通过辅助损失控制的负载均衡是一把双刃剑,如果 alpha 未经仔细调整,可能会对模型性能产生负面影响。在实践中,由于资源限制,在 LLM 训练期间调整 alpha 具有挑战性,这进一步使优化过程复杂化。

上图还将所提出的无损失方法置于相同的困惑度-MaxVio 图中,其中它实现了低困惑度和低 MaxVio,显示了无损失方法的有效性。

专家选择 (Expert Choice)

我们在此要提及的另一项先前工作是“专家选择” [7],它通过将路由策略从“token 选择”转向“专家选择”,提出了一种简单而有效的负载均衡方法。

更具体地说,MoE 路由中的门控分数通常是基于亲和力矩阵之上应用 Softmax 来计算的,如图 2 所示。传统的路由方法从 token 维度应用 Softmax,为每个 token 选择专家,因此这些方法被称为“token 选择”。问题在于,在这种机制下,我们无法控制每个专家将接收到多少 token,这最终将导致负载不平衡问题。

另一方面,专家选择通过从专家维度应用 Softmax 来选择路由到每个 token 的 token。因此,每个专家接收到的 token 在设计上是完美平衡的,从而无需辅助损失进行负载均衡。在原始论文 [7] 中,该方法展示了更好的模型性能和更快的训练速度。

然而,专家选择的一个局限性是未来 token 泄露(future token leakage)问题,因为每个专家在看到所有 token 路由分数后才决定处理哪个 token,这违反了因果性,在文本生成和机器翻译等自回归任务中可能是一个严重的问题。

DeepSeek 的无辅助损失负载均衡 (Auxiliary-loss-free Load Balancing)

为了解决负载均衡问题,同时不引入梯度干扰,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 数量与平均值之间的差值(或误差),偏置值由该差值(或误差)的符号乘以固定的更新率(update rate)来确定,该更新率是一个可调的超参数。在后续章节中,我们将看到更多关于此超参数影响的实验。

图 5:DeepSeek 无损失负载均衡算法。图片来自 [5]。

【 译注:这里的$\bar{c_i}$表示专家i历史的平均token数。 】

现在我们可以在下表中总结不同负载均衡方法的优点和局限性:

图 6:不同负载均衡方法的比较。图片来自 [5]。

我们在图 3 中已经看到,所提出的方法在模型性能和负载均衡方面取得了更好的权衡,但仍有许多方面需要检验。在下一节中,我们将更仔细地审视实验结果。

评估

基本上,有三个重要问题需要回答:

  1. 所提出的方法能否在性能和负载均衡之间实现更好的权衡?
  2. 图 5 中更新率 $u$ 的影响是什么?
  3. 我们能否进一步优化偏置更新规则(鉴于它如此简单)?

性能 vs. 负载均衡

为了回答第一个问题,作者对 10 亿和 30 亿参数模型进行了实验,比较了损失控制和无损失负载均衡的困惑度(Perplexity)和 MaxVio,结果如下图所示:

图 7:损失控制与无损失负载均衡的比较。图片来自 [5]。

上述结果与我们在图 3 中看到的情况相似,即所提出的方法实现了更低的困惑度(Perplexity)和更低的 MaxVio。

除了评估最终检查点之外,作者还展示了训练过程中的 MaxVio 曲线,以便我们能更好地了解所提出方法在整个训练过程中的表现,如下图所示:

图 8:训练过程中的 MaxVio 曲线。图片来自 [5]。

如上图所示,在 10 亿和 30 亿参数两种设置下,无损失方法在训练过程中都表现出更好的负载均衡能力,显示了该方法的稳定性。

超参数(更新率)的影响

如图 5 所示,所提出的方法引入了一个名为更新率 $u$ 的新超参数,因此一个自然的问题是,这个超参数将如何影响无损失方法的有效性。更具体地说,我们需要了解无损失方法对 $u$ 的值是敏感还是鲁棒,以及如何选择一个值来最好地利用所提出的方法。

如前所述,在门控分数中添加偏置项在概念上类似于直接对门控分数应用梯度更新,而不依赖于通过损失函数进行反向传播。在这种情况下,更新率 $u$ 的作用类似于梯度更新中的步长。从这个角度来看,我们可以预期到类似的影响:小的更新率可能导致收敛缓慢,而过大的更新率可能导致不稳定和波动。

在原始论文中,作者使用从 1e-4 到 1e-2 的不同更新率进行了实验,结果如下图所示:

图 9:更新率对训练负载均衡的影响。图片来自 [5]。

正如我们所预期的那样,当 $u$ 过小(本例中为 1e-4)时,MaxVio 下降较慢,但过大的 $u$ 也会产生负面影响,由于更大的波动性,导致整个训练过程中 MaxVio 始终较高。

替代偏置更新规则

为了回答第三个问题,作者实验了几种替代策略,并将它们与提出的版本进行比较:

  1. 变体 1: 使用 $e_i$ 的值(而非仅仅其符号)来计算偏置,即将 $b_i = b_i + u \cdot \text{sign}(e_i)$ 改为 $b_i = b_i + u \cdot e_i$。
  2. 变体 2: 使用乘法偏置而非加法偏置。

其中变体 2 可以更正式地描述如下:

他们的实验表明,变体 1 略微改善了负载均衡,但并未提升模型性能:

图 10:变体 1 的性能。图片来自 [5]。

而变体 2 甚至显示出略差的模型性能:

所有以上结果表明,最简单的策略被证明是最好的。

总结

在本文中,我们解释了 DeepSeekMoE 中使用的无辅助损失负载均衡方法,这是 DeepSeek 模型中采用的主要架构创新之一。

更具体地说,我们首先介绍了专家混合(MoE)的基础知识,强调了负载均衡的重要性并回顾了先前的解决方案,包括辅助损失方法和专家选择。然后,我们解释了 DeepSeek 的无损失负载均衡方法及其性能。

我们的主要结论是,DeepSeek 的无损失方法避免了引入梯度干扰,同时保留了因果性,其有效性已通过原始论文中的经验结果得到证实。

参考文献