用ReLU路由实现可微的专家混合模型，提升计算效率与扩展性
ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
- 以ReLU替代TopK+Softmax，实现完全可微的专家路由机制
- 在不同模型规模下均超越传统MoE，专家数量增加时性能更优
- 适合需要高效扩展、动态计算分配的大型模型场景
稀疏激活的专家混合（MoE）模型被广泛用于在不增加计算开销的前提下扩大模型容量。然而，传统的TopK路由机制训练过程不连续且不可微，限制了其性能和可扩展性。为此，我们提出ReMoE，一种完全可微的MoE架构，以ReLU作为路由器，可作为传统TopK+Softmax路由的即插即用替代方案。我们还提出了调控路由器稀疏性并平衡专家负载的方法。ReMoE的连续特性支持对令牌和层的高效动态计算分配，同时表现出领域专属性。实验表明，ReMoE在多种模型规模、专家数量和粒度级别下均持续优于传统TopK路由的MoE。此外，随着专家数量增加，ReMoE展现出更优的可扩展性。基于Megatron-LM的实现已开源：https://github.com/thu-ml/ReMoE。
原文摘要 · Abstract (English)
Sparsely activated Mixture-of-Experts (MoE) models are widely adopted to scale up model capacity without increasing the computation budget. However, vanilla TopK routers are trained in a discontinuous, non-differentiable way, limiting their performance and scalability. To address this issue, we propose ReMoE, a fully differentiable MoE architecture that offers a simple yet effective drop-in replacement for the conventional TopK+Softmax routing, utilizing ReLU as the router instead. We further propose methods to regulate the router's sparsity while balancing the load among experts. ReMoE's continuous nature enables efficient dynamic allocation of computation across tokens and layers, while also exhibiting domain specialization. Our experiments demonstrate that ReMoE consistently outperforms vanilla TopK-routed MoE across various model sizes, expert counts, and levels of granularity. Furthermore, ReMoE exhibits superior scalability with respect to the number of experts, surpassing traditional MoE architectures. The implementation based on Megatron-LM is available at https://github.com/thu-ml/ReMoE.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。