（2024|DeepSeek 北大，Loss-Free Balancing，每专家偏置，QB）MoE 的无辅助损失负载均衡策略

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. 引言

2. 背景

2.1 MoE

2.2 用于负载均衡的辅助损失

3. 无辅助损失的负载均衡策略

4. 实验

4.1 实验设置

4.2 主要结果

4.3 偏置更新算法的实证研究

5. 讨论

5.1 Loss-Free Balancing 与专家并行兼容

5.2 负载均衡与未来 Token 泄漏

6. 结论

参考

QB 算法与结果

QB‑Quantile‑Bias 关键问题整理

Q1：quantile_bias 里，分位数是固定的吗？β 是固定的吗？

Q2：bias 是动态的，被激活专家数是 k 还是动态？

Q3：QB vs QB-DA 的区别是什么？两者都是通过 s-β>0 来选激活吗？

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 是符号函数。

与其他方法的对比：下表总结了不同负载均衡方法的特性（绿色表示好，红色表示差）：

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 个专家，数量固定。

MQB 算法与结果

其中，

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