ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing
Authors: Ziteng Wang
Jun Zhu
Jianfei Chen
Venue: ICLR 2025 Poster (2025)
Keywords: Mixture-of-Experts
Differentiable Routing
Sparsity
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose ReMoE, a fully differentiable MoE with ReLU routing. ReMoE consistently outperforms vanilla TopK-routed MoE and exhibits superior scalability w.r.t. the number of experts.

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.

Published as a conference paper at ICLR 2025
REMOE: FULLY DIFFERENTIABLE MIXTURE-OF- EXPERTS WITH RELU ROUTING
Ziteng Wang, Jun Zhu, Jianfei Chen
Dept. of Comp. Sci. and Tech., Institute for AI, BNRist Center, THBI Lab,
Tsinghua- Asinghu--Bosch Bosch Joint ML Center, Tsinghua University wangzi t e22 36 ma 1s tsinguua edu.cn; {dcsz jianfeiicettsinnguaaa edu.. edu. cn
ABSTRACT
Sparsely activated Mixture-of- Expersss (MoE) models widely adopted to scale up model capacity without increasing the computation budget. However, vanilla TopK routers are trained in discontinuous, non-differentiable way, limiting their performance and scalability. To address this issue, we propose ReMoE, a fully dif- ferentiable MoE architecture that offers simple yet effective drop-in replacement for the conventional + routing, utilizing ReLU as the router instead. We further propose methods regulate the router's sparsity while balancing the load among experts. ReMoE's continuous nature enables efficient dynamic allo- cation of computation across tokens and layers, while also exhibiting domain spe- cialization. Our experiments demonstrate that ReMoE consistently outperforms vanilla TopK-routee 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 htttps//gghhbb.co///mm//////////
INTRODUCTION
Transformer models Vaswani 2017) consistently improve performance the number of parame- ters increases (Kaplan et al.J (020). However, scaling these models is constrained by computation resources. Sparsely activated Mfxture-of- (MoE) (Shazeer et al. 2017) mitigates this chal- lenge by employing a sparse architecture that selectively activates a subset of parameters during both training and inference. This conditional computation allows MoE models expand model capacity without increasing computational costs, offering a more efficient alternative to dense models.
The key component in MoE the routing network, which selects the experts activate for each token. Various routing methods (Shazeer et al. 2017: Lewis el al. 2021
Roller et al., 2021 I Zhou
et al. 2022 have been proposed, with TopK routing (Shazeer al. 2017 being the most commonly adopted. However, the vanilla TopK router introduces a discrete and non- non-differentiable training objective Shazeer et al. 2017 Zoph et al. 2022), limiting the performance and scalability.
Recent works on fully-diffes differeiffeee MoE aim to overcome this limitation. Soft MoE (Puigcerver al. 2023) introduces token merging, while SMEAR (Muqeeth et al.) (023) proposes expert merg- ing. However, both approaches break token causality, making them unsuitable for autoregressive models. Lory (Zhong et al., (2044) improves upon SMEAR and applicable to autoregressive mod- els. But it underperforms vanilla MoE with TopK routing.
In this work, we address the discontinuities by introducing ReMoE, an MoE architecture that in- corporates ReLU routing TopK routing, which computes a softmax distribution over the experts and calculates a weighted sum of the largest K experts, ReLU routing directly controls the active state of each expert through a ReLU gate. The number of active experts is determined by the sparsity of the ReLU function. To maintain the desired sparsity, propose adding load-balancing refined Lı regularization the router outputs, with an adaptively tuned coefficient. This approach ensures that ReMoE maintains the same computational costs TopK-routed MoE.
simple yet effective drop-in replacement for TopK routing. Unlike
"Corresponding author

---

Published as a conference paper at ICLR 2025
Topk Routing
Routing
.40
23
:33
42
AB
projection
Exprrt
and
Expert
add
Token
Expert
Expert
Expert
Expert
Tokon T
Token J
Tokenn
Tokon
Dpprnt
Expert
projetionn
projection
add
Softman
Topí
Figure I: Compute flows of vanilla MoE with TopK routing and ReMoE with ReLU routing. Posi- tive values are shown in orange, and negative values in blue, with deeper colors representing larger absolute values. Zeros, indicating sparsity and computation savings, are shown white. The red dash arrows in TopK routing indicate discontinuous operations, Compared with TopK routing MoE, ReMoE uses ReLU to make the compute flow fully differentiable.
Compared to TopK routing, ReLU routing is continuous and fully differentiable, as the ReLU func- tion can smoothly transition between zero and non-zero values, indicating inactive and active. Be- sides, ReLU routing manages the "on/off" state of each expert independently, offering greater flexi- bility. Moreover, the number of activated experts can vary across tokens and layers, enabling more efficient allocation of computational resources. Further analysis reveals that ReMoF effectively learns to allocate experts based on token frequency and exhibits stronger domain specialization.
Our experiments on mainstream LLaMA (Touvron et al. 2023) architecture demonstrate that ReLU routing outperforms existing routing methods including TopK routing and lullyddiffrrentiable Lory. Through an extensive investigation across model structures, we find that ReMoE consistently outper- forms TopK-routess MoE across broad range of active model sizes (182M to 978M), expert counts (4 to 128), and levels of granularity (1 to 64) (Krajewski et al., 2024). Notably, in terms of scaling behavior, observe that ReMoF exhibits steeper performance improvement the number of experts scales up, surpassing traditional MoE models.
PRELIMINARIES
2.1 MOE FOR DECODER OONYY TRANSFORMER
A typical decoder-only Transformer model consists of T, layers, each containing Self-Attention module and Feed-Forward Network (FFN) module. MoE modifies this structure by replac- ing each FFN module with an MoE module, which comprises small router and several experts FFN FFN where each expert is equivalent to the original FFN and É denotes the number of experts. Given the input x' (xt) E R of the layer I, where T is the number of tokens in batch and is the hidden size, the output y' (y!)T is computed as:
F
y Σ R(x)). (al: dy
(1)
Here, R(-) represents the routing function, and fffn is the intermediate size of the FFN, typically set to difn 4d.
2.2 TOPK ROUTING
TopK routing (Shazeer et al. 2017 Lepikhin et al. 2020 Fedus et al., 2022) is the most commonly used method for defining the routing function R?()). It introduces sparsity in the MoE computation
2

---

Published as a conference paper at ICLR 2025
by forcibly zeroing out smaller elements:
R(xl)= TopK (Tof(SSotmaa((e Wi), k)
where W E Rdx is the router's weight matrix, and TopK( (- k) retains the top k largest values while setting the rest zero. This mechanism allows for skipping the computation of the FFN, functions corresponding to the zeroed-out R(x()e values in both the forward and backward passes.
3 OUR METHOD: REMOE
3.1 MOTIVATION: FROM TOPK TO RELU
For given token (xe) after Softmax, TopK introduces jump discontinuity at the k-th largest value, denoted as T[K]. by zeroing out the values smaller than m[k]- This can be expressed as:
TopK(2z az k)c Ze 1fze= 2 t(x, k)}, t(x, k)) 2p4
(3)
where {(-} is the indicator function, returning if the condition met and 0 otherwise.
Topk(xr
Relu()))
As shown in Figure 2 the jump discontinu- ity can be eliminated by setting the breakpoint
t(x, k) = O, which actually corresponds to the ReLU function:
x(A]
0
ReL.U(z). = Fe f 2 0}
(4)
Figure 2: Comparison between TopK and ReLU.
At high level, ReLU improves upon TopK by aligning the breakpoints of all inputs and setting them to O. This ensures that the output is continuous O, where the experts transition between active and inactive. As result, the training pipeline becomes fully differentiable.
3.2 DIFFERENTIALLE RELU ROUTING
We define the ReLU routing function follows:
(5) with (1 - () being the desired sparsity of ReLU, where k is the number of active experts and E is the total number of experts. This ensures that the computational cost remains equivalent to that of TopK routing.
R(x)) = ReLU(x,WW)
In vanilla TopK routers, the Softmax outputs sum to 1, representing the probabilities of selecting each expert, after which TopK eliminates those with lower probabilities. contrast, ReLU routers discard the Softmax function, relying on ReLU's naturally non-negative outputs. The outputs of ReLU routers represent the weights assigned to each expert, which include 0. Instead of hard- coding expert selection with discontinuous TopK function, ReLU allows the router to learn which experts to activate (i.e., when to produce Os) in a fully differentiable manner.
Another key difference is that in TopK routing, each token is routed to exactly & experts, whereas in ReLU routing ReMoE, the routing decisions are independent, allowing tokens to be routed variable number of experts. This flexibility is advantageous, as not all tokens have the same level of difficulty. ReMoE can allocate more computational resources to more challenging tokens, a dynamic allocation strategy that we explore further in Section 5.1
TopK routing introduces discrete loss function when the set activated experts changes, whereas ReLU routing remains continuous and fully differentiable. For instance, in two-expert Topl- routing model, a small weight update that alters the softmax result from æi = (0.51, 0.49) to 22 (0.49, 0.51) shifts the TopK output from (0.51, 0) to 00,1. creating discontinuity. In contrast, ReLU routing only changes the activated experts when the routing output is near zero. For example, an output shift from (0.01, 0) to (0. 0.01) remains continuous. Further details on the stability analysis of these two routers can be found in Appendix A
A comparison of the compute flow between ReMoE and MoE is shown in Figure

---

Published as a conference paper at ICLR 2025
3.3 CONTROLLING SPARSITY VIA ADAPTIVE L1 REGULARIZATION
ReMoE controls computational costs by managing the sparsity of the ReLU output, targeting a sparsity level of (1 3 ). However, directly training the ReLU router often results in lower sparsity, as the model tends to activate more experts to increase capacity. To meet the desired budget, we need enforce higher sparsity in the ReLU output.
We achieve this by introducing a regularization loss, Lreg. to the loss of language model, Limi
L Lim. >eCregs
where A is an adaptive coefficient based on the current training step i. Initially, we set X to a small value and employ simple zeroth-order algorithm to update it:
A++ X; a sign( - )
Here, >> is preset update multiplier, and S denotes the average sparsity of all router outputs at the step i:
T = LTE ΣΣ> (R(x)) 00)
The idea behind Equation| 7 is that when the average sparsity S: falls below the target sparsity (1 )) we increase X; by factor of a, strengthening the regularization and encouraging higher sparsity. Conversely, if the sparsity exceeds the target, reduced. We heuristically set λo
and a = 1.2 in all our experiments, and demonstrate the robustness of these hyperparameters in Appendix B
The regularization term Lreg uses the L1 norm,, following prior work Li et al. 2022 2024.. to effectively encourage sparsity:
Song et al.
Frog
1 LT
R(a))l
R(wl).
LT
1=1
1=1
The second equation holds because the output of the ReL.U function is non-negative.
The term Lreg represents the average value of all router outputs, including zeros. By taking the derivative of X,Creg we observe that the regularization effect adds X to the gradient of each non-zero router output, effectively driving the outputs toward zero and enhancing sparsity.
1.00
Sparsity
0.75
0.50
Sparsity=0.8751=00000 Desired Sparsity=0.875
With this Li regularization, we can control the sparsity around the desired level of (1- with only minor fluctuations, as shown in Figure Consequently, ReMoE ensures that, on average, tokens are routed k experts across different layers and tokens, maintaining the same FLOPPs vanilla Topted MoE from statistical perspective. Our benchmarking results in Ap- pendix demonstrate that ReMoE can achieve nearly identical training and inference through- puts conventional MoE, providing an effi- cient alternative without compromising speed.
0.25
10000 20000 30000 40000 50000 60000
Step
Figure 3: The sparsity of ReMoE with E 1 is effectively maintained around the de- sired target. Sparsity values for all steps are plot- ted without averaging sampling. The mean and standard deviation are calculated excluding the first 100 warm-up steps.
3.4 INTEGRATE LOAD BALANCING INTO L, REGULARIZATION
Load imbalance is a significant issue in MoE design, potentially leading to routing collapse (Shazeer let al. 2017; Muennighoff et al. 2024 and uneven computational distribution across multiple de- vices. The L, regularization in Equation Itreats the router output for each expert and each layer / equally, which can contribute load balancing problems.
4

---

Published as a conference paper at ICLR 2025
Stape/
Stage
Stage MV
10
Smap Stage
Stage
10
Stage iSttage Stage
0.8
Sparsity
10-1
10r1
0.6
10
Eree
Mmy WV
10-1
11--
0.2
Sparsity
Desired Sporsky
10-1
10.3
10
0.0
10
10
60000
Step
Step
Stop
(a) Sparsity S,
(b) Coefficient term X and regular- (c) Language model loss Zim and overall regularization XgGreg
ization term freg
Figure 4: Natural Three Stage Training in ReMoF.
To address this, we introduce load-balancee refinement to the L1 regularization:
E
Lreg.th 1b
R(a))
(10)
LT
E Σ (R(a)) 0} KT
f1.
(11)
Here, Ji, is non-differentiable and represents the average activation ratio of expert in layer relative to the desired ratio & This serves as a weight for the corresponding router output, modifying the added gradient of non-zero router outputs to ft.A: This mechanism penalizes experts receiving more tokens by driving their router outputs toward zero more rapidly.
E.
Although derived from regularization, this formulation identical the load- balancing loss in vanilla TopK routing (Fedus et al. 2022). In TopK routing, the outputs of Softmax sum to 1, giving the loss lower bound of 1. In contrast, ReLU routing outputs can be arbitrarily small, making Greg.18 ib trivially bounded at O. Therefore, unlike in MoE, cannot fix the coefficient A in ReMoE, as this would lead to routing collapse toward Thanks to the adaptive update of X4 we can balance sparsity control and load balancing within a single formulation, given in Equation I0
Further discussion on load balancing in ReMoE can be found in Section 5.2 and we adopt this load-balancing refined L1 regularization in our later experiments.
3.5 NATURAL THREE-STAGE TRAINING REMOE
With the regularization scheme described above, we observe a clear and naturally occurring three- stage separation during the training of ReMoE depicted in Figure
The first stage is the warm-up stage, or the dense stage. During this stage, A is small, while Cim is large and decreases rapidly. Training ReMoE this stage is nearly equivalent to training its dense counterpart with the same total number of parameters. Each expert processes more than half of the tokens, allowing the experts diversify from their random initializations.
The second stage is the sparsifying stage, or the dense to sparse stage. At this point, the sparse regularization term X; X;Creg becomes significant, causing the ReLU routers activate fewer experts. This forces the experts to become more diverse without causing an increase in Lim.
The third stage the stable stage, or the sparse stage. In this phase, the sparsity S, stabilizes at the preset target. During this stage, Lim is optimized while being softly guided along the sparse subspace by freg- Both freg and A, change very slowly, with Lreg gradually decreasing and X; gradually increasing. However, the overall regularization term, X. remains relatively constant.
should be noted that Stages and II introduce additional computational cost and memory con- sumption since more experts are activated. However, the time overhead is negligible since they generally require only ~10 iterations (~0..1%% the total steps in our setting, benchmarking re- sults are detailed in Appendix[]] The memory overhead can be minimized by temporarily reducing the micro-batch size or by employing the activation checkpointing technique that avoids storing intermediate results of activated experts by recomputing them on-the-fly during the backward pass.
5

---

Published as a conference paper at ICLR 2025
Dence
Model
ARC=
ARC-e
BooiQ
HellaSwag
LAMBADA
PEQA
RACE
Ave
Hash
Lany
Dense
19.45
43.35
54.40
78.51
31,09
61.97
28.52
38.20
Than LO33
Sprree breu2
Hash
19.28
45.45
54.95
29.68
31.44
63.05
27.06
38.79
2.2
Lory
20.31
42.97
69.54
28.75
32.35
62.24
27.75
37.70
SparseMiste-22
19.80
46.72
15.96
30.21
34.12
62.89
29.00
38.39
RC
18.86
42.97
60.21
29.14
20.26
61.92
27.37
38.53
2.0
aMoil
20.05
45.16
57.83
29.83
32.97
63.55
28.
99.67
ReMot:
20.22
46.03
54.16
30.26
35.94
63.55
29.38
40.03
1.5
TTicenstB)
Figure 5: Training curves of dif- Table 2: Zero-shot accuracy of different routing methods on ferent routing methods.
downstream tasks.
4 EXPERIMENTS
4.1 SETUP
Infrastructure We leverage Megatron-LM (Shocybi al. 2019) code base and implement ReLU routing as a drop-in replacement for the original Topk routing, supporting all forms of model parallelism: Data, Tensor, Pipeline, and Expert Parallelism (Shoeybi et al. 2019) Narayanan et al.
2021
Korthikanti et al. 2023).
Model Architecture. We experiment with the mainstream LLaMA Touvrro et al. 2023) archi- tecture, featuring grouped query attention (GQA) (Ainslie et al. (023). SwiGLU (Shazeer 2020 activation function, RoPE (Su et al. 2019). The context length ferent dense backbone sizes as shown in Table I For vanilla MoE we adopt load balancing loss
2024) position embedding, and RMSNorm (Zhang & Sennrich
1024, and the batch size is 512. We experiment with three dif-
of weight 0.01 following Fedus et al. (2022)
ReMoE we use the adaptive load balancing L1
regularization in Equation TQ
Training Settings. We train the models on The Pile (Gao et al., 2020), an 800 GB diverse corpus.
All models are trained for 60k steps (~ 30B tokens), which exceeds the compute-optimal dataset size predicted by Krajewski et al. (2024) and is enough to converge. The byte pair encoding (BPE) tokenizer (Sennrich, 2015) is used. We adopt AdamW (Loshchilov (2017) as the optimizer with 31 = 0.9, 32 = 0.999 with ZeRO optimization (Rajbhandari et al.) 2020). The learning rate is set to be 5e 4 with a cosine scheduler. All models are trained with 8 NVIDIA A100 GPUs.
Size
#Parameters
hidden_size
num num_layers
num num_heads
num-groups
GFL.Pss
Small 182M
768
12
995
Medium
469M
1024
24
16
4
2873
Large
978M
1536
24
5991
Table 1: Configurations for the dense backbones. FLOPs are calculated with a single sequence
according to Narayanan et (202I).
4.2 COMPARISON WITH OTHER ROUTING METHODS
We compare ReMoE against the following methods: (i) Token-choice dropless TopK routing
(dMoH) (Gale et al. 2023) (ii) Expert-choice TopK routing (EC) (Zhou et al., (022) (iii) Deter-
ministic hash routing (Hash) (Roller al. (021) (iv) Fully-differentiablle expert-merging routing
(Lory) Zhoog et al. 2024) (v) TopK routing with improved gradient estimate (SparseMixxrr-v2) (Liu
al. 2024b).
The performance of these methods is evaluated with active parameters N = 1822 and the expert count E = 8. We fix the active expert count to for straightforward comparison with the dense counterpart. For the Hash method, we use mod E hashing function. And for Lory, the segment length is set to 256, following the original paper.

---

Published as a conference paper at ICLR 2025
Dense RoMoE
2,05
2.0
1.986
1.271
2.042
2.042
Derie
1.55
1.536
RaMot
Vale
1.206
1.90
1.907
Vold
1,706
1.85
Demne
NoE
1.803
1.854
1.005
1.800
1.895
1.893
1.7
1.0
ReMot
1.85
1.873
152M
460M bumber Patamaters N
978M
16
64
Expert Cuant
Granubarity
(a) Scaling in N
(b) Scaling in E
(c) Scaling in G
Figure 6: Scalability of ReMoE with respect the number of active parameters (N), expert count (E), and granularity (G). Default config is N = 182M, E 8, G = 1, k = 1. The Y-axis represents the validation loss of each model after training on 30B tokens. ReMoE consistently outperforms MoE across all configurations.
These models are trained on 30B tokens, with the training curves shown Figure [ We evaluate the zero-shot performance of the trained models on the following downstream tasks: ARC (Clark
al.
2018);
BoolQ (Clark et al.] 2019); HellaSwag (Zellers et al. 2019); LAMBADA (Paperno
et al.
2016: PIQA (Bisk et al. 2020); RACE (Lai et al. 2017).
The downstream accuracy results are summarized in Table
Our results show that all MoE models outperform the dense model. Deterministic hash routing performs worse than the learned routing methods. Among the Top-K approaches, token-choice dMoE outperforms expert-choice MoE and SparseMixer-v2 in evaluation. The differentiable routing method Lory surpasses Hash routing in training but underperforms in downstream tasks, with both methods falling short of the standard Top-K K routing. Notably, ReMoE outperforms all methods, including the mainstream Top-K K routing, while benefiting from differentiability.
4.3 SCALABILITY OF REMOE
In this section, we compare ReMoE with state-of-the-an dMoE (hereinafter referred simply as MoE) across varying model parameters N, expert counts E, and granularity levels G to demonstrate its scalability and universal superiority. Since ReMoE demands more computation in both Stage I and Stage II, we increase the number of training steps for the MoE baseline to match the total computation in each setting, ensuring more equitable comparison. We present the final validation losses in Figure with comprehensive downstream evaluation results available in Appendix B
Scaling in active parameters N. To assess scalability with respect to the number of parameters N. fix E' 8 and k 1, while varying active parameters N from 182M to 975M, corre- sponding to the dense counterpart configurations in Table [ The total parameters are 777M, 2.58B, 5.73B respectively. The results, shown in Figure [aa| indicate that ReMoE consistently outperforms MoE across all model sizes. The performance gap does not diminish as the model size increases, suggesting that ReMoE maintains its advantage larger scales.
Scaling in expert count E. In this experiment, we fix the number of parameters at N = 182M and set the number of active experts i 1, while varying the total number of experts E from 4 to 128. The scaling curve in Figurelob]rent reveals that ReMoE consistently outperforms the standard MoE across all configurations of E.
Moreover, key observation is the steeper slope in ReMoE's performance as E increases, compared MOE. This suggests that ReMoE scales more effectively with the number of experts and derives greater benefits from larger expert pools. ReMoE's s differentiable routing strategy appears better suited for leveraging large expert groups, leading to significant improvements in model expressivity and generalization.
Scaling in granularity G. We also evaluate ReMoE and MoE in fine-grained settings. Fine- grained MoH (Dai et al. 2024 Krajewski al. 2024 with granularity G is constructed by dividing
7

---

Published as a conference paper at ICLR 2025
each expert into G' smaller experts, formulated below:
BG
y E
(12)
R(x!) Toop((Softma(aa Wi) kG)
(13)
Fine- sine-grained MoE outperforms vanilla MoE from a scaling law perspective (Krajewski et al. 2024)
and has been adopted in subsequent works (Dai et al.) 2024: Tan et al. 2024 Muennighofi et al.
2024). For fine-grained ReMoE, the routing function remains identical to Equation and the target sparsity is still (1 ). The only distinction lies in the shape of the weight matrix, with Wi RdxEG
We conduct experiments with N = 182M and E' = 8, varying G from 1 to 64 for both fine-grained MoE and fine-graines ReMoE. addition to comparing these models against the dense baseline with the same number of active parameters, we also evaluate their dense counterpart with the same total number of parameters. This is achieved by expanding the intermediate size of the FFN by factor of E, which we denote Densex8. This configuration represents the strict upper bound for MoE and ReMoE, as it is equivalent to a ceature-of-Expertt with all experts activated (Dai et al. 2024
As illustrated in Figure 6c fine-grained ReMoE consistently outperforms fine- ine-ggained MoE. More- fine- gree-graied ReMoE of G' 32 and G 64 reach the performance of the theoretical upper bound, Dense x8, while requiring significantly fewer FLOPs during both training and inference. In contrast, fine-grained MoE is unable match in all settings, making ReMoE a more efficient and effective choice.
DISCUSSION
5.1 DYNAMIC EXPERT ALLOCATION IN REMOE
In ReMoE, each token dynamically activates a subset of experts, allowing the model to adaptively allocate re- sources. We evaluate the performance of the N 182M, E 8,
ReMoE model and analyze the rela- tionship between token frequency and the average num- ber of active experts. As illustrated in Figure the model tends to assign higher number of experts kens, such as @', POTAL', and @'', while reducing the number of active experts for more frequent tokens like An', and the'..
Awarage Active Expert Count
3.0
Arcage Active Export Count
dralom Expert Assignment.
10
2.5
Theen Erequenc
Token Frequency
2.0
10
rarer to-
1.5
10
1.0
10
0.5
10000 20000 30000 40000
This adaptive behavior mirrors the principles of Huff- man tree Huffman| (1952), where more frequent symbols assigned shorter codes, and rarer symbols assigned longer codes. Similarly, ReMoE tends to "cluster on" common tokens by activating fewer experts, effectively compressing the "representation" of these frequent to- kens. In contrast, for rarer tokens, ReMoE activates more diverse set of experts, "encoding" them as richer linear combination at the expert level. This suggests that ReMoE learns to dynamically allocate computational resources, achieving efficient balance be- tween resource usage and the model's capacity, optimizing performance under a constrained expert budget. Dynamic expert allocation is also evident at the domain level, detailed in Appendix G
Sorted Token ID
Figure 7: Correlation between expert allocation and token frequency in Re- MoE. X-axis is sorted by average active expert count and token frequency in log-scale.
5.2 THE ROLE OF LOAD BALANCING IN REMOE
Load imbalance can lead to routing collapse in the vanilla TopK-route MoE, where the router tends to assign the same expert to all inputs, in which scenario the training objective becomes continuous and fully differentiable. As is shown in Figure Sa there is a significant performance gap between MoE models with and without load balancing (LBB).
8

---

Published as a conference paper at ICLR 2025
Deperio
2.6
1.00
Danse
2.4
MoE
0.95
Spartit
Thais
2.2
2.0
AnMis Wo.
Arp
0.75
(Tb
(a) Training curves of MoE and (b) Average routed to- (c) Average routed to- (d) Sparsity across different
ReMoF with and without load kens ratio of ReMoF kens ratio of ReMoE layers in ReMoR balancing
w.o.
w.
Figure 8: Observations the role of load balancing in MoE and ReMoE. White squares in (b)
represent inactive experts with fewer than 1/64 tokens routed to them.
While in ReLU routing, thanks to its differentiablity, even applying the Lı regularization from Equa- tion 2 without load balancing yields comparable results with well-tuned MoE with LB. However, some experts in ReMoE without LB remain inactive, illustrated as white squares in Figure &b which shows the heat map of the average routed tokens ratio (i.c., the fraction of tokens routed the e-th expert the 1-th layer) over 50M tokens in test set. This inactivity can limit the model's capacity.
When load balancing is incorporated into the refined L1 regularization (Equation T0), the experi- ments show a more even distribution of token assignments across experts, with all experts being utilized, as shown in Figure &c The final loss in ReMoE decreases after introducing load balancing.
Besides, we observe ReMoE with LB can produce a smoother sparsity distribution across layers as depicted in Figure [d. This is because ft,e is computed based the absolute number of routed tokens, meaning denser layers receive stronger penalties.
Note that even ReMoE with load balancing (LB) does not yield a perfectly even distribution. How- ever, the trade-off between load balancing and performance can be easily adjusted by modifying the L1 regularization
Equation 10 For instance, changing ft.e to f2, would make the model more sensitive to load imbalance. Additionally, device-level load balancing techniques,
proposed in Dai et al.] (2024). could also be employed. Since load imbalance in ReMoE does not lead to severe routing collapse, primarily becomes hardware utilization issue. As such, we leave the exploration of these variants for future work.
Roured Tokens Ratio
Layer
Layer
Layer 11
0.8
0.6
Github
0.2
Stack
0.0
Wiki
Expert ID
(a) Domain specialization of MoE
Routed Tokens Ratio
Layer
Layer
Layer
0.8
Anxiv
0.6
Books
0.4
C4
Github
Stack
0.0
3
4
0
1
5
5
Wiki
Expert ID
(b) Domain specialization of ReMoE
Figure 9: Average routed tokens ratio for MoE and ReMoE across 12 layers and 8 experts in different domains. The gray dashed lines indicate uniform distribution. ReMoE shows stronger domain specialization.

---

Published as a conference paper at ICLR 2025
5.3 DOMAIN SPECIALIZATION REMOE
The differentiability and dynamic allocation strategy of ReMoB facilitates the development of di- verse experts that specialize in different domains. This allows the router effectively perform ensemble learning by leveraging the expertise of various experts, as demonstrated in our experi- ments.
In Figure D we plot the average routed tokens ratio across different experts, layers, and do- Arxiv, Books, C4, Github, Stackexchange, and Wikipedia-forr MoE and ReMoE models with N = 182M, E = 8. We focus on the first, middle, and last layers (with IDs 0, 5, and 11). The results for most experts in MoE (Figure 9a show a roughly uniform distribution across all domains. In contrast, experts in ReMoE (Figure 9b) exhibit clear domain specialization, being activated with varying frequencies across different domains. For example, more than half of the tokens from Arxiv, Github, and StackExchange- domiann that emphasize structured, non-natural languages like LaTeX and Python are routed to Expert 6 in Layer 5, significantly more than in other domains, A more detailed result of domain specialization can be found in Appendix F
RELATED WORKS
6.1 MIXTURE-OF-EXPENT
Mixture-of Mixture-of-Experts (MoE) was initially proposed in the early 1990s (Jacobs et al., 1991 Jordan
& Jacobs 1994) and later introduced into large-scale neural networks as sparse submodule for efficiency (Shazeer et al.] 2017). Advances like GShard (Lepikhin et al.] (2020) and Switch Trans- former (Fedus et al.j 2022) integrated sparse MoE into Transformer models, achieving signifi- cant results. More recently. MoE has been used in commercial-scale language models such as Mixtral-8x7B (Jiang
2024). DeepSeekMoF: 16B (Dai al. 2024), and Snowflake Arctic
17B (Snowflake]22g 2024).
6.2 ROUTING MECHANISMS MoE
Various routing methods have been developed for expert selection. Static routers, such as BASE (Lewis al. 2021), use predefined rules like combinatorial optimization, while Hash rout- ing (Roller et al. (021)) relies on deterninistic hash functions, and THOR (Zuo et al.) (202T) assigns experts randomly with regularization. Learned routers adaptively select experts based on token
put, using approaches like REINFORCE (Bengio et al. [2013; Schulman et al. 2015 Clark et al.,
2022] for reinforcement learning, and Topk routing (Shazeer et al. 2017] Zhou al., (022) for token or expert selection, though TopK introduces discontinulties that hinder gradient estimation.
6.3 DIFFEENTIABLE MIXTURE-OP-EXPERTY
Recent work on fully differentiable MoE models addresses the challenges of discrete optimization, basically through token merging and expert merging approaches. Soft MoE (Puigcerver al.) 2023) uses token merging, assigning fixed slots to each expert as a linear combination of input tokens. SMEAR (Muqeeth et al., 2023) merges experts into an ensemble via weighted averaging. However, both methods require a full probability map of input tokens, making them unsuitable for autoregres- sive models. Lory (Zhong et al. (024) preserves autoregressiveness by segmenting sentences to merge experts but underperforms compared to TopK routing.
7 CONCLUSION
In this paper, we propose ReMoE, fully diffecreniaalee MoE architecture with ReLU routing. The simple yet effective ReLU routing function acts as drop-in replacement for the conventional TopK+Softmax routing, offering (i) continuity and differentiability, and (ii) dynamic expert allo- cation across tokens and layers. With the adaptive load balancing L, regularization, ReMoE uni- versally outperforms Topted MoE across various model sizes, expert counts, and levels of granularity, demonstrating sharper performance gains as the number of experts scales.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
The authors gratefully acknowledge Chao Du and Tianyu Pang for the insightful discussions. This work was supported by the NSFC Project (No. 62376131), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University, J.Z is also supported by the XPlorer Prize.
REFERENCES
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv: 13245, 2023.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 3rxiv:1308.343: 2013.
Vincent-Piers Berges, Barlas Oğuz, Daniel Haziza, Wen-tau Yih, Luke Zettlemoyer, and Gargi Gosh. Memory layers at scale. arXiv preprint arXiv:2412.09764, 2024.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical com- monsense natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoff- mann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference on machine learning, pp. 4057- 4086. PMLR, 2022.
Christopher Clark, Kenton Lee, Ming-Wai Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint 1rXiv: 10044, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1803:v:1803.54444 2018.
Róbert Csordás, Piotr Pikekos, Kazuki Irie, and Jürgen Schmidhuber. Switchhead: Accelerating transformers with mixture-of-experts. attention. Advances in Neural Information Processing Sys- tems, 37:74411-74438. 2025.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint arXiv:2401.06066 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39 2022.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matci Zaharia. Megablocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learning and Systems, 5:288-304, 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2 2020.
Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang, and Mao Yang. Seerattention: Learning intrinsic sparse attention in your lms.. arXiv preprint arXiv:2410. 13..10.......... 2024.
Albert Gu and Tri Dao. Mamba: Linear- sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
11

---

Published as a conference paper at ICLR 2025
Xu Owen He. Mixture of A Million Experts, July 2024. URL htto://axiv.oro/// 2407. 04153 arXiv:2407.0411 [cs].
Zihao Huang, Qiyang Min, Hongzhi Huang. Defa Zhu, Yutao Zeng. Ran Guo, and Xun Zhou. Ultra- sparse memory network. arXiv preprint arXiv:2411.12366 2024.
David A Huffman. A method for the construction of minimum-redundanny codes. Proceedings of the IRE, 40(9): 1091-1101, 1952.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Albert Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04008 2024.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481-525155 2025.
Pengkun Jiao, Xinlan Wu, Bin Zhu, Jingjing Chen, Chong-Wat Wah Ngo, and Yugang Jiang. Rode: Linear rectified mixture of diverse experts for food large multi-modal models. arXiv preprint arXiv:2407.12730, 2024.
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-214 1994.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836. 2020.
Vijay Anand Korthikanti, Jared Casper. Sangkug Lym, Lawrence McAfee, Michael Andersch, Mo- hammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5:341-353, 2023.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygózdž, Piotr Sankowski, et al. Scaling laws for fine-grained mixture of experts. arXiv preprini arXiv:2402.0787,, 2024.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pp. 785-794, 2017.
Guillaume Lample, Alexandre Sablayrolles, Marc' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. Advances in Neural Information Processing Systems, 32, 2019.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 2..6..6666 16668. 2020.
Mike Lewis, Shruti Bhosale, Tim Detimers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, pp. 6275-6274. PMLR, 2021.
Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. arXiv preprint arXiv:2210.06313, 2022.
Enshu Liu, Junyi Zhu, Zinan Lin, Xuefei Ning, Matthew B Blaschko, Shengen Yan, Guohao Dai, Huazhong Yang, and Yu Wang. Efficient expert pruning for sparse mixture-of language models: Enhancing performance and reducing inference costs. arXiv preprint 2400:24070000455, 2024a.
12

---

Published as a conference paper at ICLR 2025
Liyuan Liu, Young Jin Kim, Shuohang Wang. Chen Liang, Yelong Shen, Hao Cheng, Xiaodong Liu, Masahiro Tanaka, Xiaoxia Wu, Wenxiang Hu, et al. Grin: Gradient-informed moe. arXiv preprint XXh::2409.12136, 2024b.
I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv:1711.0510/1 2017.
Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, and Hongsheng Li. Not all experts are equal: Efficient expert pruning and skipping for mixture-of-expert large language models. arXiv preprint arXiv:2402. 14800, 2024.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Wei- jia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, al. Olmoe: Open mixture-of-cxperts language models. arXiv preprint arXiv:2409.02000, 2024.
Mohammed Muqeeth, Haokun Liu, and Colin Raffel. Soft merging of experts with adaptive routing. arXiv preprint MXXx:23366037445 2023.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vi- jay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale language model training on gpu clusters using megatron-Im. In Proceed- ings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15, 2021.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 1525-1534, 2016.
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse soft mixtures of experts. arXiv preprini 2rXiv:23008.000 00951, 2023.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero; Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 14:17555- 2021.
John Schulman, Nicolas Heess, Theophane Weher, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. Advances in neural information processing systems, 28, 2015.
Rico Sennrich. Neural machine translation of rare words with subword units. arXiv preprint
M.XVv11508.0799, 2015.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202. 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le. Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 1001::17700.65..... 2017.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv: S0ti::190000003.. 2019.
Snowflake.
Arctic open:
Efficient URL
foundation
language models
www htt.:/ae...c..b....//ago arceiic-oeen-effiiiiee------ otg.:e/feeeacttoooooooooooonnnna
snowllake, April 2024.
Chenyang Song, Xu Han, Zhengyan Zhang, Shengding Hu, Xiyu Shi, Kuai Li, Chen Chen, Zhiyuan Liu, Guangli Li. Tao Yang, et al. Prosparse: Introducing and enhancing intrinsic activation spar- sity within large language models. arXiv preprint 2rXiv:2402.135166 2024.
Jianlin Su. Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
13

---

Published as a conference paper at ICLR 2025
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive Network: A Successor to Transformer for Large Language Models, August 2023. URL URL/Attp://awxiv..oor haagg//as//2333 C8621 (XXi:2307.0862) [cs].
Shawn Tan, Yikang Shen, Rameswar Panda, and Aaron Courville. Scattered mixture-of-experts implementation. arXiv preprint arXiv:2403.0822 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Xun Wu, Shaohan Huang, and Furu Wei. Mixture of lora experts. arXiv preprint arXiv:2404. 13628, 2024.
Ted Zadouri, Ahmet Üstün, Arash Ahmadian, Beyza Ermis, Acyr Locatelli, and Sara Hooker. Push- ing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning. arXiv preprini wrXiv:330.005444 2023.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 4791-4800, 2019.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen. Sageattention2 technical report: Accurate bit attention for plug-and-play inference acceleration. arXiv preprint 0rXiv:2411.10958, 2024a.
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jun Zhu, Jianfei Chen, et al. Sageattention: Accurate 8-bit attention for plug-and-play inference acceleration. arXiv preprint arXiv:2410.02367, 2024b.
Xiaofeng Zhang, Yikang Shen, Zeyu Huang. Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. arXiv preprint arXiv:2210.05144, 2022.
Zexuan Zhong, Mengzhou Xia, Danqi Chen, and Mike Lewis. Lory: Fully differentiable mixture-of- experts for autoregressive language model pre-training. arXiv preprini urXiv:2405.03/33. 2024.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7 00100311441111 2022.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08900 08900. 2022.
Simiao Zuo, Xiaodong Liu, Jian Jiao, Young Jin Kim, Hany Hassan, Ruofei Zhang, Tuo Zhao, and Jianfeng Gao. Taming sparsely activated transformer with stochastic experts. arXiv preprint 10.0022 2021.
14

---

Published as a conference paper at ICLR 2025
A STABILITY ANALYSIS OF TOPK AND RELU
We introduce two metrics, "flip rate" and "flip count", to evaluate the routing stability:
Ei= veo(M M/D1)
flip rate
(14)
LTE
flip count = E x flip rate
(15)
where M Tx denotes the 0-1 mask matrix of the output of the router at layer and training step computed using fixed calibration set of tokens.
The metric "flip rate" represents the percentage of expert activation states that change (from active to inactive conversely) in single update, while "flip count" indicates the average number of experts whose activation states change.
We measure the two metrics MoE and ReMoE with N = =82M and E' {8, 16, 32} training for 1OB tokens. The results are presented in Figure I0, indicating that the ReLU router is more stable than the TopK router:
0.05
MaE
0.04
MoE_16E
MoE 16E
MoE 32E
0.3
MoE 32E
Flip Rate
0.03
ReMoE BE
Flip Count
Re BE
ReMoE_ 16E
ReMoE_1s 16E
0.2
ReMoE_322
ReMoE 32E
0.02
0.1
0.01
0.00
0.0
0
10
2
8
10
Tokens(B)
# Tokens(B)
Figure 10: Flip rate and flip count of MoE and ReMoF
When E' = 8, we find the flip rate of MoE is higher than ReMoE, though the gap narrows as training progresses and the learning rate decreases. While for E 16 and E 32, the flip rate of MoE remains consistently 2 3x higher compared to ReMoE throughout training.
Moreover, the flip count of ReMoE is invariant with respect to E, whereas the flip count of MoE is highly sensitive to the total number of experts and keeps increasing as E grows.
Notably, the flips in Topkerouted MoE are discontinuous (0.g. 0) - 0.51)), while those in ReLU-routed ReMoE are continuous(e.g(00.000 - (0, 0.01)), further underscoring the superior- ity of the ReLU router.
B INSENSITIVITY TO Å AND a
Ao
1e-16
1e-12
1e-5
1e-1
1.05
1.1
1.2
1.3
1.5
Valid Loss
2.031
2.029
2.032
2.036
2.032
Valid Loss
2.033
2.028
2.032
2.029
2.057*
Setting time
138
136
110
55
921
Settling time
414
211
110
80
52
1 Overshoot observed in 8-92 steps.
large oscillation amplitude spursity observed.
Table Valid loss and settling time for different values of A with 1.2.
Table 4: Valid loss and settling time for different values of with Åo = -
The ReMoE adaptation algorithm in Equation 7 includes two hyperparameters: A and a. Serrling time, defined as the total number of steps required in Stage I and Stage (as outlined in Section 3.5
15

---

Published as a conference paper at ICLR 2025
influenced by these parameters. For all experiments, that performance remains stable as long as A is small and a is close to 1.
Àu -e 1e- and 1.2, but we show
Our experiments with N 182M, E = 8, G 1, and & ReMoE models trained for 20k steps (~10B tokens) reveal only minor variations in validation loss for different Ag values (Table 3 and values (Table 4), except for 1.5 which caused rapid regularization changes and excessive oscillation. Besides, although different X and a values affect settling time, the impact minor compared the overall training steps, proving the insensitivity.
C PERFORMANCE FOR LONGER TRAINING
We conduct experiments of training MoE and ReMoE for a longer duration. We experiment with N =46M,, E 8, 1 and train the models with a batch size of 4M tokens and training over indicate that the superiority of ReMoE persists in
120B tokens. The results, as shown in Table 5 longer training.
Model
Valid Loss
ARC-c
ARC-e
BoolQ
HellaSwag
LAMBADA
PIQA
RACE
Avg.
MoE ReMoE
1.716
23.62
52.40
53.94 55.96
35.43
43.64
68.34
31.48
44.12
1.689
25.34
55.22
36.76
45.82
68.93
30.43
45.49
Table 5: Performance of training =469M, E 8,
models for 120B tokens.
D SPEED COMPARISON OF REMOE AND MoE
We measure the end-to-end training time for MoE and ReMoE with models of N =469M training over 120B tokens. The time consumption across stages summarized in Table 6 find Stage and Stage II account for 1.02% of the total training time and incur ~0.58% overhead.
Model
Stage
Stage II
Stage III
Total
MoE
0.12
0.41
119.12 119.25
119.65
ReMoE
0.32
0.91
120.48
Table 6: End-to-end training time comparison across stages (in hours). The time is measured on models training over 120B tokens.
N 469M, E 8, k
# Parameters
TP
Model
Train TFLOPS
Train Diff.
Infer TFLOPS
Infer Diff.
MoE
103.49
78.47
182M
11.82%
12.19%
ReMoE
105.38
80.19
MoE
138.58
107.52
469M
J1.37%
13.89%
ReMoE
136.69
111.71
MoE
160.46
153.11
978M
J1.77%
10.23%
ReMoE
157.61
152.76
MoE
133.40
118.55
978M
2
J0.68%
11.08%
ReMoE 132.49
117.27
MoE
103.61
85.96
978M
4
J2.29%
12.33%
ReMoE
101.23
87.96
Table 7: Throughput comparison between TopK-routed MoE and ReLU-routed ReMoE models. TP indicates the tensor parallel size. Train Diff. and Infer Diff. indicate the relative TFLOPS difference of ReMoE compared MoE, where 1 denotes ReMoE is faster, and I. denotes it is slower.
16

---

Published as a conference paper at ICLR 2025
We further measure the throughput of ReMoE against TopK-routes MoE across different model sizes and tensor parallel sizes during Stage III. The results, presented in Table 7 indicate that ReMoE achieves comparable training and inference speeds with MoE, with minor deviation ranging from 2.29% to +3.89%. This speed consistency desirable, as ReMoE introduces only a minimal modification to the standard MoE architecture by adjusting the routing function, thereby avoiding additional computational overhead.
E DOWNSTREAM EVALUATION RESULTS
This section provides the detailed downstream evaluation results for the main experiments of scala- bility of ReMoE in Section 4.3 and ablations on load balancing in Section 5.2
E.1 SCALING ACTIVE PARAMETERS N
The downstream evaluation results for scaling with respect to the parameter count N, discussed in Section are presented in Table These results highlight the performance comparison with increasing model parameters.
Model
N
ARC-c
ARC-e
BoolQ
HellaSwag
LAMBADA
PIQA
RACE
Avg.
182M
19.45
43.35
54.40
28.61
31.09
61.97
28.52
38.20
Dense
469M
21.50
49.12
56.88
31.12
36.74
64.47
30.53
41.48
978M
21.93
50.88
60.24
32.42
41.06
67.46
31.77
43.68
182M
20.82
45.03
57.55
29.84
31.81
63.28
28.42
39.53
MoE
469M
23,63
52.40
53,94
32.43
43.64
68.34
31.48
43,69
978M
23,81
52.90
58.90
35.01
44.42
67.90
31.48
44.91
182M
20.22
46.68
54.16
30.26
35.94
63.55
29.38
40.03
ReMoE
469M
21.67
53.16
58.75
33.80
40.66
67.95
31.20
43.88
978M
24.06
55.26
57.28
35.93
44.42
68.99
30.43
45.20
Table 8: Downstream results of scaling in active parameters N.
E.2 SCALING EXPERT COUNT E
Table the downstream evaluation results for scaling with respect the expert count E, as
examined in Section 4.3 This analysis illustrates how varying the number of experts influences the overall model effectiveness of MoE and ReMoE.
Model
ARC-c
ARC-e
BoolQ
HellaSwag
LAMBADA
PIQA
RACE
Avg.
Dense
19.45
43.35
54.40
28.61
31.09
61.97
28.52
38,20
4
20.73
44.49
59.63
29.14
31.40
63.33
29.19
39.70
20.82
45.03
57.55
29.84
31.81
63.28
28.42
39.53
16
20.90
45.29
46.36
30.50
33.22
64.96
28.33
38.50
MoE
32
19.54
47.35
52.29
31.12
35.63
64.25
28.23
39.77
64
19.88
46.63
60.06
31.47
36.33
65.07
28.04
41.06
128
20.99
47.69
56.73
32.00
36.62
65.67
28.04
41.10
4
19.88
46.46
57.43
29.64
33.57
62.95
27.66
39.66
8
20.22
46.68
54.16
30.26
35.94
63.55
29.38
40.03
16
20.90
49.28
53.36
30.85
37.09
65.83
30.05
41.05
ReMoB
32
20.56
48.11
59.54
31.42
37.84
65.18
28.42
41.58
64
20.82
50.51
57.80
32.17
36.74
65.78
27.46
41.61
128
19.97
51.05
56.97
32.40
37.92
66.70
29.86
42.12
Table 9: Downstream results of scaling in expert count E.
17

---

Published as a conference paper at ICLR 2025
E.3 SCALING GRANULARITY G
The downstream evaluation results for scaling with respect to the granularity G are shown in Ta- ble T0, based the experiments in Section 4.3
These results demonstrate the superiority of fine-
grained ReMoE over fine-grained MoE.
Model
G
ARC-c
ARC-e
BoolQ
HellaSwag
LAMBADA
PIQA
RACE
Avg.
Dense
19.45
43.35
54.40
28.61
31.09
61.97
28.52
38.20
Dense
22.78
48.11
59.66
31.11
35.65
65.02
29.57
41.70
20.82
45.03
57.55
29.84
31.81
63.28
28.42
39.53
2
21.42
46.55
54.25
29.95
32.52
64.09
28.61
39.62
4
20.99
46.09
55.90
30.52
35.16
63,98
29.28
40.27
MoE
21.59
47.73
60.70
30.83
36.41
64.69
28.04
41.42
16
19.80
48.82
57.34
30.64
36.00
64.74
28.71
40.86
32
21.67
48.78
57.85
31.27
37.10
64.69
28.52
41.41
64
20. 14
48.74
61.50
31.03
36.31
63.93
27.85
41.35
20.22
46.68
54.16
30.26
35.94
63.55
29.38
40.03
20.14
47.39
57.95
30.60
34,52
63.71
28.52
40.40
20,39
47.94
55.35
31.04
36.11
64.64
29.00
40.64
ReMoE
8
20.82
48.36
60.49
30.90
36.06
63.87
28.90
41.34
16
21.25
49.41
56.06
30.91
36.23
64.91
29.95
41.25
20.90
48.86
55.81
31.14
36.58
64.69
30.05
41.15
64
20.65
48.74
60.06
31.56
36.43
65.40
29.00
41.69
Table 10: Downstream results of scaling in granularity G.
E.4 LOAD BALANCING ABLATIONS
Table presents the downstream evaluation results for the load balancing ablations, discussed in
Section 5.2 These results compare performance with and without load balancing, offering insights
into the different roles of load balancing in MoE and ReMoE.
Model
LB
ARC-e
ARC-e
BoolQ
HellaSwag
LAMBADA
PIQA
RACE
Avg.
Dense
19.45
43.35
54.40
28.61
31.09
61.97
28.52
38.20
MoE x
19,20
44.74
50.80
28.60
30.18
62.24
27.94
37.67
MoE V
20.05
45.16
57.83
29.83
32.97
63.55
28.33
39.67
ReMoF
x 19.45
46.34
56.94
30.19
31.79
63.33
28.61
39.52
ReMoF.
20.22
46.68
54.16
30.26
35.94
63.55
29.38
40.03
Table 11: Downstream results of training with or without load balancing.
F DETAILED RESULTS FOR DOMAIN SPECIFICATION
Figure 11 shows the average routed tokens ratio of MoE and ReMoE across all layers. ReMoE demonstrates significantly stronger domain specialization compared MoE, where certain experts more frequently activated for specific domains. This suggests that ReMoE is better at learning and exploiting the unique characteristics of different domains, allowing it to allocate computational resources more effectively. In contrast, MoE exhibits more uniform expert activation across do- mains, indicating less differentiation its expert specialization.
18

---

Published as a conference paper at ICLR 2025
Arxiv
Books
C4
Github
Stack
Wiki
Layer
Layer
Layer
0.4
0.0
Layer
Layer
Layer
0.8
0.4
Routed Tokens Ratio
0.0
Layer 6
Layer
Layer
0.8
0.4
0.0
Layer
Layer 10
Layer
0.8
0.4
0.0
0
1
3
4
5
7
0
1
4
5
7
0
1
2
4
5
7
Expert ID
(a) Domain specialization of MoE
Arxiv
Books
C4
Github
Stack
Wiki
Layer
Layer
Layer
0.8
0.4
0.0
Layer
Layer 4
Layer
0.8
0.4
Routed Tokens Ratio
0.0
Layer
Layer
Layer
0.8
0.4
0.0
Layer
Layer
Layer
0.8
0.4
0.0
7
4
Expert
(b) Domain specialization of ReMoE
Figure 11: Detailed results of average routed tokens ratio for MoE and ReMoE in different domains.
19

---

Published as a conference paper at ICLR 2025
We further analyze the experts in Layer 5 of ReMoE and observe that certain highly related, domain- specific vocabularies are consistently routed to the same expert. To investigate this, we calculate the routing probabilities of different tokens based their IDs, defined the ratio of the number of times a specific expert is utilized to the total occurrences of the token. The results are summarized in T2
Our findings reveal that the vocabularies exhibit clear specialization, reflecting domain- char- acteristics. For example, Expert 1, which is more frequently assigned to natural language domains (e.g.. Books, C4), tends to route tokens such as husband, wife, and lover. In contrast, Expert 6, which is associated with non not-natural language domains (e.g,, Arxiv, Github, StackExchange), predominantly routes code-related tokens like variable, env, and HEAD.
Expert ID Routed Tokens With High Probability
End(100%); folding(100%): F111(100%) FILE(100%) NULI(100%);;
(el
huslband(1009)) i ife(100%) baby(100%); human(100%); 1over(99.60%); ((99.86%)) (98.425%)
invest(100%)) Fortu(())))))))rttttee((0(1))))) exec
bye(97.82%); [f(97.74%) (99.63%))
Conversely(1009)) Kethods(100%%) flower(100%); Blosscm(999933%))
(rgoapain(|00011anuuuaa/8/000aaa (100%)
Spring(100%) Summer(100%) Autumm(100%); Winter(100%):
2easonss(99.29 (e (oc,) cold(1005)
=(100%); 8(089))))
}) (100%); gif(())%); size(100%); variable(100%); env(100%);
begin((7.95%)) HEAD(97.945) ( (97.83%)
Kuala(100%); Tus(100%); Lama(100%); Riley(98.94%)
Table 12: Routed tokens with high probability for experts in Layer ReMoE
G DOMAIN-LEVEL DYNAMIC EXPERT ALLOCATION REMOE
We measure the average active expert count across different domains, as shown in Figure and find that the computation allocation in ReMoE also varies the domain level. Furthermore, this variation increases in deeper layers closer to the output. This is reasonable because deeper layers tend to capture more abstract and domain-specific features, leading to more pronounced specialization in expert activation.
Average Active Expert Count
Anxiv
Books
C4
Github
Stack
Wiki
3.0
2.5
2.0
1.5
1.0
0.5
0.0
11
Avg
Layer ID
Figure 12: Domain-level dynamic expert allocation
20

---

Published as a conference paper at ICLR 2025
H TRAINING MOE WITH NEAR-DENSE WARMUP
In ReMoE, the training process naturally progresses through three stages, with the first two involving near-dense training where the majority of experts active. To facilitate a fairer comparison, in Section 4.3 we train the MoE model for additional tokens to match the overall computational cost. In this section, explore an alternative approach by introducing similar near-dense warmup phase for MoE, referred to as "MoE with warmup, to align its computational footprint with ReMoE across each stage. Specifically, we train the MoE with N 182M, E 8, and k 6 approximately matching the average sparsity of ReMoE during Stages I and II, depicted in Figure 4a for the for the remainder of the training process.
first 100 steps, before transitioning
Table 13 compares this warmup variant to both standard MoE and ReMoE. The results indicate that the warmup phase provides a modest improvement validation loss compared to standard MoE, despite matching the overall computational cost. Nonetheless, ReMoE consistently outperforms both variants. This suggests that the shree-stage training pipeline learned by ReMoE, with Stages and II comprising only the first 100 steps, is beneficial to overall performance.
Model
Valid Loss
ARC-
ARC-
BoolQ
Hella-
LAM-
PIQA
RACE
Avg.
c
Swag
BADA
MoE
1.936
20.82
45.03
57.55
29.84
31.81
63.28
28.42
39.53
MoE with warmup
1.928
20.73
46.38
52.35
30.28
33.90
63.76
27.66
39.29
ReMoE
1.921
20.22
46.68
54.16
30.26
35.94
63.55
29.38
40.03
Table 13: Performance of MoE with near-dense warmup
We further extend our experiments with MoE using warmup to configurations with larger E, which increases the computational cost of near-dense training. The results, summarized in Table [4] show that as É increases, the warmup setting consistently improves performance. However, ReMoE still outperforms both variants, maintaining steeper performance scaling with respect to E.
Model, E' =8
Valid Loss
Avg. Acc.
Model, E' =32
Valid Loss
Avg.
Model, E' = 128
Valid Loss
Avg.
Acc.
Acc.
MoE
1.936
39.53
MoE MoE with warmup
1.874
39.77
MoE
1.852
41.10
MoE with warmup
1.928
39.29
1.869
40.06
MoE with warmup
1.841
41.34
ReMoE
1.921
40.03
ReMoE 1.852
41.58
ReMoE
1.815
42.12
Table 14: Results for MoE with warmup under different expert count E
To further investigate the impact of warmup steps on MoE performance, we vary the number of
warmup steps for the E 8 MoE configuration among 50, 100, 500, and 1000. The training curves of these models, along with standard MoE and ReMoE, are shown Figure [3] and the final validation losses are summarized in [5)
Our results reveal that performance does not improve monotonically with an increasing number of warmup steps, despite the additional computation. This behavior arises due to the discrepancy between the training objectives of & 6 (warmup phase) and & = 1 (post-warmup phase). For instance, when warmup concludes after 100 steps, the transition between phases is smooth, with the loss changing minimally from 6.491 - 6.751. However, extending warmup to 500 or 1000 steps leads to more pronounced loss gap of 3.101 5.827 and 2.695 4.428, respectively.
21

---

Published as a conference paper at ICLR 2025
MoE
ReMoE
2.4
MoE Warmup 50
Model
Warmup Steps
Valid Loss
Train Loss
MoE Warmup 100
0
1.937
MoE Warmup 500
50
1.930
2.2
MoE Warmup 1000
MoE
100
1.928
500
1.930
1000
1.931
2.0
ReMoE
1.921
0
5
10
15
20
25
30
Table 15:
Final validation loss of
#Tokens(B)
MoE with different warmup steps
Figure 13: Training curves of MoE with different warmup steps
In summary, near-dense warmup can enhance the performance of TopK MoE when training from scratch by providing better initialization for the experts. However, the warmup phase should conclude while the language model loss is still decreasing rapidly. Prolonging the warmup can exacerbate the gap between the warmup and subsequent training phases, ultimately degrading per- formance. In contrast, ReMoE naturally determines the appropriate warmup steps and sparsity levels due to its continuous and differentiable training dynamics.
FUTURE DIRECTIONS
This work be advanced in the following ways:
ReLU Routing for Mixture-of-LoRA. (MoLoRA), MoLoRA (Zaourri al. 2023 Wu ct al. 2024 Jiao et al. 2024) integrates MoE architectures to manage multiple Low-Rank Adaptation (LoRA) experts, dynamically activating task-specific adapters during infer- ence. ReMoE's fully differentiable routing mechanism could enhance MoLoRA by en- abling smoother transitions between LoRA experts, particularly when adapters trained on diverse tasks. Using ReLU straightforwardly in MoLoRA is explored in RoDE (Jiao et al. 2024), which be further enhanced by scaling the expert count while controlling the sparsity as in ReMoE.
ReLU Routing in Product- Product-Key-No Key-Meeor-r (PKM) Networks. PKM (ample et al. 2019
He 2024 Berges al. 2024 Huang al. 2024 architectures treat individual neurons
as ultra-tineggriieed experts, leading to routing complexity at unprecedented scales (e.g., millions of experts). ReMoE's differentiable routing and steep scaling properties par- ticularly suited to address PKM's optimization challenges.
Synergy with Efficient Attention Algorithms. Merging ReMoE's sparse, conditional feed-forward computation with efficient attention variants-such quantized (Zhang et al. 2024b a), linearized (Sun et al. 2023) Gu Dao 2023), sparse (Jiang et al.) 2025 Gao al. 2024), mixture-of-attemm
(Zhang et 2022] Csordás et al.) (025) mecha- fisms-could enable Transformers to scale efficiently in both sequence length and model capacity without incurring additional computational overhead.
Dynamic Expert Pruning for ReMoE. ReMoE's differentiable training inherently pro- motes expert specialization, with significant variance in expert importance across domains.
This property makes ReMoE more amenable expert pruning (Lu et al. 2024 Liu al. 2024a) compared to traditional TopK-routed MoE architectures.
22


Reviews:
[{"review_id": "5MPRf4WfEq", "reviewer": "Reviewer_VcWp", "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.", "questions": "- Why ReLU based router can enhance the domain specialization of MoE?\n\n- Whether the ReLU-based router can be further enhanced with shifted-ReLU, even with a learnable threshold?\n\n- How does the stage I/II affects the training performance?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Various ablation studies and visualization are conducted, analysing the effectivenss of the proposed methods.\n\n- 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.\n\n- The perplexity improvements over dMoE are not significantly. whether downstream performance, such as on common-sense reasoning tasks, would shows significant enhancement?\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GExcLftBBy", "reviewer": "Reviewer_qrQT", "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.", "questions": "One of the benefits of ReLU routing is that there is no inherent requirement for each token to be processed by $k$ experts. \n\n- Can ReMoE leverage the flexibility it is said to have with the regularization terms used?\n\n\n\n\n\n**Suggested writing changes:**\n- Scaling in parameters N --> Scaling in active parameters N", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n- The proposed method outperforms all other methods compared to in the study.\n- 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.\n- 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:\n\n*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?*\n\nA 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. \n\n\n**Other weaknesses**\n- 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?\n- 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).\n- 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.\n- I miss an evaluation of the performance of ReMoEs on LM evaluation benchmarks.\n- A batch size of 512k tokens is small for LLM pre-training. \n- 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?\n- 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.\n\n\n\n\n\n[1] SPARSE BACKPROPAGATION FOR MOE TRAINING, https://arxiv.org/pdf/2310.00811\n\n[2] GRIN: GRadient-INformed MoE https://arxiv.org/abs/2409.12136\n\n[3] Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models https://arxiv.org/pdf/2404.05567\n\n\n\n**I would be happy to raise my score if some concerns are addressed.**", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wPqLAmgFVA", "reviewer": "Reviewer_URhd", "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.", "questions": "See above.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\n2. Experimental results presented looks promising. It is interesting to see how a relatively simple idea work so much better. \n\n3. 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.\n\n1. 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.\n\n2. 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.\n\n3. 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.\n\n4. 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\nsteps (∼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. \n\n5. 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.\n\n6. 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.\n\nI am looking forward to additional experiments in rebuttal to update my rating.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "crks0KASa9", "reviewer": "Reviewer_yJh9", "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.", "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.\n- 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).", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "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. \n- The submission clearly illustrates the background information and the proposed method. \n- The performance gain by replacing the traditional routing with ReLU is interesting.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6ee3wLmu3j", "reviewer": "Reviewer_rQKz", "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.", "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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)