[DeepSeek-V3-技术报告阅读] Complementary Sequence-Wise Auxiliary Loss

前情提要

DeepSeek V3是一个MOE结构的模型，本文仅解读文中介绍的序列内负载均衡损失：Complementary Sequence-Wise Auxiliary Loss，至于序列间负载均衡Auxiliary-Loss-Free Load Balancing，由于简单易懂本文不做解读，但是本文的内容需要其作为前置知识，因此，在阅读本文前需要先阅读过原文的Auxiliary-Loss-Free Load Balancing部分。

Complementary Sequence-Wise Auxiliary Loss

DeepSeek V3技术报告中提到了一个Complementary Sequence-Wise Auxiliary Loss，这个Loss的目的是在序列内部均衡专家路由过程，避免在一个序列的不同token中，部分专家被疯狂路由，另一些专家不被路由。这里的序列，指训练时的一个长文本，一般为定长的文本，可能包含多个句子，也可能只包含半句话。Complementary Sequence-Wise Auxiliary Loss的定义如下：

\[\begin{align} \mathcal{L}_{\mathrm{Bal}} & = \alpha \sum_{i=1}^{N_r}{f_i P_i}, \label{eq:1}\\ f_i = \frac{N_r}{K_r T} \sum_{t=1}^{T} \mathbb{1} & \left( s_{i,t} \in \operatorname{Topk} ( \{ s_{j, t} | 1 \leq j \leq N_r \}, K_{r} ) \right), \label{eq:2}\\ s^{\prime}_{i,t} & = \frac{s_{i,t}}{\sum_{j=1}^{N_r} s_{j,t}}, \label{eq:3}\\ P_i & = \frac{1}{T} \sum_{t=1}^{T}{s^{\prime}_{i,t}}, \label{eq:4} \end{align} \]

其中\(N_r\)是路由专家的数量，\(K_r\)是激活专家的数量，\(s_{i, t}\)是第\(i\)个专家和第\(t\)个token之间的亲和度（affinity），\(T\)是序列中token的数量。

这个Loss的作用是控制专家的负载均衡，为了方便研究，首先定义专家负载的衡量方式：

1. 绝对负载：路由到该专家的token的数量 / 全部的token的数量
   绝对负载的取值范围是[0, 1]
2. 相对负载：实际绝对负载 / 理想绝对负载
   相对负载的取值范围是[0, +\(\infty\)]

\(f_i\)的含义

首先看公式\eqref{eq:2}，\(f_i\)的含义是：第\(i\)个专家的“相对负载”。实际上\(f_i\)就等于：

\[f_i = \frac{\mathrm{专家}i\mathrm{实际分配到的token的比例}}{理想情况下\mathrm{专家}i\mathrm{分配到的token的比例}} \]

为了理解原始公式的含义，将原始公式写成如下形式：

\[f_i = \frac{\frac{1}{T} \cdot \sum_{t=1}^{T} \mathbb{1} \left( s_{i,t} \in \operatorname{Topk} ( \{ s_{j, t} | 1 \leq j \leq N_r \}, K_{r} ) \right)}{\frac{K_r}{N_r}} \]

分子中的\(\sum_{t=1}^{T} \mathbb{1} \left( s_{i,t} \in \operatorname{Topk} ( \{ s_{j, t} | 1 \leq j \leq N_r \}, K_{r} ) \right)\)是第\(i\)个专家分配到的token的数量（负载），乘以\(\frac{1}{T}\)就是该专家分配到的token的比例。最前面的\(\frac{N_r}{K_r}\)是：理想情况下专家\(i\)的负载。为什么是这个呢？

我们先来考虑一个数学问题：
问：一共有\(T\)个token，\(N_r\)个专家，每个token都需要从中随机选择\(K_r\)个专家，请问对于token \(t\)，它选择专家\(i\)的概率是多少？
答：\(K_r / N_r\)。

你问我为什么？那我问你，你从10个美女中随机选择3个做女朋友，请问1号美女被选中的概率是多少？当然是3/10。还不理解？那我用组合数给你证明一下：

\[\begin{equation*} P = \frac{\mathrm{C}_{N_r - 1}^{K_r - 1}}{\mathrm{C}_{N_r}^{K_r}} = \frac{K_r}{N_r} \end{equation*} \]

在完全随机的情况下，这个概率其实有这么几种含义：

1. 表示一个token被路由到专家i的概率
2. 表示一个token对一个专家i产生的“负载”
3. 表示一个专家分配到的token的比例

这个概率的前两层含义都好理解，第三层含义是，在纯随机的情况下，一个专家分配到的token的比例是\(K_r / N_r\)。为什么呢？很好理解，每个token分配给该专家的概率都是\(K_r / N_r\)，该专家分配到的token数量的期望\(E_i = T \cdot K_r / N_r\)，在所有token中占据的比例就是\(\frac{E_i}{T} = \frac{K_r}{N_r}\)。在这种理想情况下，\(f_i \equiv 1, \forall i \in \{1,\cdots,N_r\}\)，不论\(K_r\)的取值是多少，负载都是完全均衡的。

此时，\(f_i\)的意义就呼之欲出了：

\[f_i = \frac{\mathrm{专家}i\mathrm{实际分配到的token}}{理想情况下\mathrm{专家}i\mathrm{分配到的token}} \]

也就是：第\(i\)个专家实际承受的负载与最佳负载的倍数比，承受最佳负载的1倍是最好的。\(f_i\)越接近1，在该专家上的负载就越均衡，\(f_i > 1\)，则说明该专家负载过重，\(f_i < 1\)，则该专家负载过轻。

如果公式\eqref{eq:2}不除以\(K_r / N_r\)，当要激活的专家数量不同时，专家分配到的token比例就不同，反映在loss上就是在同样是均衡的负载下，\(f_i\)不一致。因此，除以\(K_r / N_r\)可以让\(f_i\)在反应负载均衡度时，不受激活专家的数量影响。

\(P_i\)的含义

公式\eqref{eq:3}，\(s^{\prime}_{i,t}\)是归一化后的token-专家 亲和度，其反应的是token \(t\)在被路由到专家\(i\)的概率。概率越大，token \(t\)就越可能被路由到专家\(i\)，导致专家\(i\)的“负载”就越可能增加一点。

然后再看公式\eqref{eq:4}，\(P_i\)将\(s^{\prime}_{i,t}\)在所有token上做了平均，表示：每个token被路由到专家\(i\)的平均概率。

首先我们明确一点，公式\eqref{eq:4}和\eqref{eq:2}没有直接关系，假设从5个专家中选3个，那么\(s^{\prime}_{\cdot, t}=\{0.1,0.1,0.2,0.3,0.3\}\)和\(\{0.001,0.001,0.002,0.002,0.994\}\)的结果是一样的。

在最理想的情况下，先不管\(s^{\prime}_{i,t}\)的取值，\(P_i\)的取值应该是\(1/N_r\)。

\(\sum_{i=1}^{N_r}{f_i P_i}\)的含义

\(f_i P_i\)是想做什么？我们反向推理，既然这个东西是Loss，那么理想情况下它应该取得最小值。而负载均衡的理想情况就是每个专家的负载都完全相同，也就是\(f_i = 1, P_i = 1/N_r\)，此时\(\sum_{i=1}^{N_r}{f_i P_i}=1\)。也就是\(\min(\mathcal{L}_{\mathrm{Bal}}) = \alpha\)

那怎么证明呢？显然，要证明\(\min(\mathcal{L}_{\mathrm{Bal}}) = \alpha\)，只需要证明\(\min(\sum_{i=1}^{N_r}{f_i P_i})=1\)。关于这个证明，在Switch Transformer的论文中只说了一句话：

The auxiliary loss of Equation 4 encourages uniform routing since it is minimized under a uniform distribution.
公式4的辅助损失鼓励均匀路由，因为它在均匀分布下取得最小值

真是说了等于没说，跟放屁一样。
证明如下：

先将证明问题转化为一个完整的数学问题，易知

\[\begin{equation} \begin{aligned} \sum_{i=1}^{N_r} P_i &= 1 \\ \sum_{i=1}^{N_r} f_i &= N_r \end{aligned}\label{eq:5} \end{equation} \]

你问我上面这块怎么证明？自己举个例子看看就懂了

那么，需要证明的问题就变成了以下数学问题：

已知：

\[\begin{equation} \begin{aligned} \sum_{i=1}^{N_r} P_i &= 1, 其中 0 \le P_i \le 1\\ \sum_{i=1}^{N_r} f_i &= N_r, 其中 0 \le f_i \le N_r \end{aligned}\label{eq:6} \end{equation} \]

问，\(\sum_{i=1}^{N_r}{f_i P_i}\)的最小值是多少？在什么时候取得最小值？

解：

使用拉格朗日乘数法求解上述问题，构造拉格朗日函数：

\[L(f_i; P_i; \lambda_1, \lambda_2) = \sum_{i=1}^{N_r}{f_i P_i} + \lambda_1(\sum_{i=1}^{N_r}f_i - N_r) + \lambda_2(\sum_{i=1}^{N_r}P_i - 1) \]

令其对各变量的偏导为0，可得方程组：

\[\begin{equation*} \left\{ \begin{aligned} \frac{\partial L}{\partial f_i} = P_i + \lambda_1 = 0 \\ \frac{\partial L}{\partial P_i} = f_i + \lambda_2 = 0 \\ f_1 + \cdots + f_{N_r} - N_r = 0 \\ P_1 + \cdots + P_{N_r} - 1 = 0 \end{aligned}\right. \end{equation*} \]

求解后可得：

\[\begin{equation*} \left\{ \begin{aligned} f_1 = f_2 = \cdots = f_{N_r} &= 1 \\ P_1 = P_2 = \cdots = P_{N_r} &= \frac{1}{N_r} \\ \lambda_1 &= -\frac{1}{N_r} \\ \lambda_2 &= -1 \end{aligned}\right. \end{equation*} \]

代入后可得，\(\sum_{i=1}^{N_r}{f_i P_i} = 1\)

求解完毕。