ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing

稀疏激活的专家混合（Mixture-of-Experts, MoE）模型被广泛采用，在不增加计算预算的情况下扩展模型容量。然而，传统的TopK路由是以不连续且不可微的方式训练的，这限制了其性能和可扩展性。为了解决这一问题，我们提出了ReMoE，这是一种完全可微的MoE架构，提供了一个简单而有效的替代方案，用于替换传统的TopK+Softmax路由，使用ReLU作为路由机制。我们进一步提出了一些方法来调节路由的稀疏性，同时在各专家之间平衡负载。ReMoE的连续特性使得能够在令牌和层之间高效地动态分配计算资源，并表现出领域专业化的特点。我们的实验表明，ReMoE在各种模型大小、专家数量和粒度级别上始终优于传统的TopK路由MoE。此外，ReMoE在专家数量方面的可扩展性更优，超越了传统的MoE架构。基于Megatron-LM的实现可在https://github.com/thu-ml/ReMoE获取。

论文试图解决Mixture-of-Experts (MoE) 模型中传统的TopK路由器训练方式不连续、非可微的问题，这限制了其性能和可扩展性。这是一个现有问题，但通过提出新的解决方案来改进它。

关键思路是提出ReMoE，一种完全可微的MoE架构，用ReLU作为路由器替代传统的TopK+Softmax路由方法。此外，还提出了调节路由器稀疏性和平衡专家负载的方法。相比当前研究，ReMoE的连续特性允许更有效的动态计算资源分配，并展示了领域专业化的能力。

该论文值得关注的地方包括：1) ReMoE在各种模型大小、专家数量和粒度水平上均优于传统的TopK路由MoE；2) 实验设计验证了ReMoE在不同配置下的优越性能和可扩展性；3) 使用Megatron-LM框架实现了开源代码，方便复现和进一步研究；4) 提出的方法为未来研究提供了新方向，如更复杂的路由器设计和优化算法。

最近在这个领域的相关研究包括：1) 'Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity' - 提出了Switch Transformer，使用稀疏激活机制扩大模型规模；2) 'GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding' - 探讨了通过条件计算和自动分片来扩展大型模型；3) 'Rethinking the Value of Network Pruning' - 重新思考网络剪枝的价值，与本文提出的稀疏激活机制有相似之处。