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

论文地址:https://arxiv.org/abs/2408.15664
ICLR 2025:https://openreview.net/forum?id=y1iU5czYpE
学术交流群:922230617
目录
[1. 引言](#1. 引言)
[2. 背景](#2. 背景)
[2.1 MoE](#2.1 MoE)
[2.2 用于负载均衡的辅助损失](#2.2 用于负载均衡的辅助损失)
[3. 无辅助损失的负载均衡策略](#3. 无辅助损失的负载均衡策略)
[4. 实验](#4. 实验)
[4.1 实验设置](#4.1 实验设置)
[4.2 主要结果](#4.2 主要结果)
[4.3 偏置更新算法的实证研究](#4.3 偏置更新算法的实证研究)
[5. 讨论](#5. 讨论)
[5.1 Loss-Free Balancing 与专家并行兼容](#5.1 Loss-Free Balancing 与专家并行兼容)
[5.2 负载均衡与未来 Token 泄漏](#5.2 负载均衡与未来 Token 泄漏)
[6. 结论](#6. 结论)
[QB 算法与结果](#QB 算法与结果)
[QB‑Quantile‑Bias 关键问题整理](#QB‑Quantile‑Bias 关键问题整理)
[Q1:quantile_bias 里,分位数是固定的吗?β 是固定的吗?](#Q1:quantile_bias 里,分位数是固定的吗?β 是固定的吗?)
[Q2:bias 是动态的,被激活专家数是 k 还是动态?](#Q2:bias 是动态的,被激活专家数是 k 还是动态?)
[Q3:QB vs QB-DA 的区别是什么?两者都是通过 s-β>0 来选激活吗?](#Q3:QB vs QB-DA 的区别是什么?两者都是通过 s-β>0 来选激活吗?)
[MQB 算法与结果](#MQB 算法与结果)
1. 引言
混合专家模型(Mixture-of-Experts, MoE)已成为扩展大语言模型(LLMs)参数规模的同时控制计算成本的有效方案。近年来,MoE 在 Transformer-based 模型中的应用取得了显著成功。然而,训练 MoE 模型始终面临负载不均衡 的问题,这可能导致路由坍塌 (Routing Collapse)或计算开销增加。
现有方法通常采用辅助损失 (Auxiliary Loss)来鼓励均衡的专家负载。虽然辅助损失可以缓解训练中的负载不平衡,但它也会引入与语言建模目标相冲突的干扰梯度(Interference Gradients),从而损害模型性能。因此,现有 MoE 方法总需要在负载均衡 和模型性能之间进行权衡。
本文提出了 Loss-Free Balancing(无损失负载均衡),一种无需辅助损失的负载均衡策略。其核心思想是在 Top-K 路由决策之前,为每个专家施加一个专家维度的偏置(bias),并根据专家最近的负载情况动态更新该偏置。该方法不产生任何干扰梯度,从而提升了 MoE 训练的性能上限。
2. 背景
2.1 MoE
当前主流的 MoE 架构将 Transformer 中的标准 FFN 层替换为 MoE 层。在 MoE 层中,采用 Top-K 路由为每个 token 选择专家。设 u_t 为第 t 个token的输入,输出 h_t 计算如下:

其中,g_{i,t} 为路由权重,仅当专家 i 的分数 s_{i,t} 属于 Top-K 时取值为 s_{i,t},否则为 0。G 是非线性门控函数,e_i 是第 i 个专家的质心向量。
2.2 用于负载均衡的辅助损失
无控制的路由策略容易遭遇负载不均衡,带来两个主要问题:
-
路由坍塌:模型持续选择少数专家,导致其他专家训练不足。
-
计算瓶颈:专家分布在多个设备时,负载不均衡会加剧计算瓶颈。
为此,常用辅助损失(Auxiliary Loss)来控制负载平衡。对于长度为 T 的序列,辅助损失定义为:

其中:
-
f_i 表示路由到专家 i 的 token 比例。
-
P_i 表示专家 i 的平均门控分数。
-
α 控制辅助损失强度的超参数。
负载均衡与模型性能之间的困境:辅助损失虽然能促进负载均衡,但也作为额外的正则化项干扰了语言建模训练。辅助损失系数 α 过小会导致负载不均衡(甚至路由坍塌),过大则会损害模型性能。
下图清晰地展示了这一矛盾------传统方法无法同时达到最优的负载均衡和模型性能,而本文提出的 Loss-Free Balancing 则成功打破了这一困境。

3. 无辅助损失的负载均衡策略
为了在不引入干扰梯度的前提下实现负载均衡,本文提出 Loss-Free Balancing,其核心是直接根据各专家的负载状况调整门控分数。
具体做法是:为每个专家 i 添加一个偏置项 b_i,加到原始门控分数 s_{i,t} 上,然后使用偏置后的分数 s_{i,t} + b_i 来决定 Top-K 选择:

关键点 :偏置 b_i 仅用于影响 Top-K 选择,并不参与最终 MoE 层输出的加权计算。因此,它不会产生额外的梯度干扰。

偏置的更新策略如下(见算法1):
-
初始化所有 b_i = 0。
-
在每个训练 batch 之后,统计每个专家被分配的 token 数量 c_i,并计算平均值 ‾c_i。
-
计算负载偏差误差 e_i = ‾ci − ci。
-
按照 b_i = b_i + u⋅sign(e_i) 更新偏置,其中 u 是**偏置更新率,**sign 是符号函数。
与其他方法的对比:下表总结了不同负载均衡方法的特性(绿色表示好,红色表示差):
| 负载均衡方法 | 专家负载均衡 | 干扰梯度 | 未来 token 泄漏 |
|---|---|---|---|
| Loss-Controlled(强辅助损失) | ✅ 均衡 | ❌ 存在(强) | ✅ 无泄漏 |
| Loss-Controlled(弱辅助损失) | ❌ 不均衡 | ❌ 存在(弱) | ✅ 无泄漏 |
| Expert Choice (EC) | ✅ 均衡 | ✅ 无 | ❌ 存在泄漏 |
| Loss-Free (Ours) | ✅ 均衡 | ✅ 无 | ✅ 无泄漏 |
Loss-Free Balancing 在保持负载均衡的同时,既不引入干扰梯度,也不破坏因果约束,理论优势明显。
4. 实验
4.1 实验设置
模型架构:
- 采用 DeepSeekMoE 作为骨干网络,该架构将专家细粒度分割,并隔离部分专家作为共享专家。
- 主要实验使用 sigmoid作为门控函数(优于 softmax)。
- 模型参数量为 1B 和 3B。
训练数据:
- 包含网络文本、数学材料、代码脚本、出版文献等多语言语料。
- 词表大小为 32K,采用BPE分词器。
- 1B 模型训练 100B tokens,3B 模型训练 200B tokens。
基线方法:传统辅助损失控制方法,设置 α=0.001 以达到合理的性能-均衡折衷。
评估指标:
-
模型性能:验证集困惑度(Perplexity, PPL)。
-
负载均衡:最大违反度(MaxVio),定义如下:

其中,Load_i 为专家 i 分配到的 token 数,‾Load_i 为完美负载均衡下的期望负载。MaxVio 有全局和 batch 两种变体。
4.2 主要结果

实验结果表明:
-
在 1B 和 3B 模型上,Loss-Free Balancing 相比辅助损失方法取得了更低的验证困惑度(更好的模型性能)。
-
同时,Loss-Free Balancing 的全局负载均衡指标(MaxVio)显著优于辅助损失方法(约0.04 vs 0.52~0.72),几乎达到完美均衡。
-
从训练过程曲线来看,Loss-Free Balancing 在整个训练过程中持续保持优越的负载均衡。
这充分验证了 Loss-Free Balancing 打破了负载均衡与模型性能之间的困境。
4.3 偏置更新算法的实证研究

更新率 u:过低(0.0001)导致收敛慢,早期负载不均;过高(0.01)导致后期偏置震荡,损害均衡。最优值为 u=0.001。
更新规则:对比了符号更新 b_i = b_i + u⋅sign(e_i) 和数值更新 b_i = b_i + u⋅e_i。符号更新在保持相似负载均衡的同时,模型性能略优。


乘法偏置 vs 加法偏置:乘法偏置形式 s_{i,t}⋅b_i 的负载均衡效果相近,但模型性能略差于加法偏置 s_{i,t} + b_i。因此,加法偏置是更优选择。
5. 讨论
5.1 Loss-Free Balancing 与专家并行兼容
超大规模 MoE 模型通常采用专家并行( Expert Parallelism**)** 来降低显存需求。专家并行下,每个计算步骤包含 micro_batch_size * ep_data_parallel_size 个样本,称为计算batch(Computation Batch)。

实验表明:
-
随着计算batch大小的增加,Loss-Free Balancing 的负载均衡效果持续改善。
-
而辅助损失控制方法的负载均衡在计算batch较大时基本维持不变。
由于专家并行会显著增大计算batch,Loss-Free Balancing 在大规模 MoE 训练中具有天然优势,且其负载均衡优势会随着专家并行度的增加而进一步放大。
5.2 负载均衡与未来 Token 泄漏

对于因果语言模型,负载均衡方法必须遵守因果约束,避免未来 token 泄漏(Future Token Leakage)。
-
Auxiliary-controlled balancing 和 Loss-Free Balancing均遵守因果约束。
-
Expert Choice (EC) 则违反了因果约束。EC 通过确保每个专家分配完全相同的 token 数量来实现完美均衡,但这会导致未来 token 影响先前 token 的专家分配。
理论分析:对于一个稀疏比 R=K/N 的MoE层,EC 每个 token 最多泄漏比特信息:

对于一个 9 层MoE、16 个专家、每个 token 激活 2 个专家的模型,总泄漏量超过 50 比特,足以让每个 token 确定其后继 token 的身份。

实验证据:
-
将 EC 的 Top-K 选择块大小从 8192 减小到 512 时,观察到约 10% 的异常损失下降,表明存在泄漏。
-
在 Top-K 选择前对 token 进行打乱,异常损失下降得到缓解。
结论:未来 token 泄漏会破坏模型的泛化能力和可靠评估,因此 Loss-Free Balancing 相比 EC 在大规模 MoE 扩展中更为安全。
6. 结论
本文提出了 Loss-Free Balancing ,一种无需辅助损失的新型 MoE 负载均衡方法。该方法通过动态调整专家维度的偏置来影响路由决策,既不引入干扰梯度,也不破坏因果约束。在 1B 和 3B MoE 模型上的实验表明,Loss-Free Balancing 在模型性能 和负载均衡两方面均优于传统辅助损失控制方法,并天然兼容专家并行,避免了未来 token 泄漏问题。
Loss-Free Balancing为MoE模型的训练提供了一个简洁、高效且安全的负载均衡方案,有望在更大规模的MoE模型训练中发挥重要作用。
参考
https://spaces.ac.cn/archives/10699 (MoE 的几何意义)
https://spaces.ac.cn/archives/10735 (负载均衡)
https://spaces.ac.cn/archives/10757 (Loss-Free)
https://spaces.ac.cn/archives/10815 (动态激活 Loss-Free)
https://spaces.ac.cn/archives/10945 (Shared Expert、Fine-Grained Expert)
(2024|ACL|DeepSeek & 北大,MoE,细粒度专家分割,共享专家隔离)DeepSeekMoE:迈向极致专家专精的 MoE 语言模型
https://spaces.ac.cn/archives/11619 (Quantile Balancing, QB;交替迭代)
https://spaces.ac.cn/archives/11626 (动态专家激活 Quantile Balancing, QB-DA;单步生成)
【此处为简洁起见,私将其缩写为 QB-DA】
https://spaces.ac.cn/archives/11760 (序列级均衡-局部中心;Moving Quantile Balancing,MQB-DA;滑动窗)
https://spaces.ac.cn/archives/11782 (门控归一化的概率解释)
QB 算法与结果


QB‑Quantile‑Bias 关键问题整理

在浏览博客的时候,对动态激活有一个疑问:
- s - β > 0,就激活专家。从这里看,是动态激活数量。
- 但是,β 是从固定值 k 和 n 得来的分位数,也是固定的。所以最终激活的专家数是否应该也是固定的?
在博主的建议下,运行了博客里附带的代码,从数值方面加深了理解,整理如下。
Q1:quantile_bias 里,分位数是固定的吗?β 是固定的吗?
A: 分位数的位置 τ 是固定常数:τ = 1 - k / n;但是,β(分位数取出的数值)是动态变化的。
beta = np.quantile(s, 1 - k / n, axis=0)
- 1 - k / n:固定不变,由超参 k、n 决定,训练全程不会改。
- β:输出出来的 bias 数值(β)是动态的,每一批 s 打分分布一变,β 就跟着变。
举例子 n=256,k=8。则,τ = 1 - k / n = 1 - 8 / 256 = 0.96875
τ 永远等于 0.96875,不会变。 但是:
- 这一批专家 A 打分整体偏高 → 它的 0.96875 分位数数值就大 → β 变大,压低专家 A;
- 下一批专家 A 打分整体偏低 → 它的 0.96875 分位数数值就小 → β 变小,抬高专家 A。
"固定分位数" 指取哪个百分比位置固定;不是说 bias 的数值固定死。
Q2:bias 是动态的,被激活专家数是 k 还是动态?
A:动态激活模式(判断 s‑β>0 即激活,不是 top‑k) 只要满足 s‑β>0 就被激活,被激活的专家数不是固定的 k,而是动态变化的。
但值得注意的是,这些动态变化的专家数的统计期望依然是 k。即,有的输入 token 激活的专家数大于 k,而有的,可能是小于 k。
Q3:QB vs QB-DA 的区别是什么?两者都是通过 s-β>0 来选激活吗?
A: ❌ QB 并不是用 s-β>0 判断激活,二者激活规则完全不同。
1)QB-DA
beta = np.quantile(s, 1 - k / n, axis=0)
求解 β:仅按**列 (axis=0,专家维度)**一步分位数算出 β。
激活规则:
- 阈值路由 ,直接以 s-β>0 作为激活判断条件**,不做 top‑k**。
- β 直接充当每个专家的激活阈值。
数学保证:P(s_i - beta_i > 0) = k/n,每个专家打分超过 β 的样本占比严格等于 k/n。
样本激活数:每个样本激活专家数量随机,统计期望为 k,单样本可大于 / 小于 k。
2)QB
for _ in range(T):
alpha = np.quantile(s - beta, 1 - k / n, axis=1, keepdims=True) # 按样本行求分位数
beta = np.quantile(s - alpha, 1 - k / n, axis=0, keepdims=True) # 按专家列求分位数
求解 β:固定同样的 τ = 1 - k / n;按行 (axis=1)、列 (axis=0) 交替迭代 T 轮收敛得到 β。
激活规则:
- top‑k 选取 。循环内部的 s-beta 、s-alpha 只是迭代求解 bias 的中间数学运算;forward / 评估阶段不会判断 s-beta>0,直接对偏移后打分(s-beta)做 top‑k。
- β 只是打分 logit 偏移量,不是激活阈值。
数学保证:不再具备 P(s_i - beta_i > 0) = k/n;优化目标是:经过 bias 偏移后,top‑k 选中的各专家统计频次尽量逼近均衡 k/n。
样本激活数:每个样本严格选出 k 个专家,数量固定。
| 项目 | QB-DA | QB |
|---|---|---|
| 求解 β 方式 | 一步,仅 axis=0 |
T 轮交替:axis=1 ↔ axis=0 |
| Forward 激活规则 | 阈值路由:s-beta>0 | Top‑k 选取 |
| 是否用 s-beta>0 做激活 | ✅ 是,真实激活判断 | ❌ 否;仅求解 bias 的中间计算 |
| 数学等式 | P(s_i - beta_i > 0) = k/n | 无该等式;追求 top‑k 选中频次均衡 |
| 单样本激活专家数 | 随机,期望 k | 严格固定 k |
MQB 算法与结果


其中,
- 第 3 行,t 为整数,取值 0, b-1,b 为分桶的数目。
- 第 4 行,+1/2 取桶 m 的中心点,/b 后归一化到 0,1,与 s 数值范围一致。

