banner
约 5,100 字
17 分钟

Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts论文阅读

摘要

本文分析 DeepSeek-AI 提出的无辅助损失 MoE 负载均衡方法。该方法以专家偏置构成反馈控制回路,只改变 Top-K 路由选择,不向语言建模目标引入额外梯度,并在 1B 和 3B 模型上同时改善困惑度与专家负载均衡。

Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts论文阅读

写在前面

Mixture-of-Experts,混合专家模型(MoE)的核心价值是稀疏激活:模型可以扩展总参数量,但每个 Token 只调用少量专家,从而控制单次前向计算量。这个机制成立的前提是路由器能把 Token 相对均匀地分配给不同专家。若少数专家持续过载,MoE 会出现计算瓶颈,未被充分使用的专家也难以得到有效训练。

传统方法在训练目标中加入负载均衡辅助损失。辅助损失过弱时,路由仍可能失衡;辅助损失过强时,额外梯度会干扰语言建模目标。本文提出的 Loss-Free Balancing 将负载均衡从损失函数中移出,改为一个基于历史负载的反馈控制过程:过载专家降低路由偏置,欠载专家提高路由偏置,偏置只参与 Top-K 选专家,不参与最终输出加权。

全文按照以下主线展开。

部分

核心问题

结论

研究问题

为什么辅助损失难以同时保证均衡和性能

负载均衡与语言建模目标共享梯度,存在直接冲突

核心方法

如何在不增加辅助损失的条件下调节专家负载

用专家偏置和批次负载统计构成闭环控制

实验验证

方法是否同时改善模型性能与负载均衡

1B 和 3B 实验中困惑度更低,最大负载偏差降至 0.04

工程分析

方法在专家并行和不同门控函数下是否成立

适合较大的计算批次,也可迁移到 Softmax 门控,但更新更敏感

边界与启发

结论能否直接推广到更大模型和生产系统

方法方向明确,但仍缺少超大规模训练、端到端吞吐与下游任务验证

速读卡片

项目

信息

论文题目

Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

作者

Lean Wang、Huazuo Gao、Chenggang Zhao、Xu Sun、Damai Dai

机构

DeepSeek-AI;北京大学多媒体信息处理全国重点实验室、计算机学院

文档类型

arXiv 预印本,2024 年 8 月 28 日提交

研究对象

自回归语言模型中的 MoE 专家负载均衡

核心方法

Loss-Free Balancing:基于历史专家负载迭代更新专家偏置

实验规模

1B 模型训练 100B Token;3B 模型训练 200B Token

核心结果

1B 与 3B 模型的验证困惑度分别由 9.56、7.97 降至 9.50、7.92,全局最大负载偏差均降至 0.04

论文链接

arXiv 页面论文 PDF

一句话概括:本文不再通过辅助损失强迫专家均衡,而是把专家偏置作为控制变量,根据上一批次的负载误差持续修正下一批次的路由决策。

一、研究问题:辅助损失为什么不够理想

1.1 MoE 路由与负载失衡

设第 tt 个 Token 的隐藏状态为 ut\mathbf{u}_t,MoE 层包含 NN 个专家。门控网络为每个专家计算分数 si,ts_{i,t},再从中选择分数最高的 KK 个专家:

ht=ut+i=1Ngi,tFFNi(ut)\mathbf{h}_t = \mathbf{u}_t + \sum_{i=1}^{N} g_{i,t}\operatorname{FFN}_i(\mathbf{u}_t)
gi,t={si,t,si,tTopK({sj,t}j=1N,K)0,其他情况g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t} \in \operatorname{TopK}(\{s_{j,t}\}_{j=1}^{N}, K) \\ 0, & \text{其他情况} \end{cases}

其中,Feed-Forward Network,前馈神经网络(FFN)表示专家网络。Top-K 路由使每个 Token 只激活少量专家,但门控网络没有天然的均衡约束。训练早期的微小分数差异可能被不断放大,使少数专家持续接收更多 Token,最终形成路由坍缩。

负载失衡会产生两个直接后果。第一,专家参数的训练机会不均等,模型容量没有被充分利用。第二,在 Expert Parallelism,专家并行(EP)中,不同专家分布在不同设备上,单个过载专家会延长整个计算步骤,其他设备即使已经完成计算也必须等待。

1.2 辅助损失带来的目标冲突

常见做法是在语言建模损失之外增加负载均衡损失:

LBalance=αi=1NfiPi\mathcal{L}_{\text{Balance}} = \alpha \sum_{i=1}^{N} f_i P_i

fif_i 表示分配给专家 ii 的 Token 比例,PiP_i 表示该专家的平均门控分数,α\alpha 控制辅助损失强度。该目标能够压制热门专家,但它也会通过梯度修改路由器。此时模型同时优化两个不完全一致的目标:语言建模希望选择最适合当前 Token 的专家,均衡目标则要求减少热门专家的选择次数。

辅助损失下负载均衡与模型性能的矛盾
辅助损失下负载均衡与模型性能的矛盾

图 1:辅助损失系数变化时,全局负载偏差与验证困惑度之间存在明显权衡。来源:论文 Figure 2。

图 1 中,α=0\alpha=0 时没有额外干扰,但负载偏差最大;提高 α\alpha 后负载逐步改善,困惑度却先降后升。Loss-Free Balancing 位于图的左下区域,说明作者希望同时获得更低困惑度和更小负载偏差,而不是在两者之间继续调参折中。

二、核心方法:把负载均衡改写为反馈控制

Loss-Free Balancing 的关键不是设计另一种可微损失,而是把路由选择和负载调节拆成两条链路。前向计算仍由原始门控分数决定输出权重;负载控制器只根据上一批次的统计结果修正后续 Top-K 选择。

无辅助损失负载均衡闭环
无辅助损失负载均衡闭环

图 2:Loss-Free Balancing 的前向计算与偏置反馈回路(图由AI辅助绘制)。

2.1 偏置只改变专家选择

作者为每个专家维护一个偏置 bib_i。Top-K 选择依据从 si,ts_{i,t} 改为 si,t+bis_{i,t}+b_i,但被选中专家的输出权重仍是原始分数 si,ts_{i,t}

gi,t={si,t,si,t+biTopK({sj,t+bj}j=1N,K)0,其他情况g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t}+b_i \in \operatorname{TopK}(\{s_{j,t}+b_j\}_{j=1}^{N}, K) \\ 0, & \text{其他情况} \end{cases}

这一区分是整个方法的核心。偏置可以改变离散路由结果,却不会直接放大或缩小专家输出。模型仍使用原始门控分数表示当前 Token 与专家之间的匹配程度,负载控制信号只负责在候选专家接近时调整选择边界。

2.2 根据批次负载更新偏置

完成一个训练批次后,系统统计每个专家接收的 Token 数量 cic_i,并计算理想平均负载 cˉ\bar{c}

ei=cˉcie_i = \bar{c} - c_i

偏置按误差符号更新:

bibi+usign(ei)b_i \leftarrow b_i + u\cdot \operatorname{sign}(e_i)

uu 是偏置更新率。若专家 ii 过载,则 ci>cˉc_i>\bar{c}ei<0e_i<0,偏置下降;若专家欠载,则偏置上升。控制器不要求精确预测下一批次的分布,只需要不断修正负载方向。

论文中的偏置反馈机制
论文中的偏置反馈机制

图 3:训练阶段先在门控分数上加入专家偏置,再根据观察到的专家负载更新偏置。来源:论文 Figure 1。

从控制系统角度看,专家负载是被控量,理想平均负载是目标值,bib_i 是控制量,uu 决定控制器的响应速度。该机制本质上是带有一个批次延迟的离散反馈控制器。

2.3 为什么不产生辅助损失干扰

传统方法将均衡目标加入总损失,其梯度会直接改变模型参数。Loss-Free Balancing 的 bib_i 由统计规则更新,不经过反向传播,也不属于语言建模损失。它仍会通过改变专家集合间接影响训练数据分配,但不会产生一条持续拉动门控分数的辅助梯度。

因此,论文中的 loss-free 应理解为不使用负载均衡辅助损失,而不是整个训练没有损失函数,也不是负载控制对模型完全没有影响。

2.4 为什么不会引入未来 Token 信息

偏置更新使用已经完成的历史批次,当前 Token 的路由只依赖当前隐藏状态、当前门控分数和历史偏置。未来 Token 不会反向改变当前 Token 的专家分配,因此仍满足自回归语言模型的因果约束。

三、与其他负载均衡方法的关系

论文将四种策略放在同一框架下比较。表中的无干扰梯度只表示不增加负载均衡梯度,不表示模型没有主任务梯度。

方法

专家负载

额外干扰梯度

未来 Token 泄漏

主要问题

强辅助损失

较均衡

语言建模性能可能下降

弱辅助损失

易失衡

可能出现路由坍缩和计算瓶颈

Expert Choice,专家选择(EC)

可严格均衡

不适合直接用于因果语言模型训练

Loss-Free Balancing

较均衡

需要选择合适的偏置更新率

Loss-Free Balancing 不是让每个专家在每个微批次中接收完全相同数量的 Token,而是通过反馈使负载逐步接近均衡。它放弃了逐批次严格相等,换取因果安全、较低目标干扰和较简单的分布式实现。

四、实验设计与评价指标

4.1 模型与训练设置

实验使用 DeepSeekMoE 作为骨干。其特点是将路由专家细粒度化,并保留共享专家处理通用知识。主实验使用 Sigmoid 门控,附录补充 Softmax 门控结果。

设置

1B 模型

3B 模型

总参数量

1B

3B

隐藏维度

1024

1280

MoE 层数

9

11

路由专家数

64

64

每个 Token 激活的路由专家数

6

6

共享专家数

2

2

训练 Token 数

100B

200B

最大序列长度

2048

2048

验证集约包含 71M Token。作者只在 1B 模型上搜索偏置更新配置,3B 模型直接继承该配置,用于检验方法能否跨规模迁移。

论文摘要、正文实验设置和附录训练配置均写明 3B 模型使用 200B Token;结论部分出现了 300B Token。由于三个相互独立的位置均支持 200B,本文将结论中的 300B 视为排版错误。

4.2 困惑度与最大负载偏差

模型性能使用 Perplexity,困惑度(PPL)衡量,数值越低越好。负载均衡使用 Maximal Violation,最大负载偏差(MaxVio):

MaxVio=maxiLoadiLoadLoad\operatorname{MaxVio} = \frac{\max_i \operatorname{Load}_i-\overline{\operatorname{Load}}}{\overline{\operatorname{Load}}}

该指标衡量最繁忙专家超过理想平均负载的比例。MaxVio=0.04\operatorname{MaxVio}=0.04 表示最繁忙专家的负载比理想均值高约 4%。论文使用三个统计尺度:

指标

统计范围

主要含义

MaxVioglobal\operatorname{MaxVio}_{\text{global}}

完整验证集

长期专家利用是否均衡

MaxViobatch\operatorname{MaxVio}_{\text{batch}}

单个训练批次

训练过程中是否存在瞬时过载

MaxViocomputation-batch\operatorname{MaxVio}_{\text{computation-batch}}

专家并行中的单次计算批次

不同设备之间是否会相互等待

只报告全局指标可能掩盖瞬时拥塞。一个专家可以在前半段过载、后半段欠载,最终全局统计接近均衡,但训练过程仍会产生计算瓶颈。因此,论文同时报告全局、训练批次和计算批次三个尺度。

五、主要结果:均衡改善并未以性能为代价

1B 与 3B 模型的主要实验结果
1B 与 3B 模型的主要实验结果

图 4:表 2 给出最终困惑度与全局负载偏差,Figure 3 给出训练期间的批次负载曲线。来源:论文 Table 2 与 Figure 3。

模型

方法

验证 PPL

MaxVioglobal\operatorname{MaxVio}_{\text{global}}

相对结论

1B

Loss-Controlled

9.56

0.72

基线

1B

Loss-Free

9.50

0.04

PPL 下降约 0.63%,最大偏差下降约 94.4%

3B

Loss-Controlled

7.97

0.52

基线

3B

Loss-Free

7.92

0.04

PPL 下降约 0.63%,最大偏差下降约 92.3%

困惑度改善并不大,但方向在两个规模上保持一致。更重要的是,负载偏差从 0.72 和 0.52 降到 0.04,说明作者没有通过牺牲均衡换取较低困惑度。

训练曲线进一步说明两类指标不能互相替代。Loss-Free 的批次负载偏差在训练早期迅速下降,随后维持在较低范围;辅助损失基线在 1B 和 3B 模型上分别长期停留在约 0.7 和 0.5。最终全局偏差接近零,不代表每个批次严格均衡,但足以显著缓解长期热点专家问题。

六、偏置更新的消融实验

6.1 更新率决定响应速度与稳定性

偏置更新率与更新规则消融实验
偏置更新率与更新规则消融实验

图 5:不同更新率的训练曲线,以及符号更新与误差幅值更新的对比。来源:论文 Figure 4 与 Table 3。

uu 太小,控制器无法及时纠正训练早期形成的负载偏差;uu 太大,偏置会在均衡点附近反复越界,训练后期出现明显振荡。实验中:

更新率

训练现象

验证 PPL

0.0001

早期收敛过慢

9.51

0.001

收敛速度与后期稳定性较好

9.50

0.01

后期负载波动增大

9.51

作者最终选择 u=0.001u=0.001。这一结果说明偏置更新不是越快越好,它需要与门控分数尺度、批次大小和训练阶段共同匹配。

6.2 更精确的负载控制不一定带来更好模型

默认方法只使用误差方向:

bibi+usign(ei)b_i \leftarrow b_i + u\operatorname{sign}(e_i)

作者还测试了按误差幅值更新:

bibi+ueib_i \leftarrow b_i + ue_i

误差幅值更新在 u=0.01u=0.01 时将全局最大偏差进一步降到 0.028,但验证 PPL 上升到 9.53;符号更新的全局最大偏差为 0.044,PPL 则达到最优的 9.50。负载指标更低并不等价于语言模型更好,工程调参应把负载均衡视为约束,而不是唯一优化目标。

6.3 加法偏置优于乘法偏置

作者还将 si,t+bis_{i,t}+b_i 改为 si,tbis_{i,t}b_i。乘法偏置可以得到相近的均衡程度,但最优 PPL 为 9.52,仍弱于加法偏置的 9.50。加法形式直接平移专家选择边界,更新方向稳定,也不受原始分数绝对值的乘性放大,因此更适合作为默认方案。

6.4 Softmax 门控仍然可用,但更难调节

Softmax 会把所有专家分数耦合在同一归一化分母中。修改一个专家的分数后,其他专家的相对间隔也会变化,因此简单的符号更新不如 Sigmoid 场景稳定。附录改用误差幅值更新,设置 u=0.001u=0.001

Softmax 门控方法

验证 PPL

MaxVioglobal\operatorname{MaxVio}_{\text{global}}

Loss-Controlled

9.604

0.937

Loss-Free

9.599

0.027

这组结果说明方法不依赖 Sigmoid,但控制器参数和更新规则需要适配门控函数,不能机械复用。

七、专家并行中的实际意义

在专家并行中,不同专家被放置到不同设备。单次计算步骤处理的样本数可写为:

纯文本
computation_batch_size = micro_batch_size × ep_data_parallel_size

计算批次越大,单批次 Token 分布越接近长期分布,历史反馈偏置也越能稳定发挥作用。

专家并行下计算批次大小与负载均衡
专家并行下计算批次大小与负载均衡

图 6:随着计算批次增大,Loss-Free Balancing 的计算批次负载偏差持续下降。来源:论文 Figure 5。

图 6 显示,在较小计算批次下,Loss-Free 的瞬时均衡不一定优于辅助损失;达到论文标注的中等计算批次后,Loss-Free 曲线低于基线,并随批次继续增大而改善。1B 和 3B 模型对应的交叉点约为 18 和 36。

这支持方法与大规模专家并行相容,但不能直接等价为训练吞吐已经提高。论文测量的是负载偏差,没有给出端到端吞吐、设备空闲时间、All-to-All 通信开销或单位 Token 训练成本。工程部署仍需通过真实硬件指标验证收益。

八、Expert Choice 的未来 Token 泄漏

Expert Choice 不再让每个 Token 选择专家,而是让每个专家从一个 Token 块中选择固定数量的 Token,因此可以严格控制专家容量。问题在于,较早 Token 是否被某个专家选中,会取决于同一块中较晚 Token 的门控分数。

Expert Choice 中的未来 Token 泄漏机制
Expert Choice 中的未来 Token 泄漏机制

图 7:后续 Token 的分数会改变前序 Token 的专家分配,路由结果因而携带未来信息。来源:论文 Figure 6。

图 7 的例子中,前序词 the 的专家分配会受到后续词 cat 或 dog 的竞争结果影响。模型可以从前序 Token 进入哪个专家,间接推断后续 Token 的身份。对于因果语言模型,这相当于把标签信息编码进路由决策。

设每个 Token 平均激活 KK 个专家,专家总数为 NN,稀疏率 R=K/NR=K/N。论文给出的单层潜在泄漏量下界为:

Klog21RRbits/TokenK\log_2\frac{1-R}{R}\quad \text{bits/Token}

N=16N=16K=2K=2 时,单层约可携带 5.61 bit 信息;9 个 MoE 层累计超过 50 bit/Token。这个计算描述的是路由组合的潜在信息容量,不代表模型一定利用了全部容量,但足以说明风险不能忽略。

Expert Choice 不同块大小与打乱策略的训练损失
Expert Choice 不同块大小与打乱策略的训练损失

图 8:块大小为 512 时出现异常损失下降;在分块前打乱 Token 后,该现象基本消失。来源:论文 Figure 9。

作者在 2B MoE 模型上训练 100B Token,并比较 512、2048、8192 三种块大小。块大小为 512 时训练损失出现约 10% 的异常下降;把 Token 打乱后再分块,下降明显减弱。打乱破坏了未来 Token 与原始上下文的对应关系,因此这组对照实验支持未来信息泄漏解释,而不是简单的优化加速。

这一分析是论文容易被忽略的贡献。负载完全相等不是最高优先级;对于自回归训练,任何依赖同一序列未来位置的批量路由决策,都必须先检查因果性。

九、论文的核心贡献

贡献

具体内容

价值

方法贡献

用专家级偏置和历史负载统计替代辅助损失

避免负载均衡梯度直接干扰语言建模目标

机制贡献

将选专家分数与输出加权分数分离

只改变离散路由边界,保留原始门控分数的语义

实验贡献

在 1B/100B Token 和 3B/200B Token 设置中验证

同时改善 PPL 和全局、批次负载均衡

工程贡献

分析计算批次大小与专家并行的关系

说明方法在较大专家并行批次中的适用条件

安全性贡献

理论分析并实验证明 Expert Choice 的未来信息泄漏

将因果约束纳入 MoE 负载均衡方法评价

十、局限性与未来工作

10.1 当前证据的边界

第一,最大模型规模为 3B,训练量为 200B Token。结果证明方法能够从 1B 迁移到 3B,但还不能单独证明在数百亿或更大参数规模下具有同样的稳定性。

第二,模型质量主要用验证困惑度衡量,没有报告知识、推理、代码或下游任务结果。较低 PPL 通常是积极信号,但不能替代完整能力评估。

第三,专家并行实验报告了负载指标,没有报告端到端训练速度和硬件利用率。负载更均衡可能减少尾部等待,也可能引入额外统计或通信开销,实际收益仍需系统测量。

第四,固定更新率依赖具体训练设置。批次大小、专家数、门控分数尺度和数据分布变化都可能改变最合适的 uu。论文给出了有效经验值,但没有证明偏置动态的收敛范围。

10.2 值得继续研究的方向

后续工作可以从四个方向展开。其一,根据负载误差、波动幅度和训练阶段自适应调整更新率,减少早期响应慢与后期振荡之间的矛盾。其二,为不同层或不同专家设置归一化控制尺度,避免统一 uu 无法适配门控分布。其三,在大规模异构集群中联合测量负载、通信、吞吐和故障恢复成本。其四,在下游任务和长上下文训练中检查路由稳定性,确认较低困惑度能否转化为实际能力提升。

十一、对工程实践的启发

11.1 把控制信号从主目标中分离

辅助目标经常被用来满足容量、延迟、公平性或稀疏性约束,但额外梯度可能改变主任务的最优方向。Loss-Free Balancing 提供了一种更通用的思路:对离散资源分配问题,可以优先考虑基于观测反馈调节决策边界,而不是把所有工程约束都写入同一个损失函数。

11.2 决策分数与输出权重不必共用同一变量

本文只对 Top-K 选择加入偏置,输出仍使用原始分数。该设计把资源调度和语义置信度分开:前者决定调用谁,后者决定结果占多大权重。类似原则也可用于检索路由、多模型调度和稀疏模块选择。

11.3 负载指标必须按系统执行粒度定义

全局均衡、训练批次均衡和计算批次均衡回答的是不同问题。全局统计适合观察专家是否被长期闲置,批次统计反映训练稳定性,计算批次统计才接近设备等待。实际系统不能只选择最容易变好的指标。

11.4 因果约束应覆盖数据之外的控制流

未来信息泄漏不只发生在特征和标签中,也可能进入批处理、排序、路由和缓存决策。只要当前样本的执行路径依赖未来位置,就可能形成隐蔽的信息通道。对自回归模型进行系统设计时,控制流也应纳入数据泄漏检查。

小结

这篇论文解决的不是如何设计更强的均衡损失,而是是否必须使用均衡损失。作者将专家负载视为可观测状态,将专家偏置视为控制变量,用上一批次的负载误差修正下一批次的 Top-K 路由。由于偏置不参与输出加权,也不通过辅助损失反向传播,模型能够减少负载失衡,同时保留语言建模目标的优化方向。

实验中,1B 和 3B 模型的全局最大负载偏差均降至 0.04,验证困惑度也小幅改善。消融实验进一步说明,更新率和门控函数决定控制器稳定性,负载越均衡并不必然代表模型性能越好。对 Expert Choice 的分析则补充了另一个关键判断:在因果语言模型中,严格均衡若依赖未来 Token,其训练结果不具备可信度。

本文最重要的启发是将模型学习与系统控制分开处理。主损失负责学习能力,反馈回路负责资源约束,两者通过路由决策连接,但不必共享同一条梯度路径。

END