基本信息

• 📝 原文链接: https://arxiv.org/abs/2412.14711
• 👥 作者: Ziteng Wang, Jianfei Chen, Jun Zhu
• 🏷️ 关键词: Mixture-of-Experts, ReLU routing
• 📚 分类: 机器学习

摘要

中文摘要

稀疏激活的专家混合模型（MoE）被广泛采用以扩大模型容量而不增加计算预算。然而，传统的TopK路由器以不连续、非可微的方式训练，这限制了它们的性能和可扩展性。为了解决这一问题，我们提出了ReMoE，这是一种全可微的MoE架构，它提供了一个简单而有效的替代方案，用于传统的TopK+Softmax路由，使用ReLU作为路由器。我们进一步提出了调节路由器稀疏性的方法，同时在专家之间平衡负载。ReMoE的连续性使其能够高效地在标记和层之间动态分配计算，同时也表现出领域专长。我们的实验表明，ReMoE在各种模型大小、专家数量和粒度级别上始终优于传统的TopK路由MoE。此外，ReMoE在专家数量的可扩展性方面表现出优越性，超过了传统的MoE架构。基于Megatron-LM的ReMoE实现可在以下网址获取：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.

论文解读

一句话总结

本文提出了一种名为ReMoE的混合专家模型，通过ReLU路由代替传统的TopK路由，实现了可微分混合专家模型，提高了模型性能和可扩展性。

问题1：这篇论文想要解决什么具体问题？

• 问题背景：稀疏激活的混合专家（MoE）模型被广泛用于扩大模型容量而不增加计算预算。然而，传统的TopK路由器以不连续和非可微分的方式进行训练，限制了其性能和可扩展性。
• 现有方案不足：传统的TopK路由器存在训练过程中的不连续性和非可微性，这限制了其性能和可扩展性。
• 研究目标：提出一种名为ReMoE的MoE架构，利用ReLU作为路由器，实现可微分的MoE模型，提高模型性能和可扩展性。

问题2：论文的核心创新点是什么？

• 技术创新：提出了一种基于ReLU的路由方法，用于选择激活的专家，实现了可微分的MoE模型。
• 方法改进：引入了自适应L1正则化，以控制路由器的稀疏性，同时平衡专家之间的负载。
• 优势：与传统的TopK路由器相比，ReMoE具有连续性和可微性，可以更有效地分配计算资源，并表现出更强的领域专业化。

问题3：实验结果如何验证了方法的有效性？

• 关键实验：在LLaMA架构上进行的实验，比较了ReMoE与TopK路由器和其他路由方法的性能。
• 性能提升：ReMoE在各种模型大小、专家数量和粒度级别上均优于传统的TopK路由器。
• 对比结果：ReMoE在零样本下游任务上的性能优于其他方法，包括TopK路由器、全可微分的Lory和稀疏Mixer-v2。

问题4：这个研究的实际应用价值是什么？

• 应用场景：ReMoE适用于需要扩大模型容量而不增加计算预算的领域，如自然语言处理和计算机视觉。
• 实施建议：ReMoE可以通过Megatron-LM进行实现，并支持多种模型并行形式。
• 局限与展望：ReMoE在训练初期需要更多的计算资源，但随着专家数量的增加，其性能提升更为显著。未来可以探索更有效的负载平衡技术和更高效的训练方法。