(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. 引言](#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 balancingLoss-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-betas-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=1axis=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 数值范围一致。
相关推荐
2601_9667990418 小时前
酷嗨米J300:一台相机同步推三平台
运维·服务器·负载均衡
小谈不敲代码21 小时前
Higress 生产级调优手册
云原生·kubernetes·gateway·负载均衡·higress
桃西西呀1 天前
GPT-5.6 的 ultra 模式凭什么开 4 个 Agent 并行跑?——多智能体协作,是把"一个聪明人"换成"一个团队"
人工智能·llm·agent
澄怀1 天前
Agent 说「我已经改好文件了」,这句话到底能不能信?
llm·agent
天地会珠海分舵1 天前
Qwen 3.8‑Flash完整实测!从推理编程到多模态!真的能打吗?
大语言模型·测评报告·qwen 3.8 flash
夫唯不争,故无尤也1 天前
On-Policy Distillation(OPD)和reinforcement (RL)结合
llm·agent·强化学习·rl·opd
龍德明宇1 天前
四个时代之问的存在论根基-龍德明宇
大语言模型·ai哲学·负主体性