MoE环游记：8、强制序列级均衡

到目前为止，“MoE环游记”系列已经写了7篇文章，其中5篇都是围绕着MoE的路由和负载均衡展开的。从路由的形式来看，它们可以分为静态计算和动态计算两类；从实现负载均衡的方法上看，它们又可以分为Aux-Loss和Loss-Free两类，其中Loss-Free又可以DeepSeek所提的SignSGD方案以及笔者所提的QB方案。

还有一个细节是，我们此前讨论的Loss-Free方案，都是在实现全局批次意义下的均衡，而很多时候我们会加上一个序列级的Aux Loss，来避免序列级的极端不均衡。如果我们相信每个Aux-Loss方案都有对应的Loss-Free版本，那么序列级的Loss-Free均衡，又应该怎么实现呢？这便是文本要探讨的主题。

前文回顾[#](https://spaces.ac.cn/archives/11760#%E5%89%8D%E6%96%87%E5%9B%9E%E9%A1%BE)

首先让我们回顾一下之前关于负载均衡的探讨。在文章《MoE环游记：2、不患寡而患不均》中，我们首次讨论MoE的负载均衡问题，所提的解决方案为Aux Loss，即添加额外的损失函数来促进均衡，Aux Loss的问题是有额外的权重系数要调，优点是可以随意调节均衡的颗粒度，可以实现序列、分组或者全局层次的均衡（当然越全局通信也越大）。

从《MoE环游记：3、换个思路来分配》开始，我们讨论Loss-Free形式的负载均衡策略，它通过引入一个额外的偏置项来调节均衡程度，并手工设计SignSGD形式的规则来更新偏置项，这种模式比Aux Loss更灵活且对Infra友好。在《MoE环游记：4、难处应当多投入》中，我们将它推广到了动态激活数量的MoE。

而在《MoE环游记：6、最优分配促均衡》中，我们从最优分配的视角看待负载均衡问题，通过交替求解对偶目标，我们得到了一种新的负载均衡策略Quantile Balancing（QB）。单从结果上来看，它可以视为Loss-Free方案偏置项的一种新的更准确的求解方案。接着，在《MoE环游记：7、动态激活极简解》中我们同样将QB推广到了动态MoE。

这里有一个细节，Loss-Free引入的偏置项是全局共享的，所以目前所有的Loss-Free方案，无一例外只能实现全局负载均衡。刚才说了，如果我们需要的话，Aux Loss也可以用于促进序列级均衡，那么Loss-Free是否也有对应的序列级均衡的方案呢？

局部中心[#](https://spaces.ac.cn/archives/11760#%E5%B1%80%E9%83%A8%E4%B8%AD%E5%BF%83)

无独有偶，@ChangJonathanC 在文章《Causal Routing Bias for Aux-Loss-Free MoE Training》中也进行过相关探索，作者提出了两种思路，我们先来简单学习一下，也便于跟我们后面自己的思路进行对比。设$\mathbf{\mathit{s}} \in \mathbb{R}^{l \times n}$表示一个长度为$l$的序列的Router输出，$\mathbf{\mathit{s}}_{i} \in \mathbb{R}^{n}$则表示第$i$个token的Router输出。文章所提的第一个思路是，将$\mathbf{\mathit{s}}$减去它的滑动平均（EMA）：

$$ (\text{1}) \hat{\mathbf{\mathit{s}}}_{i} = \mathbf{\mathit{s}}_{i} - \bar{\mathbf{\mathit{s}}}_{i} , \bar{\mathbf{\mathit{s}}}_{i} = \gamma \bar{\mathbf{\mathit{s}}}_{i - 1} + \left(\right. 1 - \gamma \left.\right) \mathbf{\mathit{s}}_{i} $$

然后用$\hat{\mathbf{\mathit{s}}}$去决定去激活哪些Expert（比如选Top-$k$，也可以结合QB等手段使用）。这个做法背后的思想很朴素：直观来想，如果用$\mathbf{\mathit{s}}$来做决策时Expert $j$被激活的次数多，那么意味着它的打分$\mathbf{\mathit{s}}_{: , j}$平均来说偏大，因此为了均衡，应该对它做相应的惩罚，惩罚的量正是由EMA估算的局部平均值。 这样一来，在新的$\hat{\mathbf{\mathit{s}}}$中，每个Expert的分数在局部都近似具有零均值，不至于让某个Expert更为“突出”，从而实现序列级均衡。然而，这只是一个经验做法，并没有理论确保它一定会更均衡，笔者简单做过测试，像第一层MoE之类的极端不平衡场景，它并没有带来实质帮助。也许作者也意识到了这个局限性，所以他提出了第二种方案。

测试训练[#](https://spaces.ac.cn/archives/11760#%E6%B5%8B%E8%AF%95%E8%AE%AD%E7%BB%83)

第二种方案的思路其实也挺直观。在《MoE环游记：3、换个思路来分配》中，我们不是引入了一个偏置项来控制全局的均衡性吗？这个偏置项用SignSGD来更新。既然如此，我们可以模仿TTT的“测试时训练”思想，沿着序列维度逐Token地激活Expert，并根据已经激活的分布情况，更新这个偏置项。

以$\mathbf{\mathit{s}} \in \mathbb{R}^{l \times n}$和Top-$k$版MoE为例，原本的Loss-Free方案，是引入全局的偏置向量$\mathbf{\mathit{\beta}} \in \mathbb{R}^{n}$，然后将激活机制改为逐$\mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}$选Top-$k$。而这里的想法是为每个$\mathbf{\mathit{s}}_{i}$都分配一个$\mathbf{\mathit{\beta}}_{i} \in \mathbb{R}^{n}$，激活机制改为$\mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}_{i}$选Top-$k$。$\mathbf{\mathit{\beta}}_{i}$的更新规则可以是符号梯度上升

$$ (\text{2}) \mathbf{\mathit{\beta}}_{i} = \mathbf{\mathit{\beta}}_{i - 1} + \eta \text{sign} ⁡ \left(\right. \mathbf{\mathit{f}}_{\leq i - 1} - 1 / n \left.\right) $$

这里$\mathbf{\mathit{f}}_{\leq i - 1}$是截止到第$i - 1$个Token的已激活的Expert分布。作者最后的版本做了一些调整：把符号函数去掉，并将$\mathbf{\mathit{f}}_{\leq i - 1}$改成了$\mathbf{\mathit{f}}_{i - 1}$（第$i - 1$个Token的激活分布），以更好地实现局部均衡

$$ (\text{3}) \mathbf{\mathit{\beta}}_{i} = \mathbf{\mathit{\beta}}_{i - 1} + \eta \left(\right. \mathbf{\mathit{f}}_{i - 1} - 1 / n \left.\right) $$

这个思路能够保持序列的因果性，同时也能保证实现序列级几乎绝对的均衡，但由于每步都要执行选Top-$k$并且根据结果$\mathbf{\mathit{f}}$更新偏置项，所以这本质上引入了一个非线性RNN。非线性RNN无法并行，因此可以预见序列长了会是一个瓶颈。当然我们也可以考虑按Chunk来更新（Mini-batch TTT）以提高效率，但这就略微显得不优雅了。

最优之解[#](https://spaces.ac.cn/archives/11760#%E6%9C%80%E4%BC%98%E4%B9%8B%E8%A7%A3)

接下来，本文所提的思路称为“Moving Quantile Balancing（MQB）”，顾名思义，这是在“（Quantile Balancing）QB”的基础上演化出来的。为此，我们先简单回顾一下QB。

按照笔者撰写的顺序，是先有《MoE环游记：6、最优分配促均衡》再有《MoE环游记：7、动态激活极简解》，但实际上后者概念上和方法上都简单得多，所以我们从后者讲起。还是给定$\mathbf{\mathit{s}} \in \mathbb{R}^{l \times n}$，我们引入偏置项$\mathbf{\mathit{\beta}} \in \mathbb{R}^{n}$，静态版的MoE是每个Token都激活$\mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}$最大的$k$个Expert，动态激活是指只要$\mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}} > 0$的Expert都激活，数量不定。

为了同时保证负载均衡和控制平均激活数量为$k$，需要选择适当的$\mathbf{\mathit{\beta}}$。非常奇妙的是，在这个场景下，$\mathbf{\mathit{\beta}}$的最优解是可以精确表示出来的：

$$ (\text{4}) \mathbf{\mathit{\beta}} = \text{desc}_\text{sort} ⁡ \left(\right. \mathbf{\mathit{s}} , \text{axis}=\text{0} \left.\right)_{\left[\right. l k / n : l k / n + 1 \left]\right.} = \text{quantile} ⁡ \left(\right. \mathbf{\mathit{s}} , 1 - k / n , \text{axis}=\text{0} \left.\right) $$

也就是说，将$\mathbf{\mathit{s}}$沿着序列维度（当前是$\text{axis}=\text{0}$）从大到小排列，取第$l k / n$大的元素，便是$\mathbf{\mathit{\beta}}$的最优解，这也对应于“$1 - k / n$分位数”。 实际上，这只不过是Expert Choice的另一种表示方式，我们知道，Expert Choice的问题是会破坏因果律（非Causal），这从它沿着序列维度去求$\mathbf{\mathit{\beta}}$也可以看出。如果我们是Encoder场景，那么这种非Casual的方案自然不是问题，但对于单向的自回归模型却不可取。

滑动分位[#](https://spaces.ac.cn/archives/11760#%E6%BB%91%E5%8A%A8%E5%88%86%E4%BD%8D)

前面 @ChangJonathanC 的“测试时训练”方案给了笔者一些启发。我们可以沿着序列维度，逐Token地基于QB计算偏置项，由于QB能够直接基于$\mathbf{\mathit{s}}$把偏置项的最优解表示出来，而不用知道Expert激活情况，所以这带来了一些并行的可能性。

笔者最开始的想法是$\mathbf{\mathit{\beta}}_{i} = \text{quantile} ⁡ \left(\right. \mathbf{\mathit{s}}_{\left[\right. : i \left]\right.} , 1 - k / n , \text{axis}=\text{0} \left.\right)$，即对所有不超过当前位置的Token分数算偏置项，那么每个位置都有自己的偏置项，并且QB的最优性保证了直至当前位置都是均衡的。由于它需要沿着序列维度递进地计算分位数，我们不妨称这种运算为“Cumulative Quantile”，相应的方案为“Cumulative Quantile Balancing（CQB）”。

CQB原理上是可行的，但由于Quantile无法增量更新，所以CQB每一步都要读完整的数据做Quantile计算，计算复杂度是线性增加的，总复杂度则是平方级。为了缓解这个问题，同时为了更好地体现出局部均衡，笔者的想法是设定一个窗口$w$，每个Token只在不超过$w$的窗口内求最优偏置项，即

$$ (\text{5}) \mathbf{\mathit{\beta}}_{i} = \text{quantile} ⁡ \left(\right. \mathbf{\mathit{s}}_{\left[\right. i - w : i \left]\right.} , 1 - k / n , \text{axis}=\text{0} \left.\right) $$

这样运算就从“Cumulative Quantile”变成了“Moving Quantile”，相应的方案便是“Moving Quantile Balancing（MQB）”的雏型。它跟SWA（Sliding Window Attention）很相似，复杂度是线性的，并且理论上也可以并行，之所以还是“雏型”，依然是因为Quantile无法增量更新，导致单步复杂度依然相对较高，仍需设法改进。

分桶估计[#](https://spaces.ac.cn/archives/11760#%E5%88%86%E6%A1%B6%E4%BC%B0%E8%AE%A1)

上一节所说的各种障碍，本质上都受限于“Quantile是非线性的，无法增量更新”，有没有一种可能，我们设法将它线性化呢？哪怕只是近似都好？很幸运，这里还真能做到！


直方图近似估计分位数示意图

一个关键事实是：只要知道变量的分布，就可以通过累积概率求分位数，而分布是可以增量更新的！所以，这里的关键是将分位数的估计转化成分布的估计，对此我们使用分桶计数，也称为“直方图近似”。首先，假设Router分数$\mathbf{\mathit{s}}$都在$\left[\right. 0 , 1 \left]\right.$内，这可以通过加Sigmoid等激活函数来保证；接着，将$\left[\right. 0 , 1 \left]\right.$等分为$b$个桶，将$\mathbf{\mathit{s}}$的每个值离散化为一个整数，继而转化为对应的one hot向量。

这样，我们就得到了一个$l \times n \times b$的0/1数组，如果将它沿着序列维度$l$求平均，就可以得到每个Expert的分数分布（以$b$维向量表示），继而可以求每个Expert的分位数。不过，为了保证因果性，我们不能直接对整个序列平均，只能做类似“Cumulative Average”的操作，为了避免“Moving Quantile”的严格边界带来的麻烦，同时也使得整个过程更为“平滑”，这里我们选择做EMA。

也就是说，将$\mathbf{\mathit{s}}$分桶one hot化后，我们沿着序列维度做EMA，每个Token都能得到一个局部的分布，根据这个分布就可以求对应的$1 - k / n$分位数，这就是MQB的关键变量$\mathbf{\mathit{\beta}}_{i}$。由于EMA这个操作是线性的、可并行的，因此我们就以一种相对高效的、Loss-Free的方式实现了序列级均衡，这便是最终版的MQB。

$$ \text{Moving Quantile Balancing }(\text{MQB}) \\ 输入:\text{ }打分矩阵\textrm{ } \mathbf{\mathit{s}} \in \left[\right. 0 , 1 \left]\right.^{l \times n} \\ 输出:\text{ }校正后的打分矩阵\textrm{ } \hat{\mathbf{\mathit{s}}} \in \mathbb{R}^{l \times n} \\ \begin{matrix}1 : & \mathbf{\mathit{h}}_{i , j} = \text{one}_\text{hot} ⁡ \left(\right. \lfloor s_{i , j} \times b \rfloor \left.\right) \in \left{\right. 0 , 1 \left.\right}^{b} \\ 2 : & \bar{\mathbf{\mathit{h}}}_{i , j} = \gamma \bar{\mathbf{\mathit{h}}}_{i - 1 , j} + \left(\right. 1 - \gamma \left.\right) \mathbf{\mathit{h}}_{i , j} \\ 3 : & m_{i , j}^{*} = min \left{\right. m \left|\right. \sum_{t = 0}^{m} \bar{h}_{i , j , t} \geq 1 - \frac{k}{n} \left.\right} \\ 4 : & \beta_{i , j} = \left(\right. m_{i , j}^{*} + 1 / 2 \left.\right) / b \\ 5 : & \hat{\mathbf{\mathit{s}}}_{i} = \mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}_{i} ,\end{matrix} $$

一般情况[#](https://spaces.ac.cn/archives/11760#%E4%B8%80%E8%88%AC%E6%83%85%E5%86%B5)

到目前为止，我们的讨论都是基于动态激活版MoE的，即激活$\mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}_{i} > 0$的Expert，数目不定。从《MoE环游记：6、最优分配促均衡》我们知道，如果是Top-$k$版MoE，它的$\mathbf{\mathit{\beta}}$并没有简单的解析解，需要交替迭代。

那么Top-$k$版MoE又该如何以Loss-Free的方式去实现序列级负载均衡呢？实际上，MQB在动态激活下能够实现序列级均衡，表明它已经把Router局部的异常高峰“削平”了，这时候我们对$\mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}_{i}$取Top-$k$，它也许是不均衡的，但这种不均衡至多是全局级别的，那么我们只需给它补上一步全局均衡的QB，就可以重新恢复均衡性。

也就是说，对于Top-$k$版MoE，我们的强制序列级均衡方案是MQB+QB，在$\mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}_{i}$序列的基础上再做一次QB（当然换SignSGD也行），即可实现序列级负载均衡。但要注意，完美的序列级负载均衡往往会明显损伤效果，一般情况下全局均衡够用了，序列级均衡只是用来预防极端情况，所以施加程度不宜过重。

为了降低对主任务的影响，Aux Loss方案可以选择调低损失系数，而MQB则可以考虑将前后的分数做个加权平均：

$$ (\text{6}) \lambda \left(\right. \mathbf{\mathit{s}}_{i} - \mathbf{\mathit{\beta}}_{i} \left.\right) + \left(\right. 1 - \lambda \left.\right) \mathbf{\mathit{s}}_{i} = \mathbf{\mathit{s}}_{i} - \lambda \mathbf{\mathit{\beta}}_{i} , \lambda \in \left[\right. 0 , 1 \left]\right. $$

也就是说，通过引入$\lambda < 1$来削弱序列均衡的程度，然后在此基础上叠加QB来保证全局均衡（由于$\lambda < 1$是无法保证均衡的，所以这时候无论是动态版还是Top-$k$版都需要补上QB）。

其他细节[#](https://spaces.ac.cn/archives/11760#%E5%85%B6%E4%BB%96%E7%BB%86%E8%8A%82)

经过分桶后，MQB的核心运算也变成EMA，所以兜兜转转，我们又回到了EMA，形式跟 @ChangJonathanC 所提的第一种方案类似，不同的是EMA的对象。@ChangJonathanC 的方案是直接对$l \times n$大小的原始分数做EMA，MQB则是将它扩张成$l \times n \times b$后再做EMA。

这个变化有点像是线性注意力通过外积来扩大容量，能够记忆的信息更多，而且QB的最优性也保证了它能实现序列级均衡，不再停留在经验做法。另一点跟线性注意力的相似性是，EMA本质上也是个RNN，所以推理阶段要多存一个$n \times b$大小的State向量，实践发现$b = 100$已经能取得不错的效果，所以这个大小的State应该尚可接受。

还需要指出的是，不管QB还是MQB，它们都只用于Expert的激活决策，最后乘上Expert的门控函数，还是基于减去偏置项之前的、原始的Router分数来计算，这是所有Loss-Free方案共同的特点，确保均衡干预项不直接改变MoE形式，以求对主模型的影响最小化。

最后，关于分桶方式，笔者目前的实验是通过Sigmoid激活然后$\left[\right. 0 , 1 \left]\right.$均匀离散化来分桶，而如何更精细地分桶，可能还有一定的实验空间，这些交给有兴趣的读者完成了。

实验结果[#](https://spaces.ac.cn/archives/11760#%E5%AE%9E%E9%AA%8C%E7%BB%93%E6%9E%9C)

对于MQB，笔者简单做了些实验，配置是一个总参数量约3B、128选4的MoE模型，MQB的滑动平均系数为0.99、分桶数默认为100，用MaxVio作为负载均衡指标。由于第一层MoE是最不均衡的，所以下面我们都只演示第一层的效果。

首先是$\lambda = 1$的完全体MQB：


MQB(λ=1)与QB、SignSGD的效果比较

可见$\lambda = 1$能实现非常完美的负载均衡，代价是Loss差了0.06，这是非常可观的掉点了，表明过于序列均衡是对效果不利的。但如果将$\lambda$调到0.3左右，那么Loss基本不掉，并且依然能促进负载均衡，如下图


不同的λ对应的负载均衡情况

延伸思考[#](https://spaces.ac.cn/archives/11760#%E5%BB%B6%E4%BC%B8%E6%80%9D%E8%80%83)

从上面的实验我们可以进一步确认，完美的序列级均衡（$\lambda = 1$）对效果损伤比较严重，如果取$\lambda = 0.3$，那么能够大致保证效果无损，同时改善每一层的均衡情况。

总的来说，本文只是回答“怎么用Loss-Free的方式实现序列级负载均衡”，但是否需要序列级负载均衡、为什么需要序列级负载均衡、需要何种程度上的序列级负载均衡，对此我们无法给出标准答案。一个初步的理解是，序列级的极度不均衡，可能会对序列本身的效果有所影响（坍缩为一个小模型），所以希望鼓励一定程度的序列级均衡。

除了MQB外，其实还有一些别的方案也能Loss-Free地实现序列级均衡，比如我们在《DeepSeek V4的tid2eid是怎么来的？》所介绍的Hash Routing，也算是一种强制序列均衡的简单方案。但Hash Routing已经重新定义了路由形式，属于非常规路线。

MQB这种方案的价值在于，它更好地兼容常规的MoE形式，并且将“调节Aux Loss的系数”转化成了“调节一个可解释的局部偏置项”，让我们可以更直观地控制干预强度，同时可并行性也保证了效率。由于它只影响Router决策而不注入梯度，所以对主任务的效果影响也尽可能小。

文章小结[#](https://spaces.ac.cn/archives/11760#%E6%96%87%E7%AB%A0%E5%B0%8F%E7%BB%93)

本文围绕“如何用Loss-Free方式实现序列级负载均衡”展开，从原本的Quantile Balancing（QB）出发，逐步推导出了一种名为Moving Quantile Balancing（MQB）方法，成功实现了这一目的。但序列级均衡是否真的必要、要做到什么程度，仍然是一个开放的问题。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11760](https://spaces.ac.cn/archives/11760 "MoE环游记：8、强制序列级均衡")_

_更详细的转载事宜请参考：_[《科学空间FAQ》](https://spaces.ac.cn/archives/6508#%E6%96%87%E7%AB%A0%E5%A6%82%E4%BD%95%E8%BD%AC%E8%BD%BD/%E5%BC%95%E7%94%A8 "《科学空间FAQ》")