ReMoE：具有ReLU路由的全可微混合专家模型
ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
AI总结：
针对传统TopK路由不可微导致性能受限的问题，提出全可微MoE架构ReMoE，以ReLU替代TopK+Softmax路由，实现高效动态计算分配，实验表明其在多种规模下优于传统MoE并具更优可扩展性。
摘要：
稀疏激活的混合专家（MoE）模型被广泛采用，以在不增加计算预算的情况下扩展模型容量。然而，传统的TopK路由器以不连续、不可微的方式进行训练，限制了其性能和可扩展性。为解决这一问题，我们提出了ReMoE，一种全可微的MoE架构，它为传统的TopK+Softmax路由提供了一种简单而有效的替代方案，使用ReLU作为路由器。我们进一步提出了调节路由器稀疏性并平衡专家间负载的方法。ReMoE的连续性使其能够在令牌和层之间高效动态分配计算，同时展现出领域专门化。我们的实验表明，ReMoE在各种模型规模、专家数量和粒度级别上始终优于传统的TopK路由MoE。此外，ReMoE在专家数量方面表现出卓越的可扩展性，超越了传统MoE架构。基于Megatron-LM的实现可在https://github.com/thu-ml/ReMoE获取。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
