title: ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
summary: 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.
paper_authors._id: 676a25362d7ae887c4f20b6d
paper_authors.name: Ziteng Wang
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-12-30T19:35:36.994Z
paper_authors._id: 6151815f33293feca12a6d44
paper_authors.avatarUrl: /avatars/e92c8de879c542d2f843b742cd204065.svg
paper_authors.fullname: Ziteng Wang
paper_authors.type: user
paper_authors.user: thuwzt
paper_authors._id: 676a25362d7ae887c4f20b6e
paper_authors.name: Jianfei Chen
paper_authors._id: 676a25362d7ae887c4f20b6f
paper_authors.name: Jun Zhu
paper_title: ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
paper_summary: 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.
paper_ai_keywords: Mixture-of-Experts (MoE) TopK routers ReMoE differentiable MoE architecture conventional TopK+Softmax routing ReLU router sparsity load balancing among experts dynamic allocation of computation domain specialization Megatron-LM