ReMoE: 基于ReLU路由的完全可微混合专家模型
机器学习
2025-02-28 v2
摘要
稀疏激活的混合专家(MoE)模型被广泛用于在不增加计算预算的情况下扩展模型容量。然而，传统的TopK路由是以不连续、不可微的方式训练的，这限制了其性能和可扩展性。为了解决这个问题，我们提出了ReMoE，一种完全可微的MoE架构，它作为传统TopK+Softmax路由的简单而有效的替代方案，利用ReLU作为路由器。我们还提出了调节路由器稀疏性同时平衡专家间负载的方法。ReMoE的连续性使其能够在token和层之间高效地动态分配计算，同时展现出领域专业化。我们的实验表明，在各种模型大小、专家数量和粒度下，ReMoE始终优于传统的TopK路由MoE。此外，ReMoE在专家数量方面表现出卓越的可扩展性，超越了传统的MoE架构。基于Megatron-LM的实现可在https://github.com/thu-ml/ReMoE获取。
关键词
引用
@article{arxiv.2412.14711,
title = {ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing},
author = {Ziteng Wang and Jun Zhu and Jianfei Chen},
journal= {arXiv preprint arXiv:2412.14711},
year = {2025}
}
备注
ICLR 2025