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

We greatly thank the reviewer for recognizing our method as conceptually clean and for acknowledging that it produces good results. We sincerely hope that our responses will further clarify our contributions and address any remaining uncertainties. Below, we provide detailed replies to the reviewer’s questions:

Q1: Conceptual Intuitions on Regularization

Thank you for the thoughtful advice! We will clarify the conceptual intuitions for $L_1$ regularization, load balancing, and dynamic adjustment as follows:

• $L_1$ Regularization: This introduces a force to each non-zero router output by adding $\frac{\lambda_i}{LT}$ to their gradients, driving the outputs toward zero (inactive) and enhancing sparsity.
• Load Balancing: By adding a weight $f_{l,e}$ to the regularization, the added gradient of non-zero router outputs is modified to $\frac{f_{l,e} \lambda_i}{L T}$, which depends for each expert. This penalizes over-utilized experts by driving their outputs toward zero more quickly.
• Dynamic Adjustment: The coefficient $\lambda_i$ dynamically adjusts the magnitude of the regularization force applied to router outputs, ensuring the model stabilizes at the target sparsity level.

We have incorporated these conceptual intuitions into Sections 3.3 and 3.4 to clarify the narrative further.

Q2: What issue is being overcome with ReLU routing

The superiority of ReLU over TopK can be summarized into the following two aspects:

• Continuity and differentiability: ReLU is fully differentiable and stays continuous when the active experts change, where TopK fails.
• Dynamic expert allocation: ReLU fertilizes dynamic expert allocation, so that tokens are not necessarily routed to $k$ experts and more computations are allocated to rarer tokens.

We have modified the Conclusion Section to highlight the above statements.

Q3: Why is TopK routing struggling beyond just "being discontinuous" and how is ReLU overcoming it?

• Theoretically, the discontinuity in TopK causes the router to receive false gradient when active experts change. For example, consider a simple MoE $f(x)=\tilde\theta_1f_1(x)+\tilde\theta_2f_2(x)$, where $\tilde{\theta_1},\tilde{\theta_2}=\text{Top1}(\theta_1,\theta_2)$. Depending on the comparison, we have $f(x)=\theta_1f_1(x)$ if $\theta_1>\theta_2$ and $f(x)=\theta_2f_2(x)$ if $\theta_1<\theta_2$. The output $f(x)$ is discontinuous and thus non-differentiable at $\theta_1=\theta_2$ since $f_1\ne f_2$. When active expert changes from $f_1$ to $f_2$, the sparse gradient computed with $f_1$​ no longer accurately reflects the true gradient for the router, which is now related to $f_2$​. This issue does not arise with ReLU, as its gradient is inherently sparse and aligned with its output.
• Empirically, we further measured the percentage of expert activation states that change in a single update on a calibration set ("flip rate" of the router outputs). Our results show that the flip rate of the TopK router is consistently higher than that of the ReLU router, with the gap increasing as $E$ grows (up to $2-3\times$ more flips for TopK at $E=32$). This indicates that the ReLU router produces more stable routing decisions. For more details, please refer to Appendix A.

Table: Flip rate comparison with $N=$ 182M, $E=32$, $k=1$

Q4: Attribution of the success: continuity, dynamic expert allocation or multiple phases of training.

Thank the reviewer for the insightful question!

The multiple phases of training is not critical to the performance improvement, as Stage I/II takes only ~0.17% of the total training and Stage III is just the same with TopK. But it does provide a stable initialization to the experts of ReMoE and is indisposable.

As for the other two improvements—continuity and dynamic expert allocation—it is challenging to determine which contributes more significantly to the effectiveness of the ReLU router. These two aspects are inherently intertwined in the ReLU solution: a continuous and fully differentiable router ensures continuity at zero for each router output, allowing the outputs to transition smoothly between zero and non-zero independently, which leads to dynamic expert allocation. Conversely, dynamic expert allocation suggests that the router can flexibly update the activation states of experts, which inherently requires a degree of smoothness around zero, thereby ensuring continuity. Therefore, both continuity and dynamic expert allocation are essential to ReLU router, working together to enhance its effectiveness.

We thank the reviewer for considering our work as interesting and insightful and for the valuable feedback. We hope that our responses provide additional clarity. Below we address the concerns:

Q1: Small experiment scale

• Model scale: The largest model in our experiments is the 978M × 8 model, comprising a total of 5.7B parameters.
• Data scale: We extend the longest training duration from 30B tokens to 120B tokens for the 469M $\times$ 8 model with total parameters of 2.5B in Appendix C. The results demonstrate that ReMoE continues to outperform MoE:

Table: Results for active parameters $N=$ 469M and total parameters of 2.5B models trained on 120B tokens

Q2: Downstream evaluation results:

In the revised version, we have included downstream evaluation results for all experiments. These results demonstrate that ReMoE consistently outperforms MoE in downstream tasks, aligning with the perplexity results. For a detailed breakdown of the results, please refer to Appendix E.

Q3: Analysis on why fully differentiable ReLU router will help

The ReLU router helps mainly in the following two aspects:

• Differentiability: In TopK, expert activation state switches are discontinuous (e.g.$(0.51,0)\rightarrow(0,0.51)$, even though the pre-TopK outputs transition as $(0.51,0.49)\rightarrow(0.49,0.51)$)In contrast, the ReLU router provides continuous and differentiable transitions (e.g.$(0.01,0)\rightarrow(0,0.01)$), ensuring smooth gradient flow.
• Dynamism: The ReLU router allows tokens to dynamically route to a variable number of experts, which can be learned by the model. In comparison, TopK enforces the selection of exactly $k$ experts per token, regardless of differences in token importance.

To further analyze this, we measure the percentage of expert activation states that change in a single update on a calibration set ("flip rate" of the router outputs). Our results show that the flip rate of the TopK router is consistently higher than that of the ReLU router, with the gap increasing as $E$ grows (up to $2-3\times$ more flips for TopK at $E=32$). This indicates that the ReLU router produces more stable routing decisions. For more details, please refer to Appendix A.

Table: Flip rate comparison with $N=$ 182M, $E=32$, $k=1$

Q4: Meaning of $N$

We apologize for the confusion. $N$ refers to the active parameters. The largest model of 978M $\times$ 8 has 5.7B total parameters. We have updated the subtitle "Scaling in parameters $N$" to "Scaling in active parameters $N$" and added the total parameters count to the main paragraph in Section 4.3 to avoid misunderstandings.

Q5: Dynamic allocation based on semantic information

We sincerely thank the reviewer for the insightful question! To inspect this, we measured the number of activated experts across different domains, with the results presented in Appendix G. Our findings show that ReMoE effectively learns to allocate computations dynamically based on semantic information. It assigns varying numbers of experts to different domains, with these differences becoming more pronounced in the deeper layers. Thanks again for your advice!

Table: Average active expert count of ReMoE on different domains:

Dear Reviewer yJh9:

We really appreciate your insightful and constructive comments that help us improve this paper. We have taken our maximum effort to address all your concerns.

As the ICLR public discussion phase concludes in two days, we kindly ask if our responses have fully addressed your concerns. If there are any remaining issues, we’d be happy to provide further clarifications. If you feel that your concerns have been resolved, we would greatly appreciate it if you could consider raising the score.

Thank you once again for your time and thoughtful input. We hope you have a wonderful day!

Best wishes,

The Authors

We thank the reviewer for recognizing our work as interesting and promising and for offering thoughtful and constructive feedback. We sincerely hope that our responses can provide additional clarity and would greatly appreciate it if the reviewer could reconsider the score. Below we address the reviewer’s concerns in detail:

Q1: Training speed analysis

We have added a detailed analysis of training speed for MoE and ReMoE in Appendix D. This includes an end-to-end training time comparison with a breakdown across the three stages, and throughput measurements for both training and inference under varying active parameters and parallelism configurations. The key results are briefly summarized in the following two aspects:

• Overhead of Stages I/II: Stages I and II do introduce additional computational overhead due to activating more experts. However, these stages account for only 0.17% of the overall training process, resulting in a negligible overhead of just 0.58%:

Table: End-to-End Training Time Comparison of MoE and ReMoE (Unit: Hours).

• Overhead of dynamic expert selection in ReMoE: the dynamic expert selection in ReMoE doesn't introduce overhead with state-of-the-art MoE implementation of grouped-GEMM. We don't trade speed for performance: the training and inference speeds of ReMoE and MoE are comparable:

Table: relative throughput of ReMoE over MoE.

Q2: Downstream evaluation results

In the revised version, we have included downstream evaluation results for all experiments. These results demonstrate that ReMoE consistently outperforms MoE in downstream tasks, aligning with the perplexity results. For a detailed breakdown of the results, please refer to Appendix E.

Q3: GPU utilization concern of dynamic expert selection:

In practical MoE computation, an all-to-all dispatch and combine mechanism is employed, grouping tokens routed to the same expert. As a result, the number of experts chosen per token is inconsequential, as the tokens are ultimately combined into input matrix blocks for each expert with variable shapes.

For ReMoE, we implement a similar all-to-all token dispatcher as in TopK MoE with capacity constraints. Like ReMoE, TopK MoE with capacity also results in uneven expert selection for tokens, where tokens are not routed to exactly K experts. This does not introduce GPU utilization issues.

Empirically, we measure the throughput of both training and inference of ReMoE and MoE and find they are comparable in speed, as stated in Q1.

Q4: Concern on introduced hyperparameters:

We would like to emphasize that across all our experiments, spanning various model sizes, expert counts, and granularities, we consistently used the same set of heuristically preset hyperparameters, specifically $\alpha=1.2,\lambda_0=1e^{-8}$, without any additional tuning. This consistency demonstrates the robustness of the hyperparameters.

Regarding the oscillation, the presented case of $\alpha = 1.5$ was intended as an abnormal value in our ablation study. With $\alpha = 1.5$, the regularization coefficient changes by 1.5 times with each update, representing an unusually large step that leads to oscillations. However, for more reasonable $\alpha$ values close to 1—specifically in the range of 1.05 to 1.3—the validation loss remains stable and consistent.

To further validate robustness, we replicated the ablation studies on larger models with active parameters $N =$ 469M and total parameters of 2.5B. The results confirm that the performance remains robust across hyperparameter settings. We hope these results can alleviate the reviewer's concern.

Table: Robustness of introduced hyperparameters on ReMoE with $N=$ 469M active parameters.

Dear Reviewer URhd:

We really appreciate your insightful and constructive comments that help us improve this paper. We have taken our maximum effort to address all your concerns.

As the ICLR public discussion phase concludes in two days, we kindly ask if our responses have fully addressed your concerns. If there are any remaining issues, we’d be happy to provide further clarifications. If you feel that your concerns have been resolved, we would greatly appreciate it if you could consider raising the score.

Thank you once again for your time and thoughtful input. We hope you have a wonderful day!

Best wishes,

The Authors

Dear Reviewer URhd,

We sincerely appreciate the time and effort you have dedicated to reviewing our paper, as well as your thoughtful feedback and recognition of our work. We have carefully addressed your questions to the best of our ability.

As the extended discussion period deadline approaches, please do not hesitate to reach out if you have any additional questions or require further clarification. We would be delighted to continue the conversation and address any remaining concerns.

Thank you once again for your valuable insights!

Best wishes,
The Authors

We would like to express our gratitude to the reviewer for recognizing the novelty and potential of our work, as well as for providing insightful and constructive feedback. We sincerely hope that, after considering our responses, the reviewer will reconsider the score. Below, we address the reviewer’s concerns in detail:

Q1: Greatest concern on attribution of the success: is it due to the extra cost in the first 100 steps?

We sincerely thank the reviewer for raising this important question and appreciate the suggestion to evaluate this aspect. However, the first 100 steps are not particularly expensive. Our measurements of the end-to-end time consumption for different stages of training MoE and ReMoE reveal that the overhead of Stage I/II (the first 100 steps) accounts for only 0.58% of the total training time.

Table: End-to-End Training Time Comparison of MoE and ReMoE (Unit: Hours).

Following the reviewer's advice, we constructed a dMoE baseline where training starts with $k = \text{int}(0.8 \cdot E) = 6$ experts out of $E = 8$ for the first 100 steps, then switches to $k = 1$ thereafter. We refer to this configuration as "MoE with warmup." The results of this experiment are as follows:

Table: Comparison of MoE, MoE with warmup, and ReMoE

The results show that "MoE with warmup" achieves a slightly lower validation loss but slightly worse downstream performance compared to the standard MoE. Notably, ReMoE outperforms both. This observation suggests that the first 100 steps, accounting for only 0.16% of the total training procedure, do not provide a critical improvement to overall performance.

To further investigate, we extended the warmup setting to our scaling experiments with larger $E$. In these experiments, we used more experts during Stages I and II, mirroring ReMoE, ensuring that the compute budgets for MoE and ReMoE were equivalent across all three stages:

Table: Results for MoE with warmup under different expert count $E$

We find that when $E$ increases, the warmup setting does yield slightly larger improvements since more compute is used. But ReMoE still performs better, and the the steeper slope in $E$ preserves.

From another perspective, as shown in Figure 5, the training loss of ReMoE at 25B tokens is comparable to that of dMoE with 30B tokens. This gap is far larger than the overhead in Stage I/II, proving the superiority of ReLU routing.

Due to time constraints, we are unable to apply warmup to all our experiments during the rebuttal phase. Additionally, MoE with warmup is not a standard approach, and its results have not been commonly reported in previous literature. However, for the camera-ready version, we will include updated results with extended training steps for MoE to match the additional compute used in Stages I and II of ReMoE across all scaling experiments.

Dear Reviewer qrQT:

We really appreciate your insightful and constructive comments that help us improve this paper. We have taken our maximum effort to address all your concerns.

As the ICLR public discussion phase concludes in two days, we kindly ask if our responses have fully addressed your concerns. If there are any remaining issues, we’d be happy to provide further clarifications. If you feel that your concerns have been resolved, we would greatly appreciate it if you could consider raising the score.

Thank you once again for your time and thoughtful input. We hope you have a wonderful day!

Best wishes,

The Authors

Thank you for taking the time to run these experiments.

Cost of the first 100 steps I agree with the author's point that 0.58% of total training time is negligible. For longer industry-scale training runs this number will also shrink. The only obstacle I see to adoption, here, are logistical issues related to the need for re-sharding checkpoints if significant increases in parallelism are required to fit the temporarily more dense model in memory (this should be the case for the largest dense models). Such issues are addressable, however, and should not stand in the way of publication.

Attribution of success I read the new results as follows: 1) while appreciated for reproducibility, downstream average accuracy is noisy and unreliable at this model size, 2) validation accuracy can be relied upon at this model scale, and 3) it seems that warmup does help the TopK MoE improve over the baseline, but the improvement relative to ReLU MoE diminishes with more experts. This addresses my concern that the improvement of ReLU MoE is entirely due to the near-dense phase at the beginning of training. If longer near-dense training allows TopK to match the performance of ReLU MoE, the authors can still claim their methods require less near-dense training.

Follow-up questions/comments

• How do the training curves look for TopK with warmup vs ReMoE? You included plots for other experiments during the rebuttal phase why not these?
• If the number of near-dense training steps for TopK increases say from 100 to 500 or 1000 does the final performance of the TopK MoE also improve?

We thank the reviewer for finding our method clear and easy to follow. We value the constructive feedback and have carefully considered the concerns raised. Below we provide detailed responses to address the reviewer's concerns:

Q1: Lack of Novelty?

• While replacing TopK+Softmax with ReLU may appear simple, the underlying reasons for its effectiveness are nontrivial: This approach addresses the issue of discontinuity in TopK-routed MoE and facilitates dynamic expert allocation.
• Additionally, controlling the sparsity of the ReLU router and integrating load balancing mechanisms pose nontrivial challenges that we successfully address.
• Notably, ReMoE is the first fully differentiable MoE architecture applicable to autoregressive models that outperforms TopK-routed MoE.

Q2: Improvement Not Significant?

While the improvement in final validation loss may seem small in absolute terms (~0.8% decrease from 1.937 to 1.921 in Figure 6(a)), it is significant in three ways:

• Token efficiency: The perplexity result translates to approximately 16.7% fewer tokens required to achieve the same perplexity: ReMoE needs only about 25B tokens to reach a validation loss of 1.937, compared to 30B tokens for dMoE.
• Expert efficiency: As shown in Figure 6(b), ReMoE with 32 experts outperforms MoE with 128 experts.
• Scalability: ReMoE exhibits a steeper scaling slope with expert counts ($E$), suggesting it offers greater potential for improvement as $E$ increases compared to MoE.

Q3: Downstream evaluation results

In the revised version, we have included downstream evaluation results for all experiments. These results demonstrate that ReMoE consistently outperforms MoE in downstream tasks, aligning with the perplexity results. For a detailed breakdown of the results, please refer to Appendix E.

Q4: Effect of stage I/II on efficiency, memory and performance

• Efficiency: Stages I and II indeed activate more experts and require additional compute. However, they comprise only ~100 steps, which is approximately 0.17% of the total training process. As such, the time overhead is negligible. To provide further clarity, we included an end-to-end training time comparison in Appendix D. The results indicate that the time overhead from Stages I and II amounts to just 0.58%.

Table: End-to-End Training Time Comparison of MoE and ReMoE (Unit: Hours)

• Memory: Stages I and II activate more experts compared to MoE, resulting in higher activation memory consumption. However, this issue can be significantly mitigated using the activation checkpointing technique, which avoids storing intermediate results of extra activated experts by recomputing them on-the-fly during the backward pass. With activation checkpointing, the memory overhead is reduced proportionally to the number of layers.

  To quantify this, we measured the memory usage of MoE and ReMoE with $N = 469$M, $E = 8$, using activation checkpointing and a micro batch size of 8. The peak memory usage shows a modest increase of 2.8% for ReMoE compared to MoE (28.59GB->29.40GB).

  Clarification of the memory overhead is made in Section 3.5.

• Performance: The two stages utilize more compute and thus slightly increase the performance. But the improvement is marginal since they account for only 0.17% of the total training process. To further investigate, we conducted an experiment by training a TopK MoE with a similar warmup strategy: activating nearly all experts during the first 100 steps and switching to $k = 1$ thereafter. The results show that MoE with warmup achieves a slightly lower validation loss but slightly worse downstream performance compared to standard MoE. Both configurations underperform relative to ReMoE:

Table: Comparison of MoE, MoE with warmup (extra compute in Stage I/II), and ReMoE

Dear Reviewer VcWp:

We really appreciate your insightful and constructive comments that help us improve this paper. We have taken our maximum effort to address all your concerns.

As the ICLR public discussion phase concludes in two days, we kindly ask if our responses have fully addressed your concerns. If there are any remaining issues, we’d be happy to provide further clarifications. If you feel that your concerns have been resolved, we would greatly appreciate it if you could consider raising the score.

Thank you once again for your time and thoughtful input. We hope you have a wonderful day!

Best wishes,

The Authors