Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts论文地址https://arxiv.org/abs/2408.15664ICLR 2025https://openreview.net/forum?idy1iU5czYpE学术交流群922230617目录1. 引言2. 背景2.1 MoE2.2 用于负载均衡的辅助损失3. 无辅助损失的负载均衡策略4. 实验4.1 实验设置4.2 主要结果4.3 偏置更新算法的实证研究5. 讨论5.1 Loss-Free Balancing 与专家并行兼容5.2 负载均衡与未来 Token 泄漏6. 结论参考QB 算法与结果QB‑Quantile‑Bias 关键问题整理Q1quantile_bias 里分位数是固定的吗β 是固定的吗Q2bias 是动态的被激活专家数是 k 还是动态Q3QB 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 是符号函数。与其他方法的对比下表总结了不同负载均衡方法的特性绿色表示好红色表示差负载均衡方法专家负载均衡干扰梯度未来 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 tokens3B 模型训练 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导致后期偏置震荡损害均衡。最优值为 u0.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个样本称为计算batchComputation Batch。实验表明随着计算batch大小的增加Loss-Free Balancing 的负载均衡效果持续改善。而辅助损失控制方法的负载均衡在计算batch较大时基本维持不变。由于专家并行会显著增大计算batchLoss-Free Balancing 在大规模 MoE 训练中具有天然优势且其负载均衡优势会随着专家并行度的增加而进一步放大。5.2 负载均衡与未来 Token 泄漏对于因果语言模型负载均衡方法必须遵守因果约束避免未来 token 泄漏Future Token Leakage。Auxiliary-controlled balancing和Loss-Free Balancing均遵守因果约束。Expert Choice (EC)则违反了因果约束。EC 通过确保每个专家分配完全相同的 token 数量来实现完美均衡但这会导致未来 token 影响先前 token 的专家分配。理论分析对于一个稀疏比 RK/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 BalancingMQB-DA滑动窗)https://spaces.ac.cn/archives/11782 (门控归一化的概率解释)QB 算法与结果QB‑Quantile‑Bias 关键问题整理在浏览博客的时候对动态激活有一个疑问s - β 0就激活专家。从这里看是动态激活数量。但是β 是从固定值 k 和 n 得来的分位数也是固定的。所以最终激活的专家数是否应该也是固定的在博主的建议下运行了博客里附带的代码从数值方面加深了理解整理如下。Q1quantile_bias里分位数是固定的吗β 是固定的吗A 分位数的位置 τ 是固定常数τ 1 - k / n但是β分位数取出的数值是动态变化的。beta np.quantile(s, 1 - k / n, axis0)1 - k / n固定不变由超参 k、n 决定训练全程不会改。β输出出来的 bias 数值β是动态的每一批 s 打分分布一变β 就跟着变。举例子 n256k8。则τ 1 - k / n 1 - 8 / 256 0.96875τ 永远等于 0.96875不会变。 但是这一批专家 A 打分整体偏高 → 它的 0.96875 分位数数值就大 → β 变大压低专家 A下一批专家 A 打分整体偏低 → 它的 0.96875 分位数数值就小 → β 变小抬高专家 A。“固定分位数” 指取哪个百分比位置固定不是说 bias 的数值固定死。Q2bias 是动态的被激活专家数是 k 还是动态A动态激活模式判断s‑β0即激活不是 top‑k只要满足s‑β0就被激活被激活的专家数不是固定的 k而是动态变化的。但值得注意的是这些动态变化的专家数的统计期望依然是 k。即有的输入 token 激活的专家数大于 k而有的可能是小于 k。Q3QB vs QB-DA 的区别是什么两者都是通过 s-β0 来选激活吗A ❌QB 并不是用 s-β0 判断激活二者激活规则完全不同。1QB-DAbeta np.quantile(s, 1 - k / n, axis0)求解 β仅按列 (axis0专家维度)一步分位数算出 β。激活规则阈值路由直接以 s-β0 作为激活判断条件不做 top‑k。β 直接充当每个专家的激活阈值。数学保证P(s_i - beta_i 0) k/n每个专家打分超过 β 的样本占比严格等于 k/n。样本激活数每个样本激活专家数量随机统计期望为 k单样本可大于 / 小于 k。2QBfor _ in range(T): alpha np.quantile(s - beta, 1 - k / n, axis1, keepdimsTrue) # 按样本行求分位数 beta np.quantile(s - alpha, 1 - k / n, axis0, keepdimsTrue) # 按专家列求分位数求解 β固定同样的 τ 1 - k / n按行 (axis1)、列 (axis0) 交替迭代 T 轮收敛得到 β。激活规则top‑k 选取。循环内部的s-beta、s-alpha只是迭代求解 bias 的中间数学运算forward / 评估阶段不会判断 s-beta0直接对偏移后打分s-beta做 top‑k。β 只是打分 logit 偏移量不是激活阈值。数学保证不再具备 P(s_i - beta_i 0) k/n优化目标是经过 bias 偏移后top‑k 选中的各专家统计频次尽量逼近均衡 k/n。样本激活数每个样本严格选出 k 个专家数量固定。项目QB-DAQB求解 β 方式一步仅axis0T 轮交替axis1↔axis0Forward 激活规则阈值路由s-beta0Top‑k 选取是否用 s-beta0 做激活✅ 是真实激活判断❌ 否仅求解 bias 的中间计算数学等式P(s_i - beta_i 0) k/n无该等式追求 top‑k 选中频次均衡单样本激活专家数随机期望 k严格固定 kMQB 算法与结果其中第 3 行t 为整数取值 [0, b-1]b 为分桶的数目。第 4 行1/2 取桶 m 的中心点/b 后归一化到 [0,1]与 s 数值范围一致。