面向专家混合模型的无辅助损失负载均衡策略

原文标题：Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts arXiv：2408.15664 ｜发布：28 Aug 2024 作者/机构：Lean Wang、Huazuo Gao、Chenggang Zhao、Xu Sun、Damai Dai（DeepSeek-AI；部分作者兼属北京大学计算机学院多媒体信息处理国家重点实验室） 原文 PDF：04-Aux-Loss-Free-Load-Balancing-2408.15664.pdf

一、摘要（全文翻译）

对于专家混合（Mixture-of-Experts, MoE）模型而言，专家负载不均衡会导致路由崩塌（routing collapse）或计算开销增加。现有方法通常采用辅助损失（auxiliary loss）来鼓励负载均衡，但较大的辅助损失会在训练中引入不可忽视的干扰梯度，从而损害模型性能。为了在控制负载均衡的同时不在训练中产生不需要的梯度，我们提出了 Loss-Free Balancing，其特色是一种无辅助损失的负载均衡策略。具体而言，在 top-K 路由决策之前，Loss-Free Balancing 会先对每个专家的路由分数施加一个专家级别的偏置（expert-wise bias）。通过根据每个专家近期的负载动态更新其偏置，Loss-Free Balancing 能够持续维持专家负载的均衡分布。此外，由于 Loss-Free Balancing 不产生任何干扰梯度，它也提升了 MoE 训练所能带来的模型性能上限。我们在参数量最高 3B、训练 token 数最高 200B 的 MoE 模型上验证了 Loss-Free Balancing 的性能。实验结果表明，与传统由辅助损失控制的负载均衡策略相比，Loss-Free Balancing 同时取得了更好的性能和更好的负载均衡。

二、核心贡献

提出无辅助损失的负载均衡策略（Loss-Free Balancing）：把负载均衡从「损失函数上的正则项」改为「路由分数上的可调偏置」，从根源消除辅助损失引入的干扰梯度，破解均衡与性能不可兼得的矛盾。
给出可迭代的专家偏置更新算法：在 top-K 选择前为每个专家引入偏置 $b_i$，每个 batch 后按「重负载降偏置、轻负载升偏置」更新，并实证确定更新率 $u=0.001$。
打破「均衡 vs. 性能」的两难：论文在 \alpha\in\{0,10^{-4},10^{-3},10^{-2}\} 上展示辅助损失系数越小越不均衡、越大性能越差（Figure 2），本方法在两方面均优于该基线。
均衡度呈量级改善：1B 模型验证困惑度由 9.56 降至 9.50、\text{MaxVio}_{\text{global}} 由 0.72 降至 0.04；3B 模型困惑度由 7.97 降至 7.92、\text{MaxVio}_{\text{global}} 由 0.52 降至 0.04。
与专家并行（Expert Parallelism）天然兼容：computation batch 越大均衡越好，而辅助损失方法在大 batch 下基本停滞。
揭示 Expert Choice（EC）的未来 token 泄漏：理论上每 token 泄漏量可超过 K\log_2\frac{1-R}{R} bit，并有实验佐证；本方法遵守因果约束。

三、方法与技术要点

3.1 背景：MoE 层与辅助损失的困境

MoE 层用多个专家 FFN 替换标准 Transformer 的 MLP 层，对第 t 个 token 的输入 $\mathbf{u}_t$，输出为

\mathbf{h}_t=\mathbf{u}_t+\sum_{i=1}^{N}g_{i,t}\operatorname{FFN}_i(\mathbf{u}_t),\quad s_{i,t}=G(\mathbf{u}_t^{\top}\mathbf{e}_i),

其中 g_{i,t} 仅在 s_{i,t} 落入 top-K 时为 $s_{i,t}$，否则为 0。负载不均衡的两类危害是路由崩塌（只选少数专家，其余专家训练不足）与多设备下的计算瓶颈。为此主流做法引入辅助损失：

\mathcal{L}_{\text{Balance}}=\alpha\sum_{i=1}^{N}f_i P_i,\quad f_i=\frac{N}{KT}\sum_{t=1}^{T}\mathbbm{1}(\text{Token }t\text{ selects Expert }i),\quad P_i=\frac{1}{T}\sum_{t=1}^{T}s_{i,t}.

痛点是：它本质上是额外正则项，会与语言建模目标产生冲突梯度。\alpha 太小则负载不均衡甚至路由崩塌，太大则明显拖累性能——即论文所称「负载均衡与模型性能之间的两难」。

3.2 核心设计：偏置门控分数（biased gating score）

Loss-Free Balancing 不改损失函数，而是改路由的判定依据：为每个专家加一个偏置 $b_i$，用带偏置的分数决定 top-K 归属，但最终加权输出仍用原始分数：

g_{i,t}=\begin{cases}s_{i,t},& s_{i,t}+b_i\in\operatorname{Topk}\left(\{s_{j,t}+b_j\mid 1\le j\le N\},K\right),\\ 0,&\text{otherwise.}\end{cases}

关键在于：b_i 只影响「选谁」、不进入 $g_{i,t}$，因此既不改变专家输出的加权方式，也不向反向传播注入与语言建模目标无关的梯度；均衡由偏置数值承担，性能由主损失主导，两者解耦。

3.3 偏置更新算法

偏置按「重负载压低、轻负载抬高」的原则迭代更新。对每个 batch：先按上式正常前向与训练，然后统计各专家被分配的 token 数 c_i 及其均值 $\overline{c_i}$，得到负载违背误差 $e_i=\overline{c_i}-c_i$，再更新

b_i \leftarrow b_i + u\cdot\operatorname{sign}(e_i).

作者强调只使用历史 batch 的负载信息：用当前序列的负载会破坏语言建模的因果约束，导致未来 token 信息泄漏。这一「事后更新」性质使本方法兼具控制力与因果安全性。

3.4 与其他负载均衡方法的对比

论文从三个维度刻画各方案：是否均衡、是否引入干扰梯度、是否泄漏未来 token。

负载均衡方法	专家负载均衡	干扰梯度	未来 token 泄漏
辅助损失控制（强）	均衡	强	无泄漏
辅助损失控制（弱）	不均衡	弱	无泄漏
Expert Choice	均衡	无	有泄漏
Loss-Free（本文）	均衡	无	无泄漏

EC 让每个专家被分配完全相同数量的 token 以实现完美均衡，但每个 token 的目标专家取决于同序列/批次中的未来 token，破坏因果约束；辅助损失方法则始终需权衡均衡与性能。本文方法是唯一同时满足三者的方案。

3.5 算法变体的实证选择

论文对偏置更新算法做了消融，据此选定主实验配置：

更新率：u=0.0001 收敛太慢（训练早期均衡差），u=0.01 令 b_i 后期剧烈波动、反而恶化均衡，u=0.001 兼顾两者。
更新规则：换成 $b_i+u\cdot e_i$（违背更严重的专家偏置变化更快）虽略微改善均衡，但困惑度都更差，故保留 sign 版本。
加性 vs. 乘性偏置：乘性形式 s_{i,t}*b_i 参与 top-K（初始化为 1），均衡相近而性能略差，故采用加性偏置。
softmax 门控：softmax 的归一化使两个专家间的分数差受其他专家影响，偏置更难调，此时改用 b_i+u\cdot e_i 变体且 $u=10^{-3}$。

四、主要实验结果

基准 / 任务	本模型	对比对象	说明
1B MoE，验证困惑度	9.50（Loss-Free）	9.56（辅助损失控制，$\alpha$=0.001）	1B 模型训练 100B tokens
1B MoE，`\text{MaxVio}_{\text{global}}`	0.04	0.72	均衡度提升约 18 倍
3B MoE，验证困惑度	7.92（Loss-Free）	7.97（辅助损失控制）	3B 模型训练 200B tokens
3B MoE，`\text{MaxVio}_{\text{global}}`	0.04	0.52	均衡度提升约 13 倍
1B，更新规则消融（$u$=0.001）	9.50 / 0.044（sign 版本）	9.53 / 0.028、9.51 / 0.036、9.51 / 0.040（`u\cdot e_i` 变体）	sign 版本性能最优
1B，偏置形式消融	9.50 / 0.044（加性）	9.52 / 0.041、9.52 / 0.036、9.54 / 0.048（乘性 Bias）	加性偏置性能最好
1B，softmax 门控	9.599 / 0.027（Loss-Free）	9.604 / 0.937（Loss-Controlled，$\alpha$=0.0003）	困惑度略优，均衡显著更优
Expert Choice 泄漏实验（2B MoE / 100B tokens）	chunk=512 出现约 10% 异常 loss 下降	chunk=2048 / 8192 无此现象；打乱 token 后异常消失	证明未来 token 泄漏真实存在
理论泄漏量（9 层 MoE，16 专家，平均激活 2）	每 token 超过 50 bit	—	足以让每个 token 推断其后继 token 的身份

在 1B 与 3B 两个规模上，本方法同时取得更低困惑度与更好的全局负载均衡：困惑度领先 0.05–0.06，而 \text{MaxVio}_{\text{global}} 从 0.72/0.52 降到 0.04，属量级差异。训练曲线（Figure 3）与 computation batch 扫描（Figure 5）显示其均衡在绝大部分训练步中更优、且 batch 越大优势越明显，而辅助损失方法在大 batch 下基本不再改善。论文未报告下游任务评测。

五、结论、局限与影响

结论：论文提出 Loss-Free Balancing——一种不引入辅助损失梯度的 MoE 负载均衡控制方法，解决了辅助损失在强制均衡时引入额外梯度、损害性能的问题。在 1B/3B MoE 上，相比辅助损失训练，本方法同时取得更好的性能与负载均衡。
局限性：论文未设单独的局限性章节，可归纳的自述边界为：（1）实验规模仅到 3B 参数、最多 200B tokens；（2）主实验固定用 sigmoid 门控（其基线优于 softmax），softmax 场景需改用 u\cdot e_i 变体并单独调参；（3）3B 模型训练 token 数前后不一致——摘要、正文与附录 B 写 200B，第 6 节结论写 300B；（4）未报告下游任务指标，仅以验证困惑度与 \text{MaxVio} 衡量。
影响：该策略已被 DeepSeek-V3 采用，成为大规模 MoE 训练中替代辅助损失的实用方案：它把负载均衡从「损失层面」移到「路由层面」，解除均衡强度与性能的绑定；又因与专家并行天然兼容，并行规模越大优势越显著。论文对 EC 泄漏的分析也警示了 MoE 路由的因果安全性。

六、关键术语对照表

English	中文
Mixture-of-Experts (MoE)	专家混合
Loss-Free Balancing	无辅助损失负载均衡
Auxiliary Loss / Loss-Controlled	辅助损失 / 辅助损失控制
Interference Gradient	干扰梯度
Routing Collapse	路由崩塌
Expert-wise Bias	专家级偏置
Biased Gating Score	偏置门控分数
Top-K Routing	Top-K 路由
Bias Update Rate	偏置更新率
Maximal Violation (MaxVio)	最大违背度
Expert Choice (EC)	专家选择路由
Future Token Leakage	未来 token 泄漏
Expert Parallelism	专家并行
Computation Batch	计算批次
DeepSeekMoE	深度求索 MoE 架构

10 KiB Raw Blame History Unescape Escape