无辅助损失的 Mixture-of-Experts 负载均衡策略
机器学习
2024-08-29 v1 计算与语言
摘要
对于 Mixture-of-Experts（MoE）模型，专家负载不均会导致路由崩溃或增加计算开销。现有方法通常采用辅助损失来鼓励负载均衡，但大型辅助损失会引入显著的干扰梯度，从而影响模型性能。为在控制负载均衡的同时避免产生不必要的梯度，本文提出一种无辅助损失的负载均衡策略。具体而言，在进行 top-K 路由决策之前，Loss-Free Balancing 首先对每个专家的路由得分施加专家级偏置。通过根据每个专家的最近负载动态更新偏置，Loss-Free Balancing 可以持续维持专家负载的平衡分布。此外，由于 Loss-Free Balancing 不产生任何干扰梯度，因此也提高了从 MoE 训练中获得的模型性能上限。我们在参数最高达 30 亿、训练数据最高达 2000 亿 token 的 MoE 模型上验证了 Loss-Free Balancing 的性能。实验结果表明，Loss-Free Balancing 在保持更好性能的同时，实现了更好的负载均衡，优于传统的基于辅助损失控制的负载均衡策略。
引用
@article{arxiv.2408.15664,
title = {Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts},
author = {Lean Wang and Huazuo Gao and Chenggang Zhao and Xu Sun and Damai Dai},
journal= {arXiv preprint arXiv:2408.15664},
year = {2024}
}