ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
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.
676a25362d7ae887c4f20b6d
Ziteng Wang
claimed_verified
2024-12-30T19:35:36.994Z
6151815f33293feca12a6d44
/avatars/e92c8de879c542d2f843b742cd204065.svg
Ziteng Wang
user
thuwzt
676a25362d7ae887c4f20b6e
Jianfei Chen
676a25362d7ae887c4f20b6f
Jun Zhu
ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
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, a differentiable MoE architecture using ReLU-based routing, outperforms traditional TopK routing in terms of scalability and performance across various models and expert counts.
Mixture-of-Experts (MoE) TopK routers ReMoE ReLU sparsity regulation domain specialization Megatron-LM