114 lines
10 KiB
Markdown
114 lines
10 KiB
Markdown
# 面向专家混合模型的无辅助损失负载均衡策略
|
||
|
||
> **原文标题**：Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
|
||
> **arXiv**：[2408.15664](https://arxiv.org/abs/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](./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 同时取得了更好的性能和更好的负载均衡。
|
||
|
||
## 二、核心贡献
|
||
|
||
1. **提出无辅助损失的负载均衡策略（Loss-Free Balancing）**：把负载均衡从「损失函数上的正则项」改为「路由分数上的可调偏置」，从根源消除辅助损失引入的干扰梯度，破解均衡与性能不可兼得的矛盾。
|
||
2. **给出可迭代的专家偏置更新算法**：在 top-K 选择前为每个专家引入偏置 $b_i$，每个 batch 后按「重负载降偏置、轻负载升偏置」更新，并实证确定更新率 $u=0.001$。
|
||
3. **打破「均衡 vs. 性能」的两难**：论文在 $\alpha\in\{0,10^{-4},10^{-3},10^{-2}\}$ 上展示辅助损失系数越小越不均衡、越大性能越差（Figure 2），本方法在两方面均优于该基线。
|
||
4. **均衡度呈量级改善**：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。
|
||
5. **与专家并行（Expert Parallelism）天然兼容**：computation batch 越大均衡越好，而辅助损失方法在大 batch 下基本停滞。
|
||
6. **揭示 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 架构 |
|