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

2024-12-19

arXiv.org

AI Summary

In conventional sparse Mixture-of-Experts (MoE) models, the Top-K routing mechanism is non-differentiable, impeding end-to-end optimization and limiting scalability. This paper proposes the first fully differentiable MoE architecture, replacing the discrete Top-K + Softmax routing with a continuous, ReLU-based differentiable routing mechanism. The method incorporates sparsity regularization and load-balancing constraints to ensure controlled expert activation and equitable expert utilization, while enabling dynamic expert selection and domain specialization. Extensive experiments across varying model scales, numbers of experts, and granularity levels demonstrate consistent and significant improvements over baseline Top-K MoE—particularly in expert scalability and downstream task performance. The architecture maintains high sparsity without sacrificing trainability or inference efficiency. Code is publicly available.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Mixed Discrete/Continuous Search
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Differentiable MoE architecture
• ReLU routing mechanism
• Scalable expert load balancing

Innovation

• Fully differentiable MoE architecture
• Utilizes ReLU for routing
• Regulates router sparsity effectively