Mixture-of-Experts的无辅助损失负载均衡策略
Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
AI总结：
针对MoE模型中辅助损失负载均衡引发的干扰梯度问题，提出无辅助损失的Loss-Free Balancing策略，通过动态调整专家偏置实现负载均衡，在3B参数模型上验证了其优于传统方法的性能与均衡效果。
摘要：
对于Mixture-of-Experts (MoE)模型，专家负载不均衡会导致路由崩溃或计算开销增加。现有方法通常采用辅助损失来促进负载均衡，但较大的辅助损失会在训练中引入不可忽视的干扰梯度，从而损害模型性能。为了在训练时控制负载均衡且不产生不良梯度，我们提出了Loss-Free Balancing，其特征在于一种无辅助损失的负载均衡策略。具体而言，在top-K路由决策之前，Loss-Free Balancing会首先对每个专家的路由得分应用专家级偏置。通过根据每个专家的近期负载动态更新偏置，Loss-Free Balancing能够持续保持专家负载的均衡分布。此外，由于Loss-Free Balancing不产生任何干扰梯度，它还提升了从MoE训练中获得的模型性能上限。我们在参数量高达3B、训练数据量高达200B tokens的MoE模型上验证了Loss-Free Balancing的性能。实验结果表明，与传统的辅助损失控制的负载均衡策略相比，Loss-Free Balancing实现了更好的性能和更好的负载均衡。
英文摘要：
For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
