Qwen刘大一恒提出PID-Ba稀疏MoE负载均衡方法

智猩猩AI整理

编辑：林夕

9月30日，Qwen团队研究者刘大一恒（Dayiheng Liu）等人提交了一篇新论文 《ID Balancing: Stable Training of Extremely Sparse MoE via PID-Based Load Control》，提出了一种新的MoE负载控制方法——ID Balancing。

这套方法把MoE的专家负载问题，与经典控制理论中的PID控制联系起来。

论文在768个专家的极稀疏MoE设置下进行了实验。当路由从Top-10进一步变成Top-5、Top-3时，ID Balancing依然能够保持较稳定的专家负载。

在Top-3-of-768设置下，ID Balancing将最差情况下的 MaxVio降低超过50%，训练过程中平均的 MinVio降低超过12%，同时保持了具有竞争力的语言模型和下游任务表现。

MoE的核心机制，是把不同Token分配给不同专家。

假设模型拥有768个专家，一个Token并不需要让全部专家参与计算，而是通过Router选择Top-K个专家。

这样可以在保持超大参数规模的同时，把每个Token的实际计算量控制在较低水平。

但随着专家数量增加、Top-K降低，专家之间的负载差距也会越来越明显。

有些专家可能收到大量Token，成为训练中的热门专家；另一些专家则长期处于低负载状态。

过去比较常见的办法，是增加辅助负载均衡损失，让Router尽量把Token平均分配给不同专家。

但这类方法也存在一个问题：负载均衡本身是一个随着训练不断变化的动态过程，固定的辅助目标未必能够很好地适应不同训练阶段。

因此，近年来也出现了不依赖辅助损失的方案。

论文重点分析了两种代表性方法：DeepSeek的Loss-Free方法，以及Kimi K3中的Quantile Balancing。

团队发现，这两种方法实际上都可以放进一个不完整的PID控制框架中理解。

DeepSeek Loss-Free更接近一个固定步长的积分控制器，而Quantile Balancing则可以看作一种广义比例控制器。

在此基础上，论文提出ID Balancing，也就是Integral-Derivative Balancing。

它加入了积分项和微分项。

• 积分项负责根据当前的负载误差进行持续修正，而且误差越大，调整幅度越大。

• 微分项则关注负载误差的变化趋势：如果当前的失衡正在进一步恶化，就提前加大修正力度；如果负载已经趋于稳定，则减少额外干预。

简单理解，就是让Router不仅知道现在谁太忙，还知道谁正在越来越忙。

论文首先验证了积分项的作用。

在18.9B模型、Top-3-of-768路由、120B Token的训练设置下，团队将幅度感知的积分控制与DeepSeek Loss-Free方法进行比较。

结果表明，如果积分更新幅度能够随着负载误差大小变化，就能更有效地降低训练早期的专家过载，同时减少部分专家持续处于低负载的情况。

进一步的消融实验也表明，积分增益 (K_i) 会直接影响负载控制效果。

但仅仅知道当前差多少还不够。

如果某个专家当前已经超载，而且超载程度还在快速增加，系统需要更早进行干预。

这正是微分项发挥作用的地方，团队比较了加入微分项前后的训练过程。

结果显示，加入微分控制后，可以进一步降低早期专家过载，同时减少Token向少数热门专家集中的情况。

因此，ID Balancing最终形成了一套比较清晰的控制机制：积分项负责持续纠偏，微分项负责抑制正在扩大的失衡。

这也是这篇论文和此前无辅助损失方法相比，一个比较核心的区别。

方法提出之后，真正需要验证的问题是：当专家数量达到768个，每个Token只选择其中3个专家时，这套方法还能不能真正控制住负载？

论文给出的结果比较明确。

团队首先在18.9B模型上进行实验，固定使用768个专家，并分别测试Top-10、Top-5和Top-3三种路由方式，训练数据规模为120B Token。

Figure 1实际上同时回答了两个问题。

在图1(a)和图1(b)中，团队比较了Top-10、Top-5、Top-3-of-768三种设置下的负载情况。

其中，MaxVio用于衡量最严重的专家过载程度，MinVio则反映专家欠载情况。

随着Top-K从10降低到5，再降低到3，路由越来越稀疏，但ID Balancing依然保持了较稳定的负载控制。

尤其是在最极端的Top-3-of-768设置下，它仍然明显优于其他方法。

图1(c)则进一步把模型规模扩大。

团队将总参数量从18.9B增加到69.9B，激活参数规模从约1.03B增加到3.2B，同时固定使用Top-10-of-768路由。

随着模型不断扩大，ID Balancing的最坏Backbone MaxVio基本没有明显恶化。

相比辅助损失方法，其最坏负载偏差降低了约89.6%。

从具体数字来看，Top-3-of-768设置下，ID Balancing的Worst Backbone MaxVio为15.33。

相比之下，Quantile Balancing为31.11，DeepSeek Loss-Free达到211.19。

在平均欠载指标上，ID Balancing的Avg. MinVio为0.5480，也低于辅助损失方法的0.6232、DeepSeek Loss-Free的0.7441以及Quantile Balancing的0.6373。

与此同时，语言建模损失并没有因为负载控制而明显恶化。Top-3设置下，ID Balancing的LM loss为1.7426。

论文还进一步测试了MMLU、MMLU-Pro、SuperGPQA、MATH、GSM8K、BBH、MMMLU、EvalPlus、MultiPL-E等下游任务。

整体结果显示，ID Balancing在负载控制方面取得明显优势的同时，下游任务表现仍然具有竞争力。

团队还观察了训练过程中的动态变化。

在18.9B模型、Top-3-of-768路由下，Figure 2进一步比较了训练过程中的MaxVio和MinVio。

ID Balancing能够降低训练过程中出现的瞬时专家过载，同时减少部分专家长期处于低负载的情况。

团队还测试了更高学习率等不同训练条件。

在这些设置下，ID Balancing依然保持了较好的负载控制能力，说明这套机制并不只依赖某一种固定训练配置。

从这些实验可以看到，ID Balancing的优势随着MoE稀疏程度增加而更加明显。

这也正是它最有价值的地方。

因为当768个专家中每个Token只选择3个专家时，Router犯下的任何一次偏差，都更容易被放大。

热门专家不断吸收更多Token，冷门专家则可能越来越难获得训练信号。

ID Balancing通过持续反馈，让这种偏差在训练过程中不断被纠正。

事实上，负载均衡并不是Qwen团队第一次研究的问题。

此前，Qwen团队就曾提出 Global Load Balancing，尝试从更大的全局范围处理MoE中的专家负载分配。

传统的局部负载均衡方式，虽然能够让不同专家之间的Token数量更加接近，但也可能限制专家逐渐形成自己的专业化能力。

Global Load Balancing则希望在更大的范围内进行负载管理，在计算效率与专家专业化之间取得更好的平衡。

这次ID Balancing关注的则是另一个层面：

当MoE已经极度稀疏，专家数量进一步增加之后，怎样让这个庞大的路由系统本身保持稳定。

从这个角度看，两项工作关注的问题并不完全相同。

一个更强调全局范围内的负载分配，一个则进一步把专家负载视作动态系统进行控制。

而这背后其实指向了同一个趋势：MoE正在从增加专家走向管理越来越多的专家。

未来的大模型可能拥有更多专家、更大的参数规模，以及更低的激活比例。

到了那个阶段，如何设计路由器、如何控制专家负载、如何避免热门专家成为瓶颈，都会直接影响MoE模型能否继续向更大规模扩展。

这也是ID Balancing这项工作的意义所在。

它没有改变MoE最基本的稀疏计算范式，而是从一个看似基础、却越来越重要的问题入手——让大量专家在极稀疏条件下更稳定地工作。

对于正在不断扩大的MoE模型来说，专家数量可以继续增加，但如何让这些专家真正被高效利用，可能才是下一阶段更值得解决的问题。