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

Authors: Ziteng Wang, Jun Zhu, Jianfei Chen
Primary area: foundation or frontier models, including LLMs
Keywords: Mixture-of-Experts, Differentiable Routing, Sparsity

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730678126032

thread_type:
official_review

<thread_object>
creator: reviewer_rqkz

<content>
summary:
This paper proposes a new routing method for MoE routing which addresses issues related to discontinuity inherent in TopK routing while still maintaining sparsity. The key idea is to replace TopK softmax with a ReLU routing function. This leads to a natural selection of experts from the non-zero routing scores. A sparse selection of experts is achieved via L1-regularization which evolves dynamically throughout training in order to eventually reach a target sparsity. The L1-regularization can also be combined with a load balancing loss to ensure both sparse selection of experts as well as an even distribution of token routing.

strengths:
Mixture-of-experts models are a powerful new paradigm which can unlock favorable scaling. However, training such models is difficult, in large part due to the discrete nature of the TopK expert routing. The proposed routing method alleviates this difficult which enables better capabilities in MoE models. The nature of this solution is conceptually clean and gives good results. The paper is well-written and provides a useful empirical analysis.

weaknesses:
It would be helpful to delve a bit more into the conceptual intuitions of the regularization penalty, especially with the additional load balancing and the dynamic penalty adjustment. This will make it easier for the reader to grasp the key conceptual contribution. Additionally, it would be helpful to shed some light on exactly what issue is being overcome with the new strategy.

questions:
Why is TopK routing struggling beyond just "being discontinuous" and how is the new strategy overcoming this? Is it really the continuity of the routing or the more flexible nature of the sparse regularization which allows for multiple "phases" of training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730685287995

thread_type:
official_review

<thread_object>
creator: reviewer_yjh9

<content>
summary:
Traditional Mixture of Experts (MoE) models use Top-K routing to select experts, which is a non-differentiable process. In this work, the authors propose using ReLU activation combined with load balancing loss and sparsity loss to address this issue. Several validation experiments are conducted to demonstrate the effectiveness of the proposed approach.

strengths:
- The authors propose a simple, intuitive, and straightforward alternative method for MoE routing, which can be integrated into existing MoE structures via a drop-in approach. 
- The submission clearly illustrates the background information and the proposed method. 
- The performance gain by replacing the traditional routing with ReLU is interesting.
- The submission provides many insightful observation, for instance, the analysis on the correlation between expert allocation and token frequency, and the effect of load balance loss.

weaknesses:
- The experiment scale is relatively small, and only perplexity is reported. I'm wondering about the performance comparison on downstream tasks.
- The performance improvement seems significant, but it would help if the authors can provide more analysis and experiment/theoritical evidence on why the fully differentiable router will help.

questions:
- What does $N$ refers to? If I understand correctly, $N$ refers to the number of activated parameters. In that case, the largest model should have about 7/8 billion parameters in total, right? I suggest the authors to explicitly mention the important hyper-parameter like the model size.
- The Figure 7 is really interesting and insightful. But I'm wondering whether the router will assign different number of experts solely based on token ids (or token frequency), or it will also capture some semantic informaton (i.e. difficulty of tasks).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691747155

thread_type:
official_review

<thread_object>
creator: reviewer_urhd

<content>
summary:
The paper introduce ReMoE, a differentiable MoE architecture that incorporates ReLU routing in replacement for TopK routing for Mixture-of-Experts. This paper further propose methods to regulate the router’s sparsity while balancing the load among experts. ReMoE enables efficient dynamic allocation of computation across tokens and layers, exhibiting domain specialization. The paper observes a natural three-stage training process in ReMoE: warm-up stage, sparsifying stage, and stable stage offering interesting insights. Perplexity based experiments demonstrate that ReMoE consistently outperforms vanilla TopK-routed MoE across various model sizes, expert counts, and levels of granularity.

strengths:
1. The paper is well-written and authors have presented some interesting ablations like domain specialization of experts, assign a higher number of experts to rarer tokens, etc.

2. Experimental results presented looks promising. It is interesting to see how a relatively simple idea work so much better. 

3. Three-stage training that naturally occurs during REMoE training is an interesting section. I recommend author to add the loss plots too in the figure to draw parallels.

weaknesses:
While there are many interesting experiments and ablation in the paper, I have several comments.

1. The authors have failed to provide the training cost of different baselines used in the experiments and also for inference. Clearly, ReMoE stage I training activates a lot more experts during initial stage. Some speed analysis is must to clearly identify the computation overhead of dynamic expert selection both from inference and training perspective. A more detailed analysis of the trade-offs between performance gains and computational costs would be beneficial.

2. Complete evaluation setup of paper is centered around training and validation loss, and perplexity which is completely not reliable to justify performance. To illustrate the performance benefit of ReMoE, it is critically important to show results translating to downstream tasks.

3. TopK routing facilitates better parallelization opportunities while scaling MoEs. What are the authors thoughts on adopting a dynamic expert selection which can lead to poor GPU utilization.

4. The authors introduces two new hyperparameters $\lambda_{0}$ and $\alpha$. MoE training is well-known to be tricky. I am concerned how these hyperparameters will further increase challenges. The authors have indeed provided some ablation on 20k
steps (∼10B tokens) for 182M token, but it is not clear how the observations replicate in large MoEs. In addition, even for this small-scale experiment, some $\alpha$ values lead to rapid regularization changes and excessive oscillation. 

5. While the paper compares ReMoE with several baselines, it would be beneficial to include comparisons with other recent MoE variants, such as those addressing the routing bottleneck or improving expert utilization.

6. Domain specialization experiments are interesting. I would recommend authors to conduct some data domain vocabulary coverage ratio related experiments for each experts to complete the full-picture.

I am looking forward to additional experiments in rebuttal to update my rating.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701697791

thread_type:
official_review

<thread_object>
creator: reviewer_qrqt

<content>
summary:
The paper proposes a new routing method for MoE transformers and provides an accompanying training recipe. The routing method (ReMoE) replaces softmax + TopK expert selection with a ReLU activation function. This results in a fully differentiable MoE transformer, which requires a new loss penalty to encourage a balanced load and a reasonable number of active experts ($k$). The authors empirically evaluate their new method in the context of autoregressive language modeling. Experiments include a study of performance w.r.t loss at different model sizes, a varied number of experts, and different expert granularity; an analysis of different training stages; and an analysis of routing behavior.

strengths:
- The paper is generally well-written and easy to understand. In particular, the method section is thorough and gives a comprehensive explanation of the modifications made to the standard MoE training pipeline. I quite like Figure 4 as it gives a precise picture of the training dynamics specific to ReMoEs.
- The proposed method outperforms all other methods compared to in the study.
- Experiments in Figure 6 studying performance improvements for different numbers of active parameters, different expert granularity, and a varying number of experts demonstrate that the proposed ReLU MoEs along with their training algorithm consistently outperforms a standard TopK MoE.
- Using the ReLU activation function as a replacement to TopK + softmax for MoE routing is a novel and potentially interesting idea. The method could allow for improved conditional computation as there is no hard requirement for each token to be processed by exactly $k$ experts.

weaknesses:
My **greatest concern** regards the attribution of the success of the method. As the paper is currently written, it suggests that using the ReLU activation function in place of TopK + softmax is the cause of the improved performance because it makes the routing function *differentiable*. However, the training algorithm is also changed. Notably, during the first ~100 steps of training the ReMoE has sparsity as low as 20% (Figure 4 (a)), requiring substantially more memory and computational cost for these first 100 steps. This leads me to ask the following question:

*Is the success of the method due to the use of ReLU or is it due to the expensive near-dense training of the MoE for the first 100 steps?*

A simple way to address this weakness would be to provide a dMoE baseline that trains with k=int( 0.8 * E ) (e.g., nearly dense) for the first 100 steps of training and switches to k=1 thereafter. 


**Other weaknesses**
- While the ReMoE method allows for more flexible routing, it could unevenly distribute the load across a sequence, causing latency issues during autoregressive generation. How does the allocation of compute vary across the sequence of tokens?
- In the introduction, you state "the vanilla TopK router introduces a discrete and nondifferentiable training objective (Shazeer et al., 2017; Zoph et al., 2022), limiting the performance". This is false. The training objective (e.g. auxiliary loss), itself, is differentiable. The difficulty is related to receiving a sparse gradient (e.g., a gradient for only a subset of activated experts).
- Recent relevant works [1-3] are not mentioned in the related work section and not compared to in the main manuscript. Specifically, the sparse-mixer-v2 method, which improves the router's gradient estimate would be a relevant baseline to compare with.
- I miss an evaluation of the performance of ReMoEs on LM evaluation benchmarks.
- A batch size of 512k tokens is small for LLM pre-training. 
- While authors claim to train on a compute-optimal number of tokens, is 30B compute optimal for all models in the study? Many recent LLMs train well beyond the compute-optimal regime (e.g. Llama3 8B was trained for more than 50x compute optimal). Do the ReMoE results hold for longer training?
- ReMoEs will result in high memory consumption early on in training. This is reasonable for smaller models, but can quickly become expensive for very sparse MoEs. This should be explicitly noted in section 3.5 line 269.





[1] SPARSE BACKPROPAGATION FOR MOE TRAINING, https://arxiv.org/pdf/2310.00811

[2] GRIN: GRadient-INformed MoE https://arxiv.org/abs/2409.12136

[3] Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models https://arxiv.org/pdf/2404.05567



**I would be happy to raise my score if some concerns are addressed.**

questions:
One of the benefits of ReLU routing is that there is no inherent requirement for each token to be processed by $k$ experts. 

- Can ReMoE leverage the flexibility it is said to have with the regularization terms used?





**Suggested writing changes:**
- Scaling in parameters N --> Scaling in active parameters N
</content>
</thread_object>

</thread>
<thread>
creation date:
1730708131819

thread_type:
official_review

<thread_object>
creator: reviewer_vcwp

<content>
summary:
This work proposed a new MoE architecture that replace the TopK with ReLU in the routing modules. The ReLU-MoE is then optimized in a three stage training framework, experiments demonstrates the effectiveness of the proposed method across language modeling on different model sizes.

strengths:
- Various ablation studies and visualization are conducted, analysing the effectivenss of the proposed methods.

- The method is clearly explained and easy to follow.

weaknesses:
- The novelty of ReMoE is modest, as it simply replaces the TopK operation with ReLU.

- The perplexity improvements over dMoE are not significantly. whether downstream performance, such as on common-sense reasoning tasks, would shows significant enhancement?

- In practical MoE implementations, an all-to-all dispatch and combine approach can be used to assign each token to the appropriate experts, thereby reducing memory usage and computational load. However, during the initial stage of training, the absence of high sparsity can impact memory and computation efficiency, as this phase effectively requires training a nearly dense model.

questions:
- Why ReLU based router can enhance the domain specialization of MoE?

- Whether the ReLU-based router can be further enhanced with shifted-ReLU, even with a learnable threshold?

- How does the stage I/II affects the training performance?
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces a novel approach to address key limitations of traditional TopK-routed MoE architectures by replacing the discrete TopK routing mechanism with a ReLU-based differentiable router. The authors present a compelling case for the ReLU router, highlighting its ability to ensure smooth gradients, dynamic expert allocation, and efficient load balancing. Their experiments demonstrate significant improvements over traditional MoE architectures, both in terms of perplexity and downstream task performance, across various model sizes, expert counts, and token granularities. This work also introduces practical enhancements, such as a three-stage training process that stabilizes the ReLU routing, which, despite its simplicity, yields promising results. The authors provide extensive empirical validation, including ablation studies and comparisons with state-of-the-art baselines like SparseMixer-v2.

While some reviewers noted that the novelty of replacing TopK with ReLU might seem modest at first glance, the authors convincingly argue that the continuous nature of ReLU enables dynamic computation scaling and a steeper scalability curve for expert counts, making it a meaningful advancement for MoE models. The paper’s clear presentation and rigorous rebuttal process, where all major concerns were satisfactorily addressed, reflect its readiness for publication.