Layerwise Recurrent Router for  Mixture-of-Experts
Authors: Zihan Qiu
Zeyu Huang
Shuang Cheng
Yizhi Zhou
Zili Wang
Ivan Titov
Jie Fu
Venue: ICLR 2025 Poster (2025)
Keywords: Mixture of Expert
Router
LLMs
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce the Layerwise Recurrent Router for Mixture-of-Experts, efficiently establishing dependencies between routing decisions across consecutive layers and bringing consistent improvement.

Abstract: The scaling of large language models (LLMs) has revolutionized their capabilities in various tasks, yet this growth must be matched with efficient computational strategies. 
The Mixture-of-Experts (MoE) architecture stands out for its ability to scale model size without significantly increasing training costs. 
Despite their advantages, current MoE models often display parameter inefficiency. 
For instance, a pre-trained MoE-based LLM with 52 billion parameters might perform comparably to a standard model with 6.7 billion. 
Being a crucial part of MoE, 
current routers in different layers independently assign tokens without leveraging historical routing information, potentially leading to suboptimal token-expert combinations and the parameter inefficiency problem.
To alleviate this issue, we introduce the Layerwise Recurrent Router for Mixture-of-Experts (RMoE). 
RMoE leverages a Gated Recurrent Unit (GRU) to establish dependencies between routing decisions across consecutive layers.
Such layerwise recurrence can be efficiently parallelly computed for input tokens and introduces negotiable costs.
Our extensive empirical evaluations demonstrate that RMoE-based language models consistently outperform a spectrum of baseline models. 
Furthermore, RMoE integrates a novel computation stage orthogonal to existing methods, allowing seamless compatibility with other MoE architectures. 
Our analyses attribute RMoE's gains to its effective cross-layer information sharing, which also improves expert selection and diversity.

Published as a conference paper at ICLR 2025
LAYERWISE RECURRENT ROUTER FOR MIXTURE-OF-EXPERT
Zihan Qiu Zeyu Huang" SShaang Cheng
'Yizhi Zhou
Ziii Wang
2,6 Ivan Titov 'Jie Fu'
1 Alibaba Group, ?University of Edinburgh, SICT, Chinese Academy of Sciences,
*Nanjing University, IINF Technology UUniversity of Amsterdam "Shanghai Lab qzh1
6288 gmail. zeyu. huang@ed ed uk, fujie@pjlab.org. org.c
ABSTRACT
The scaling of large language models (LLMs) has revolutionized their capabili- ties in various tasks, yet this growth must be matched with efficient computational strategies. The Mixture-of-Experts (MoE) architecture stands out for its ability scale model size without significantly increasing training costs. Despite their advantages, current MoE models often display parameter inefficiency. For in- stance, pre-trained MoE-based LLM with 52 billion parameters might perform comparably to standard model with 6.7 billion parameters (Rajbhandari et al., 2022). Being crucial part MoE, current routers in different layers indepen- dently assign tokens without leveraging historical routing information, potentially leading to suboptimal token-expert combinations and the parameter inefficiency problem. To alleviate this issue, introduce the Layerwise Recurrent Router for Mixture-of-Eaae -fpeExs (RMoE). RMoE leverages a Gated Recurrent Unit (GRU) establish dependencies between routing decisions across consecutive layers. Such layerwise recurrence can be efficiently parallelly computed for input tokens and introduces negotiable costs. Our extensive empirical evaluations demonstrate that RMoE-based language models consistently outperform a spectrum of base- line models. Furthermore, RMoF integrates novel computation stage orthogonal to existing methods, allowing seamless compatibility with other MoE architec- tures. Our analyses attribute RMoE's gains to its effective cross-layer informa- tion sharing, which also improves expert selection and diversity. Our code is at nttps://githinn com/ qwwzh2/ /RMoE.
INTRODUCTION
In the era of large language models (LLMs), scaling the model parameters and training data up has unlocked remarkable model capabilities, such as in-context learning (Brown et al., 2020; Dong et al., 2022), nuanced conversations (Ouyang et al., 2022), and even complex code (Guo et al., 2024) and math (Imani et al., 2023) tasks. These advancements showcase the profound impact of increasing model size. The quest to enhance neural networks' capacity while ensuring training and inference efficiency spurred the development of computation-efficien transformer architectures. The Mfxture-of--xpeent (MoE) framework is of such efficient architectural recipes (Shazeer et al., 2017: Lepikhin et al., 2021; Fedus et al., 2022: Zhang et al., 2022; Dai et al., 2024). Most MoE modules comprise one router and group of expert networks. The router, usually parametrized as one linear layer, conditionally and sparsely assigns each input token to its corresponding experts, i.e., the FeedForward Network (FFN) in the transformer layer. Therefore, MoE can significantly scale the model size and keep computational costs nearly unchanged (Smith et al., 2022).
Despite efficiently increasing the model size, most current pre-trained MoE models are not on par with standard models of the same size, demonstrating their parameter inefficiency. For example, Rajbhandari al. (2022) shows that with the same training data, an MoE with 52B parameters and 1.3B activated ones for each token performs similarly to a 6.7B standard model. Komatsuzaki
Equal contribution
1 "ooresppnniigg author

---

Published as a conference paper at ICLR 2025
Hidden forward
Experts
GRU fonward: Share GRU
Experts
Expert scores forward
X
E
E
Hidden
Hidden
Proj
Proj;
gi
1
g;
Router
Router
h;
Shared GRU
Shared GRU
Layenwise
Layer 1
recurrence
Layer i
Figure 1: Recurrent router for Mixture-of -f-eeetsss.......eoo In the i-th layer, the hidden state x; is I. projected to x' with alower hidden dimension (Eq. 4), II. combined with previous layer's GRU output h;-1, and processed through the cross-layer-shared GRU to produce the current layer's GRU output, h; (Eq. 5). III. layer i's router uses this output to select experts and exccutes standard MoE computation (Eq. 6). Such operation doesn't t introduce sequence-level recurrence and can be efficiently implemented, as shown in Tab. and Tab. 3.
et al. (2023) demonstrates that upcycling a standard T5-base (248M) into its MoE counterpart (2B) by copying existing FFN can bring some improvements, but it still lags behind the T5-large with 783M parameters. Similarly, Dai et al. (2024) use fine-grat fine-grained and shared experts improve the effectiveness, but the 16B MoE performs comparably with the 7B standard model (Bi et al., 2024).
One potential bottleneck for the current MoE could be the router. Typically, the router is parame- terized as one lightweight linear layers, which may limit its capacity to explore the optimal token- expert combination. Previous works also reveal such limitations. For instance, Xue et al. (2024) finds the routing results converge to the token-id-base of pre-training, which means the token-expert combination is far from well-explored. Some works even show hash functions (Roller et al., 2021), stochastic routing policy (Zuo et al., 2021), and ixee-random router (Chen et al., 2023) achieves competitive performance with the learnable router, illustrating that the learnable router component in MoE needs further enhancement.
routing very quickly during the early phase
Despite some enhancements for router (Chi et al., 2022; Shen et al., 2023; Do et al., 2023; Chen et al., 2023), current routers in different MoE layers still operate independently without comprehen- sive investigations into the decisions of other layers. This isolation may lead to suboptimal expert utilization, each layer manages its routing based solely on local information, potentially leading to inefficiency of model parameters. Though vanilla MoE models could technically share the rout- ing information via hidden states residual, this information may be overshadowed by the language modelling loss, requiring routing-relevant information to "compete" for its representation.
To this end, we introduce dedicated component to capture and pass routing information for each layer. The proposed architecture, Recurrent Router for Mixture-of- Experts (RMoB), is shown in the Fig. 1. Concretely, we regard routing decisions in consecutive layers sequence in which the routing results of the i-th layer should be conditioned on previous layers' decisions. We thus intro- duce lightweight Gated Recurrent Unit (GRU) (Dey & Salem, 2017) to capture this dependence and simulate the information flow between routers across layers. Intuitively, GRU has a reset and an update gate to control the information flow across time steps. Hence, such layerwise recurrence will inform the router to which experts the current token was assigned in previous layers, potentially supporting cross-layer collaborations. Furthermore, the introduced GRU is especially for routing. thus helps to disentangle the states relevant to model prediction and routing decisions.
We validate RMoE's performance with various model sizes, architectures, datasets, and training settings (per-training and supervised fine-tuning), demonstrating that RMoE outperforms range of baselines. Moreover, RMoE's introduction of a novel computation stage during routing makes it orthogonal to and compatible with most existing methods. We further analyze RMoE and elu- cidate the primary contributors to its improvement. Our findings indicate that while the GRU in RMoE shares essential cross-layer information, it also enables additional gradient propagation for the router. Our analysis shows that layerwise recurrence provides cross-layer information, fostering router exploration and optimizing expert utilization. Consequently, the selected experts are lever- aged more effectively, leading to increased diversity of experts. We believe that our innovative router design and massive analysis can offer insights into the development of future MoE models.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORKS: VARIOUS ROUTERS FOR MOE
In this section, we review previous approaches to improve router design in SMoE. For example, XMoE (Chi et al., 2022) first projects hidden states into lower-dimension space and computes their cosine-similarity to low-dimension expert embeddings, which can prevent the hidden states from collapsing to linear combination of expert embeddings. Moduleformer (Shen et al., 2023) uses an MLP router with ReLU activation to increase router capacity. SMoE-dropout (Chen et al., 2023) utilizes fixed
linear router and gradually increases Top-k during training. HyperMoE (Do et al., 2023) introduces fixed random- initiaalizee hypemeet (Ha et al., 2016) at each layer to generate router weights condition on input and one learnable router embedding. One concurrent work (Gong et al., 2024) also introduces GRU in sequential routing stages. However, it does not view such a recurrent mechanism as a general and composable method with broad MoE fields provide relative ablation analysis. Extra discussion of related work to improve MoE from routing and training strategies, and utilize recurrent controllers can be found App. A.
METHODOLOGY
3.1 PRELIMINARIES
Mixture-of-Experts MoEs are typically implemented by replacing transformer models' original forward networks (FFNs) with a group of parallel FFNs and incorporating a router. Suppose there are N experts, denoted as En E N]. The router g(: G, k),, defined by its parameters G e (R,,, and an integer k maps the input to score distribution over the experts, g(x; G,k) R.N. Given x E Rd, the output y E IR'h is the weighted sum of the outputs from all experts:
Σ (x: G, E,)
(1)
REN
Typically, is a simple linear layer followed by softmax and a Top-k function. The n th element x G E R. represents the gating score of expert E,,, and the th column of G can be regarded as the expert embedding for expert E,, When & for Top-k is smaller than N, only a subset of experts is involved in the computation, which is known Sparse Mixtureoof-Experts (SMoE) (Shazeer al., 2017; Fedus et al., 2022).
Recurrent Neural Networks RNNs (Medsker et al., 2001) are designed to handle sequential data by maintaining a hidden state h that holds the information from previous time steps. This hidden state is updated at each time step i based on the current input x and the hidden state at the last time step 1;-1, formulated as h, f(b:-1,
The Gated Recurrent Units (GRU) Dey Salem (2017) module an advanced variant of RNNs that addresses traditional RNNs' limitations, such as difficulty capturing long-term dependencies and gradient vanishing issues. Given an input x/ at time step i, GRU first calculates the reset gate s and the update gate to determine how much of the previous memory to keep and to forget,
8; (W,x) ,x! + U,h.-I), 2, = oW X + I.)
(2)
where a represented the sigmoid activation function and all W and U are tranable parameters. And then, the hidden state ht updated by
tanh W X S; O (Wah;-)); h
h (1 2, h + 2, h
3.2 LAYERWISE RECURRENT ROUTER
Existing routers work independently, this lack of global information may prevent routers from dis- covering more effective token-expert combinations. Therefore, we integrate GRU into the routing process, explicitly incorporating historical routing information into the current expert selection for each token. Formally, at the th layer, we first use linear layer to project the hidden state x; the dimension of the GRU state x E I²' (usually smaller than the dimension h of x,. We choose 128 for most of the settings provide further analysis in Tab. and Tab. 7):
x Proj,(x.)
3

---

Published as a conference paper at ICLR 2025
Importantly, we separate projectors for each layer since the hidden states x of different layers vary greatly (more discussion in Sec. 5). This projection output x', along with the GRU result from the previous layer, 1 is then fed into GRU unit to obtain the current GRU output h,
h: G h5-1)-
(5)
Next, h; is input into the router and then expert outputs are aggregated based on the router output:
Yi =
Σ (n,; G2 (x(())) REN
(6)
Here, y; represents the output of the i-th layer, h; is the GRU output, (h;; Gi, k) is the router output computed with routing parameter Gj in layer i. Notice that, unlike traditional RNNs, which use a shared projector together for sequential inputs when the input dimension isn't 't equal to the RNN's hidden dimension, we use different projectors Proj, in Eq. for different layers since hidden states and model weights in different layers usually various a lot (Fig. 11 and Tab. 6).
Despite capturing inter-layer dependencies between routers in different layers, RMoE potentially has other advantages: (1) Preveni representation collapse: Chi et al. (2022) identified that the single linear layer routers encourage token embeddings clustering around expert embedding. implying a trend toward representation collapse issue. And they propose XMoE to first project hidden states into a low-dimension and then calculate the gating score. Similarly, the projector (Eq. 4) and GRU (Eq. 6) in RMoE also separate hidden states from expert embeddings and can reduce this issue. (2) Additional Gradient Flow: Before the inclusion of GRU, the router's gradient mainly derive from the expert weight score In in Eq. 1. The introduction of GRU not only provides enriched information about historical routing but also an extra gradient propagation through GRU hidden states. We denote this extra gradient flow as Recurrent Gradient, and we empirically demonstrated that this Recurrent Gradient is important to RMoE. (3) Applicable with other MoE design: the proposed method introduces an additional computation stage into SMoE, it is orthogonal to most existing attempts to improve MoE and is seamlessly compatible with them.
4
EXPERIMENTS
4.1 EXPERIMENTAL SETTINGS
Langauge Modeling Tasks and Metrics Following (Pham et al., 2024), we first test on two com- mon language modeling tasks: enwiki8 (character-level language modeling, with Bits-Per-Character (BPC) as the evaluation metrics) and WikiText- 103 (word-level language modeling, with Perplexity (PPL) as the evaluation metries). We employ default mein-validation-tee splits for each dataset. We report test performances of the best validation checkpoints. More details can be found in App. A.2.
Configurations and Baselines We compare RMoE with other existing router designs. All meth- ods are based on the decoder-only standard writch-transformer architecture with post-norm. Fol- lowing (Pham et al., 2024), all routers select top-2 experts from 16 experts. Each task is trained on 2 NVIDIA A100 GPUs for about 20 hours. More training configurations can be found in App. A.2.
Our baselines include (1) SMoE: standard switch-transformers with a standard linear router. (2) HyperMoE (Do et al., 2023): the method employs a fixed, randomly initialized hypernetwork (Ha et al., 2016) to produce the weights for the linear router, subsequently allowing the generated lin- car layer to perform the routing. (3) SMoE-MLP (Shen et al., 2023): it replaces the linear router with a two-layer MLP using the GELU activation function. (4) RandomMoE: inspired by SMoE- Dropout (Chen et al., 2023) and HyperMoE, we propose to compare with a fixed randomly initial- ized linear router; this could be naive baseline for all learnable routers. (5) XMoE (Chi et al., 2022): it first down-projects the token embeddings to a lower dimension (default 16) and computes its cosine-similarity with the low-dimension expert embeddings. It also uses a learnable temperature in softmax. (5) CosineSMoE, similar to XMoE except without down-projection.
Pre-Training and SFT paradigm As pre- pre-training- then thhn-supereieed fiw.riiee--i---tt-nnngg has become the standard paradigm, also evaluate the RMoE in this setting. We conduct preliminary scale-up experiment on a setting of training 0.91B models with 40B tokens. Our pre-training corpus is a mul- tilingual data collection that spans common and specialized domains, including Wikipedia, finance,
4

---

Published as a conference paper at ICLR 2025
Table 1: Performance of RMoE and baselines on two language modeling tasks, Enwiki8 and WikiText-103.
Params means the non-embedding model parameters and (router parameters). Notice we don't t separate un- learnahle parameters in HyperMof: and RandomSMoF. Mem means the peak GPU memory usage with the same batch-size configurations. Speed is the average time for training steps. Results demonstrate that the RMoE outperforms baseline models and achieves comparable memory usage and speed as the standard SMoE.
Enwiki8 (BPC).
WikiText (PPL.)
Params
Mem
Speed
Algorithm
val
test
val
test
(M)
(GB)
(s/1k steps)
SMoE
1.152
1.128
31.279
33.061
36.08 (0.04)
47.92
960.2
HyperMol
1.162
1.139
31.918
33,374
48.41 (12.4)
49.69
962.0
SMoE-MLP
1.164
1.137
31,430
33.142
36.79 (0.75)
48.70
964.1
MandomSMoE
1.163
1.135
31.938
33.410
36.08 (0.04)
47.72
961.4
MosineMoE
1.148
1.122
31.367
33.047
36.08 (0.04)
48.68
962.4
XMoE
1,150
1.125
31,265
32.926
36, (0.09)
48.70
967.5
RMoE
1.141
1.116
30.939
32.867
36.51 (0.47)
49,46
972.9
and legal texts. Our model architecture is modified based on Llama family (Touvron et al., 2023). Specifically, we use 24-layer model and top-4 gating from 16 experts per layer following (Dai et al., 2024). This yields a model with approximately 0.53B activated / 0.91 total parameters. All different routers use the same training configurations. To ensure expert load balance, we em- play balance loss with weights 0.01 during training. These experiments are conducted using the Megablocks (Gale et al., 2023) on 8 NVIDIA A100 GPUs for about 5 days. More details can be found in App. A.2. After pertaining, we perform supervised fine-tuning (sft). All models are trained on Alpaca (Taori al., 2023) with the same configuration. We use im-evaluation-harre harness' to evaluate the tnn-ttund model. To simulate the real LLMs application scenario, we don't ' perform task-specific line-tuning and evaluation. Since the models are largely under-trained, they give almost random-guessinn guessing results on challenging tasks like MMLU (Hendrycks et al., 2020). Therefore, we only test on tasks (ARC-easy, Hellaswag, PIQA, SciQ. LAMBADA) in am--vvlluaionn-hare -hanness... More details about sft configurations and tasks can be found in App. A.2. We further justify the scalability of RMoE on the setting of training 15B activate 2.7B models with 120B / 400B tokens. Given our utilization of a high-quality pre-training corpus, pre-training on 400B tokens yields better results compared to experimental MoE like OpenMoE (Xue al., 2024). We find RMoE consistently pro- vides over one-point improvement in performance on benchmarks such as MMLU, GSM8K, and HumanEval. More details can be found in App. A.3.
4.2
MAIN RESULTS
Tab. shows the performance of RMoE and selected baselines on two language modelling tasks. Our ob- servations as follows: (1) RMoE performs best on validation and test sets of two tasks, and the re- current routing mechanism and the introduction of extra GRU block do not severely impact the train- ing speed and memory usage, making RMoE more practical. (2) Comparing SMoE-MLP and SMoE, we find that replacing the original simple linear layer with a more capable MIP does not improve perfor- mance. even underperforms the fixed random routing (RandomMoE) on Enwikik8, suggesting that naively increasing model capacity can't 't result in more powerful router. Furthermore, since RMoE introduces novel computation stages in routing and is orthogonal to most existing router de- signs, it can easily be combined with them. Tab. 2 showcases the performance of the original XMoE and XMoE with GRU router in different XMoE lower dimensions (8, 16, and 32). We observe that the GRU router benefits all of the 3 configurations of XMoE.
Table 2: Performance of combining layer-wise recurrent routing mechanism with XMoE.
Enwiki8 (BPC).
WikiText (PPL).
Algorithm
Val
test
val
test
XMoli
1.160
1.132
31.74
33.55
GRU router
1.150
1.124
31.34
32.99
XMoE (16)
1.150
1.125
31.27
32.93
GRU router
1.144
1.119
31.15
32.47
XMoB (32)
1.140
1.114
31.30
32.71
GRU router
1.136
1.112
31.25
32.55
While previous work on improving routers has not mostly been evaluated on large-scale pre- training (Dai et al., 2022; Chi et al., 2022; Do et al., 2023), we scale up RMoE to billion-level parameters and training tokens. We report SMoE and RMoE's evaluation results (both directly eval- uated and evaluated after supervised fine-tuning (sft)) in Tab. Existing works suggest freezing the router during SMoE tuning Zoph et al. (2022); we report SMoE's results under freeze and unfreeze settings. Correspondingly, for RMoE, we freeze the GRU and the linear layer under the freeze setting. From Tab. 3, we can observe that (1) Even in large-scale pre-training that requires more
5

---

Published as a conference paper at ICLR 2025
SMoE and RMoE's pre-training costs and evaluation results in selected informative Im-evaluation- harness tasks. 'sft' means supervised fine-tuning the Alpaca dataset. The task names and metrics for short names in the table are: 'ARC-e" for ARC-Easy, acc; 'Hella' is for Hellaswag, acc-norm; PPqqa" for PIQA, ace-norm; 'Lamb' for LAMBADA, acc. Each model has approximately 0.53B activated parameters out-of 0.91B parameters. RMoE introduces about 3.5M additional parameters relative to SMoE.
Table 3:
Algorithm
Training
ARC-e
Hella
Piqa
Seiq
Lamb
Avg
pre-train 20B tokens
47.14
35.51
64.69
76.2
14.61
47.63
SMoE
+sft
50,93
35.82
65.61
74.7
17.81
48,97
+sft (freeze router)
50.59
35.78
66.32
74.7
18.18
49.11
S Speed: 48.87 sistop Mem: 48,00 GB
pre-train 40B tokens
52.57
40.85
67.74
83.4
26.74
54,26
+sft
53.70
42.07
68.61
83.5
32.80
56.13
+sft (freeze router)
53.45
41.94
68.88
83.1
32.06
55.88
pre-train 20B tokens
47,01
35.91
65.23
78.7
19.13
49.20
+sft
48.53
36.90
66.21
79.6
24.74
51.20
RMoE
+sft (freeze router)
49.24
36.79
66.16
79.7
24.32
51.24
pre-train 40B tokens
51.18
41.38
67.79
83.6
32.58
55.31
Speed: 49.07 sistep Mem: 48.69 GB
+sft
53,20
43.05
68.55
83.8
37.16
57.15
+sft (freeze router)
53.11
43.16
68.77
82.8
37.57
57.08
complex parallel training strategies, RMoE brings negligible wall time and memory cost compared with vanilla SMoE. (2) In comparable settings (e.g., the same number of tokens and with/without sft), RMoE outperforms SMoE, and even the best results of SMoE are lower than those of RMoE.
ABLATION STUDIES
Which contributes more?
Table 4: Enwiki8 validation and test BPC for different routing designs. 'NP' stands for not passing recurrent states cross-layer. 'RMoE+NP' has the same parameters and FLOPs as 'RMoE'.
More Router parameters or lay- erwise recurrence. A straightfor- ward reason for RMoE improvement could be that RMoE introduces ad- ditional computation and parameters. To disentangle the effect of introduc- ing more router parameters and layer- wise recurrence, we consider the fol- lowing two extra settings: (I) SMoE+MLP: we naively increase the router parameters by replacing the original linear layer with larger MLP layer, (2) RMoE NP: we change Eq. 5 to GRU ho) to cancel the layerwise recurrence of RMOE, rendering a stateless GRU. The setting has the same parameters and computation as RMoE. From Tab. 4, we can observe that (1) in our setting, introduc- ing larger routers in SMoE doesn't it bring improvement (SMoE v.s. SMoE MLP). (2) When ablated on the layerwise recurrence in RMoE, the performance largely drops, even worse than SMoE. Both results suggest that the layerwise recurrence is the main contributor.
Algorithu
Vall
Test.
Paras (M)
Val.
Test..
Paras (M)
Small
Medium
SMoR SMoE + MLP
1.214
1.184
15.32
1.152
1.128
36.08
12214
1.183
15.73
1.164
1.137
36.79
RMRE+ NP
1.227
1.196
15.61
1.150
1.123
36.51
RMoE
1.213
1.883
15.61
1.141
1.116
36.51
Recurrent Gradient is important to RMoE Following the aforementioned analysis, we try to further disentangle the effect of the layerwise re- currence. When removing the layerwise recur- rence as in the RMoE + NP setting, we remove two information flows across layers: (1) the for- ward information about previous routers' deci- sions and (2) the backward gradient propagation through GRU hidden states in different layers. To compare the two information flows, we investigate
Table 5: Enwiki8 validation and test BPC. 'de- tach h--1' means detaching the recurrent hidden states before passing it to the next block. 'r-0.5/1.0 /1.0 means passing the routing logits of the previous block the current block. 'detach-r' means detach- ing the gradient computation of passed logits.
Algorithm
Val
Test
SMoE
1.152
1.128
RMoE
1.141
1.116
+ NP
1.150
1.123
the following settings: (1) RMoE + detach h
detach h4-1
1.159
1.133
in intermediate stage between RMoE and RMoF- NP. By detaching h to stop its gradient compu- tation in Eq. 5, each GRU cell can only use previ- ous information during feed-forward. (2) RMoB
NP + r-005
1.149
1.124
NP + r-1.0
1.150
1.124
NP + r-0.5 detach-r
1.157
1.133
NP + r-1.0 + detach-r
1.152
1.126
+ NP + r-a: inspired by Realformer (He et al.,
2020) that introduces residual attention score to facilitate attention gradient back- mack-propagatio,, investigate an intermediate stage between RMoE and RMoE-NP by adding gating logits resid- ual for the RMoE NP settings. Concretely, the gating score of i-th layer for expert is

---

Published as a conference paper at ICLR 2025
In (be; G;, k) + agn. (1; ( It is a straightforward way to supplement router information across layers based on the NP setting. In our experiments, we set a as 0.5 and 1.0. (3) Moreover, we also test detaching the gradient computation of passed logits (h G,-1), denoted as 'detach-r'. From Tab. 5, RMoE+ + detach h4-1 performs even worse than RMoE-NP, showing that the Recurrent Gradient is important. Similarly. 'NP+00.5' and 'NP+r1.0 .0' are comparable with 'NP", showing that the naive gating score residual can't provide effective cross-layer information. The performance of their detached version largely drops, demonstrating the importance of extra gradient passing.
To further validate the gradient passing hypothesis, we test 'NP' and 'NP-r0.5/ /.0" on deeper models. The re- sults are summarized in Fig. 2. As the layer increases, we can observe that (1) RMoF consistently outperforms other settings, and RMoE-NP even lags behind SMoE. The possible reason is, without passing recurrent states, RMoE-NP is similar to SMoE-MLP which simply in- creases router complexity but doesn't refine the router training. (2) RMoE-NP-rO.5 surpasses RMoE-NP, fur- ther emphasizing that SMoE's optimization benefits from the added additional gradient flow for routers. The spirit echoes the principles behind residual network, where residual connection are used to create direct paths for gra- dient propagation, thereby mitigating gradient vanishing
Test BPC Layer
1.20
SMoE
RMGE
1.20
RMLE MP
AMUE NP-1..
1.00
Moda Layer
Figure 2: Test BPC on Enwiki8 with differ- ent model sizes (6, 12, 18. 24. 32). Similar validation results are in App. A.5 Fig. 14 as lerys deepen. Similarly, the GRU and the direct logits passing help for gradient flow of routers in deep layers. Ad shown in the Fig. 2, as the layer increases. the performance gaps between them may becomes more significant (3) While providing additional gradient across layers, RMoEE-NP-r055 un- derperforms RMoE. This may because the indexes of experts in layer are not aligned with those in other layers, directly adding logits can lead to improper constraints and hurt the model performance, further highlighting that RMoE adds flexible while informative pathways in the SMoE framework.
Table 6: Ablation of RMoE design. 'L-proj" means the layerwise projector in Equ 4, 'SSproj" is the standard RNN projector. 'SMoE + L-proj GRU router' is our proposed used RMoE method.
Enwiki8 (BPC)
Wiki'Tex (PPL)..
Algorithm
Params
Val
test
SMoE
1.152
1.128
31.28
33.06
36.08
SMoE + L-proj + GRU router
1.141
1.116
30.93
32.86
36.50
SMoF + S-proj + GRU router
1.148
1.123
31.15
33.02
36.23
SMoF + L-proj + RNN router
1.145
1.119
31.18
32.72
36.44
SMoE + L-proj + LSTM router
1.148
1.122
31.19
33.04
36.54
Layerwise projector and suitable recurrent net bring the best results. This part tests the other components in RMOE, such as recurrent hidden state dimension, lay-
Table 7: Ablation the recurrent design on large scale per-training setting. p is the dimension of the recurrent state r, in Eq. 4. We report averaged tasks (the same as Tab. 3) results for pre-trained and stf models.
erwise projector, and GRU cell. As shown in Tab. 6: (1) All methods with recurrent routers outperform SMoE. (2) Layerwise projector in Eq. 5 performs better than stan- dard RNNs using single shared projector. One possible reason is that the weights and hidden states norm in differ- ent layers vary greatly (as shown in App. A.4.5 Fig. 11),
All models are trained with 20B tokens.
Algorithm
Pretrain
+sft
SMoE
47.63
49.11
RMoE (GRU. p -22)
49.20
51.32
and it would be hard for single shared projector to pro- cess them. This approach aligns with the design princi- ple of not sharing LayerNorm parameters when employ- ing shared MoE transformer blocks, as discussed by Xue et al. (2022). (3) The GRU router performs best. Moreover, we further compare RMoE variants in the larger scale settings. We compare pre-trained models with different structures and recurrent hidden dimensions in Tab. 7 (Averaged results, full results in App. A.5 Tab. 12). We can find similar results: (1) All RMoE variants outperform SMoE; (2) Simple router (RNN) and complex routers (GRU with p = 256, 512) perform worse. In short, layerwise projector and moderate recurrent cell (e.g. GRU with p= 128) effectively introduce layerwise recurrent.
RMoE (GRU, p= 226)
49.08
50.04
RMOE (GRU, p= 512)
49.19
50.02
RMOE (RNN, 256)
47.92
50.44
7

---

Published as a conference paper at ICLR 2025
Lay0
64
64
ASS
CSC
6.6
CGB
AGSA
6G
CUB
CSr
SPM
5:5M
CGR
622
6:44
S4F
CON
GET
412
Layer
Layor
Layer
[a] Mean MI SMoE
(b) Mcan Ml XMM
[c] Mean Mf HyporMoE
4.40
4.0
3PM
123
0.50
5/10
500
XMI CPM
16
40
Mean MI RMoE
(e) Mean Mi RMGE
Mean MI of RMoE MP-12.0
Figure 3: Heat maps of cross-layer mutual information (MI) for different methods. The (i-th row, j-th column) value represents MI between layers i and j. The First Row ((a) SMoE, (b) XMoE, (c) HyperMoE); All three methods have low cross-layer MI. Second Row(d)) RMoE, (e) RMoE-NP, (f) RMoE-NP-r1.0): While RMoE has high cross-layer MI when disabled layerwise recurrent states passing, MI largely drops.
6 OBSERVATIONS
Layerwise recurrence increases cross-layer mutual information. The intuition of the proposed RMoE is that current routers in different layers are isolated, and the layerwise GRU is incorporated to provide routers with global information for coordination. Therefore, we measure the Mutual Information (MI) between routing distributions in different layers for cach router in Fig. 3. The code can be found in App. A.4.2. We can observe: (1) Besides RMoE, all existing methods show low cross-layer MI, indicating that the routers of different layers work relatively independently. (2) RMoE shows higher MI than three baselines (d v.s. a, b, and c) and RMoE-NP (d v.s. e). showing the recurrent router can facilitate cross-layer information sharing. (3) While RMoE-NP's MI is largely smaller than RMoEE it still surpasses the three baseline methods. The reason can be the shared GRU in Eq. 5. (4) Intuitively, passing routing logits can directly improve MI (f v.s. e). However, directly passing logits can't ensure long-range information sharing, as the values in the right part of (f), which indicate the MI between non-neighbor layers, are smaller than those in (d).
RMoE enables moderate flat gating scores. The router's gating score is a noteworthy feature for MoE-based models. It showcases the models' training dynamics and how they ultimately exploit their experts. Ideally, the training paradigm of MoE models may have two stages: exploration and then exploitation. i.e., the router should actively explore more new expert combinations at the early stage of learning. But if the gating score converges to a sharp distribution too early, the router will learn very shallow routing mechanisms and fail to find optimal routing decisions. So We record gate entropy for each token(- (E In In 9n). In is the gating score for expert n)) and plot the entropy distribution in Fig. (left). Generally, the higher the entropy. the more evenly the router activates different experts rather than allowing one expert to dominate the layer. Thus, large density in high- entropy parts means many recorded tokens have flat gating score distributions. We can observe that (1) RandomMoF, with fixed most tokens have high entropy, as there is only one peak in the large entropy location. This indicates while RandomRMoE can highly encourage exploration, the router may be under-trained and lack exploitation. (2) SMoE and HyperMoE show low routing entropy, with many tokens having nearly zero entropy. Such low entropy means the softmax operation gives nearly one-hot results, which means the Top-k experts degrade to Top-1 and the router's 's gradient are very sparse. This can hurt the exploration of expert selection and lead to inefficient Top-k experts usage. (3) XMoE and
router, shows the largest gate entropy. Moreover,
8

---

Published as a conference paper at ICLR 2025
CosineMoE, using cosine similarity, which normalized the input and weights G before computing logits, show relatively high entropy. They also perform better than SMoE in Tab. 1, indicating the benefits of suitable exploration. (4) RMoE, with unique cross-layer information sharing, has high entropy for many tokens while low entropy for a few tokens. These moderate gating scores can achieve better balance between exploration and exploitation.
One may argue that such high entropy may come from the under-trained recurrent router in RMoE instead of capturing the dependency across layers, the unlearnable RandomMoB also gives high entropy. Therefore, we further visualize the scores of 'RMoB-NPP' and 'RMoE-NP-rO.55 0 in Fig. (right). The observations are: (1) RMoE-NP's entropy is slightly larger than SMoE's but largely smaller than RMoE''s. indicating that the larger entropy in RMoE is not from under-training but from cross-layer information sharing. (2) While RMoE-NP-r0.5 larger than SMoE and smaller than RMoF, -RMoF-NP-r1.00 is the largest. From Tab. 5 and Fig. 2, the small and large one both under-per under-perform RMM.,, These further demonstrate that the recurrent network can achieve a moderate flat gating score distribution, leading to better trade-off between exploration and exploitation.
Entrapy Layer
SMid
HyperMuE
Randemetal
Entropy of Layer 5/)
RAME
XHoE
Cas ine
RUMGE no 1.0
RUSOE no
no
10
2.5
2.0
1.5
XMol
CosinaMoE
RandomitteE
KMoE np 0.5
0.0
HyperMoE
0.0
Gate Entrapy
1.5
2.0
IMMEE
Gate Eniray
1.0
2.0
SMIE
114st:
Figure 4: Gate score entropy distribution Enwiki8 test set for different router configurations. More similar results can be found App. A.4.4 Fig. 8 Fig.
We also look into the statistics of selected experts' scores. Here we calculate the (1) Inner Balance (IB): defined as the ratio Top-1 score/Top-2 score, large IB means the first ex- pert dominates all selected experts; and (2) Outer Balance (OB), defined as EkeTop-k Ik. indicating the selected scores' ratio in the score distribution, large OB means selected expert scores dominate the gate score distribution. Because such ratio could have some extreme values, we report the median number for all tokens Tab. 8. We can observe: (1) Ran- domMoE, with fixed router, shows the lowest IB and OB. (2) Low-entropy models in the previous section (Sec. 6) have high IB and OB. (3) RMoE gives suitable IB and OB. While simply using a complex router ("RMoE-NP') shows relatively low IB and OB, RMoE is lower. Moreover, passing logits reduce IB and OB (<RMoE-NP=r-0.511.0)
Table
Expert scores balance on Enwiki8. Inner Balance (IB) repre- sents the (top-1 score top-2 score) tio, and Outer Balance (OB) represents
summed selected gate scores.
Algorithm
1B
OB
SMoE
34.68
0.915
HyperMoE
34.60
0.920
CosineMoE
7.611
0.794
XMoE
19.93
0.861
RandomMoE
2.000
0.414
RMoE
2.021
0.573
+ NP
16.58
0.842
+ NP r-0.5
2.792
0.661
All these experi- ments show sharing cross-layer router information can lead to more balanced routing decision and thus facilitate expert usage.
+ NP
1.147
0.212
Layerwise recurrence reduces the negative effeet of load balance constraint. To provide more direct analysis of the router gradient, we investigate how the gradient norm the router varies throughout the entire training process. When training MoE model, the gradient of the router has

---

Published as conference paper at ICLR 2025
Comparison Expert Similarities
Similarity
0.7
SMoE-medlle dlal
RMOE median
SMoE-mean
Nommed Expert Similantt Distributton
RMoL--menn
0.6
0.5
0.4
1k
3k 4k
5k
7k
8k 10k
Figure 5: Experts similarity distribution across layers during large-scale pre-training. We plot box plots of expert similarity from checkpoints taken every 1K training steps (approximately 4B tokens), showing the expert similarity across the 24 layers of the model (with maximum, minimum, first quartile, median, and mean).
two separate sources: (1) the language modeling (LM) loss, and (2) the load balancing (LB) loss that pushes the router assign tokens to different experts balanced manner. We empirically find (1) LB loss dominates the training of the linear router at the early training stage. This could hurt model's general performance, as Wang al. (2024) find, a high LB loss can cause balance token distribution but reduce performance. (2) On the contrary. the gradient of the RNN router from LB loss stabilises in the early stage, and the gradient from the LM loss keeps decreasing, suggesting that the RNN router is more optimised towards the LM loss. These observations suggest the recurrent router can effectively controls the influence of the LB loss. More details can be found in App. A.4.1
Layerwise recurrence encourages expert diversity One intriguing feature of MoE is that experts could modularly specialize on different inputs. Therefore, following recent works that analyze the FFNs (Geva et al., 2021; Qiu al., 2024a;b) and expert weights similarity (Wu et al., 2022; Lo et al., 2024), we use the cosine-similarity of expert's parameters to measure the expert diversity We calculate inFig. 5. To better understand the scale of similarity score, also plot one dash line showing the similarity of random initialized experts. More details about similarity calculation and explanation be found in App. A.4.3. We can observe that: (1) At the beginning of the training, the lowest expert similarities are similar to the random initialized one. (2) The expert similarity increases in the early training stages, then decreases later. This may be due the randomly initialized router in the early stages, which essentially assigns tokens randomly to different experts, leading to increased expert similarity. As the router continues learn, it gradually assigns specific tokens to the cor- responding experts, resulting in decreased expert similarity as training progresses.
SMoE RMoE in the large-scale pre-training settings, and results are shown
(3) During the entire training stages, the average similarity score between experts in RMoE lower than those in SMoF, indicating that RMoE encourages more diverse experts. This expert diversity also reasonably corresponds to the moderate flat gate scores in Sec
CONCLUSION
This work introduces a layer-wise recurrent router for existing MoE-hased language models. We validate the effectiveness of this layer-wise recurrence across various settings, tasks, and model sizes. By adding a new yet efficient computation stage the routing, RMoE stands orthogonal to most existing methods and can be flexibly integrated with them. Ablation studies reveal that this recurrent mechanism offers additional Recurrent Gradients, aiding router optimization. Further analysis validates our intuition that GRU facilitates inter-layer information sharing. We also system- atically compare RMoE's s model behavior with various baseline models, demonstrating that RMoE enhance existing SMoF methods and providing insights for future research.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
Jie Fu is supported by Shanghai Artificial Intelligence Laboratory.
REFERENCES
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding. Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek llm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.0295. 2024.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fowre AAAI Conference on Artificial Intelligence, 2020.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
Tianlong Chen, Zhenyu Zhang, Ajay Jaiswal, Shiwei Liu, and Zhangyang Wang. Sparse moe as the new dropout: Scaling dense and self-slimmable transformers. arXiv preprint arXiv:2303.01610, 2023.
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao. Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts. In NeurlPS, 2022.
Peter Clark, Isaac Cowhey, Oren Etzioni, Daniel Khashabi, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 2018.
reasoning challenge. In arXiv preprint
Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. Stable- moe: Stable routing strategy for mixture of experts. arXiv preprint axXiv:2244.883% 2022.
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.066 2024.
Rahul Dey and Fathi M Salem. Gate-variants of gated recurrent unit (gru) neural networks. In 2017 IEEE 60th international midwest symposium on circuits and systems (MWSCAS), pp. 1597-1600. IEEE, 2017.
Yifeng Ding, Jiawei Liu, Yuxiang Wei, Terry Yue Zhuo, and Lingming Zhang. Xft: Unlocking the power of code instruction tuning by simply merging upeyeled mixture-of-experts arXiv preprint arXiv:2404.1 15247, 2024.
Giang Do, Khiem Le, Quang Pham, Trungtin Nguyen, Thanh-Nam Doan, Bint T Nguyen, Chenghao Liu, Savitha Ramasamy, Xiaoli Li, and Steven Hoi. Hyperrouter: Towards efficient training and inference of sparse mixture of experts. arXiv preprint arXiv:2312.07035, 2023.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. A survey on in-context learning. arXiv preprint arXiv:2301 100344, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23:120:1-120:39 2022.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. MegaBlocks: Efficient Sparse Training with Mfxture-o- Experss.. Proceedings of Machine Learning and Systems, 5, 2023.
11

---

Published as a conference paper at ICLR 2025
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers key-value memories. In tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pp. 5484-5492 5495. Association for Computational Linguistics, 2021.
Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-
Zhuocheng Gong, Ang Lv, Jian Guan, Junxi Yan, Wei Wu, Huishuai Zhang, Minlie Huang, Dongyan Zhao, and Rui Yan. Mixture-of Reinventing transformers as dynamic assemblies of modules. arXiv preprint MxXiv:2407.066777 2024.
Daya Guo, Qihao Zhu, Dejian Yang, Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Y Wu, YK Li, et al. Deepseck-coder, When the large language model meets programming-the rise of code intelligence. arXiv preprint 1rXiv:2401.141990 2024.
Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A Smith, and Luke Zettlemoyer. Demix layers: Disentangling domains for modular language modeling. arXiv preprint arX1v:2108.05030, 2021.
Suchin Gururangan, Margaret Li, Mike Lewis, Weijia Shi, Tim Althoff, Noah A Smith, and Luke Zeullemoyer Scaling expert language models with unsupervised domain discovery. arXiv preprint arXiv:2303.14177 2023.
David Ha, Andrew Dai, and Quoe V Le. Hypernetworks. arXiv preprint wiXv::660.00910 2016.
Ruining He, Anirudh Ravula, Bhargav Kanagal, and Joshua Ainslie. Realformer: Transformer likes residual attention. arXiv preprint arXiv:2012.11747 2020.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 2020.
Quzhe Huang, Zhenwei An, Nan Zhuang, Mingxu Tao, Chen Zhang, Yang Jin, Kun Xu, Liwei Chen, Songfang Huang, and Yansong Feng. Harder tasks need more experts: Dynamic routing in moe models. arXiv preprint arXiv:2403.0760 07652, 2024.
Shima Imani, Liang Du, and Harsh Shrivastava. Mathprompter: Mathematical reasoning using large language models. arXiv preprint arXiv:2303.05396 2023.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upeycling: Training mixture-of-experts from dense checkpoints. In The Eleventh International Conference on Learn- ing Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
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. In Proceedings of the 2021 International Conference on Learning Representations (ICLR), 2021.
Bo Li, Yifei Shen, Jingkang Yang. Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenRe- view.net, 2023.
Margaret Li. Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A Smith, and Luke Zeullemoee. Branch-train-merge: Embarrassingly parallel training of expert language models. arXiv preprint 2333022288.03366 2022.
Bin Lin, Zhenyu Tang, Yang Ye, Jiaxi Cui, Bin Zhu, Peng Jin, Junwu Zhang, Munan Ning, and Li Yuan. Moe-llava: Mixture of experts for large vision-language models. arXiv preprint arXiv:2401.154477 2024.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv: 1800.0905. 2018.
12

---

Published as a conference paper at ICLR 2025
Ka Man Lo. Zeyu Huang, Zihan Qiu, Zili Wang, and Jie Fu. A closer look into mixture-of-experts in large language models. arXiv preprint arXiv:2406. 18219, 2024.
Larry R Medsker, Lakhmi Jain, et al. Recurrent neural networks. Design and Applicationn, 2, 2001.
Xiaonan Nie, Xupeng Miao, Shijie Cao, Lingxiao Ma, Qibin Liu, Jilong Xue, Youshan Miao, Yi Liu, Zhi Yang, and Bin Cui. Evomoe: An evolutional mixture-ofeexperts training framework via dense-to-sparse gate. arXiv preprint arXiv:21 14397, 2021.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models fol- low instructions with human feedback. Advances in neural information processing systems, 35: 27730 27744, 2022.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernandez. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 2016.
Hieu Pham, Melody Guan, BarretZoph, Quoe Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In International conference on machine learning, pp. 40.5 4104. PMLR, 2018.
Quang Pham, Giang Do, Huy Nguyen, TrungTin Nguyen, Chenghao Liu, Mina Sartipi, Binh Nguyen, Savitha Ramasamy, Xiaoli Li, Steven Hoi, al. Competesmoo-effectiee training of sparse mixture of experts via competition. arXiv preprint 2rXiv:2402.02526 2024.
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. arXiv preprint arXiv:2308.000 00951, 2023.
Zihan Qiu, Zeyu Huang, and Jie Fu. Unlocking emergent modularity in large language models. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies (Volume Long Papers), Mexico City, Mexico, June 2024a. Association for Computational Linguistics.
Zihan Qiu, Zeyu Huang, Youcheng Huang, and Jie Fu. Empirical study updating key-value memories in transformer feed-forward layers. arXiv preprint arXiv:2402.12233: 2024b.
Samyam Rajbhandari, Conglong Li, Z. Yao, Minjia Zhang, Reza Yazdani Aminabadi, A. Awan, Jeff Rasley, and Yuxiong He. Deepspeed-moe: Advancing mixture-of-experr inference and training to power next-generation ai scale. ArXiv, abs/2201.05596, 2022.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941 2017.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566 2021.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Sth International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017. Conference Track Proceedings. OpenReview.net, 2017.
Yikang Shen, Zheyu Zhang, Tianyou Cao, Shawn Tan, Zhenfang Chen, and Chuang Gan. Module- former: Learning modular large language models from uncurated data. CoRR, abs/2306.04640, 2023.
Samuel L. Smith, Ananya Kumar Ram, James Bradbury, Sharan Narang, Jared Casper, Matthew Johnson, Anselm Levskaya, John Schulman, Jascha Sohl-Dickstein and Barret Zoph. Using megablocks to scale language model training. In International Conference on Machine Learning, pp. 20275-20291. PMLR, 2022.
13

---

Published as a conference paper at ICLR 2025
Sainbayar Sukhbaatar, Olga Golovneva, Vasu Sharma, Hu Xu, Xi Victoria Lin, Baptiste Rozière, Jacob Kahn, Daniel Li, Wen-tau Yih, Jason Weston, et al. Branch- train-mix: Mixing expert Ilms 2024.
into mixture-of-eeperts Ilm. arXiv preprint arXiv:2403.0720
Rohan Taori, Ishaan Guirajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-followee llama model. https: //tp://gghtttt/tt com/ tat tatsu-lab/ sttanford_aaeeeeeee alpaca 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023. URL https: : /ww.iiivvvoo../// .ww.voorr//oooo /aa///b////// 2////3332221 13971.
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxiliary-loss-fre load balancing strategy for mixture-of-experts arXiv preprint arXiv:2408.15600 2024.
Johannes Welbl, Pontus Stenetorp, and Sebastian Riedel. Constructing datasets for multi-hop read- ing comprehension across documents. In arXiv preprint arXiv: 2017.
Haoze Wu, Zihan Qiu, Zili Wang, Hang Zhao, and Jie Fu. Gw-moe: Resolving uncertainty in moe router with global workspace theory. arXiv preprint 12.iv 2024.
Lemeng Wu, Mengchen Liu, Yinpeng Chen, Dongdong Chen, Xiyang Dai, and Lu Yuan. Residual mixture of experts. arXiv preprint arXiv:2204.09636 2022.
Fuzhao Xue, Ziji Shi, Futao Wei, Yuxuan Lou, Yong Liu, and Yang You. Go wider instead of deeper. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 8779-
2022.
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni. Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe: An early effort on open mixture-of-experts language models. arXiv preprint arXiv:2402.01739, 2024.
Yuanhang Yang, Shiyi Qi, Wenchao Gu, Chaozheng Wang, Cuiyun Gao, and Zenglin Xu. Enhancing efficiency in sparse models with sparser selection. arXiv preprint arXiv:2403.18926 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. In Proceedings of the 2022 Confer- ence on Empirical Methods Natural Language Processing. EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 4160-4162. Association for Computational Linguis- tics, 2022.
Hao Zhao, Zihan Qiu, Huijia Wu, Zili Wang, Zhaofeng He, and Jie Fu. Hypermoe: Towards better mixture of experts via transferring among experts. arXiv preprint arXiv:2402.12656. 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:7103-7114, 2022.
Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning, arXiv preprint 1eXiv:/61..111 2016.
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 27Xi::2002.08006 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.0422 2021.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 MORE RELATED WORKS
Routing Strategies While most MoE works follow the original success and use token choice routing, some works explore different routing approaches. In Expert-Choice Routing (Zhou al., 2022), each expert selects tokens to process across the whole batch input. This method avoids expert imbalance issues and allows different tokens be processed by flexible number of experts. Soft ofxxure--f-Expertt (Puigcerver et al., 2023) further assigns token weights for input tokens, weighted-averages them, and passes these merged tokens to different experts. This method moves one step behind the Expert-Choice Routing to allow more precise control. However, their token- selecting operations are non-causal and thus can't 't be directly used in the decoder models. Recent works (Huang et al., 2024; Yang et al., 2024) introduce dynamic top-k for each input token. While the FLOPs can be reduced, since this dynamic assignment can hurt the parallel computation of experts, more system-leve implementation must be optimized to achieve wall-time time efficiency. Some works also analyze issues in the routing of standard MoE like uncertain tokens (Wu et al., 2024) and lack of expert knowledge transfer (Zhao et al., 2024).
Training Strategies Due the unstable nature of MoE (Zoph et al., 2022), some works investi- gate special training strategies for MoE. EvoMoE (Nie et al., 2021) uses large top-k (even equal to the expert number) at the beginning of training, gradually decreasing k. StableMoB (Dai al., 2022) proposes to freeze the router after training some tokens to avoid token assignment conflicts. Residual Mixture of Experts (Wu al., 2022) initializes MoE from dense training checkpoints and finds it is an efficient method to train MoE models. Later, sparse-upcycling (Komatsuzaki et al., 2023) further trains large-scale language models from dense checkpoints, and many works follow this paradigm to efficiently utilize the power of MoE in fine tuning (Li et al., 2023), instruction tuning (Lin et al., 2024), and visual instruction tuning (Ding et al., 2024). Different from directly training MoE models, some works continue training the same pre-trained model on several differ- ent datasets to encourage specialization and combine them, either merging them into an MoE-style model (Gururangan et al., 2021; Sukhbaatar et al., 2024) or keeping group of models and intro- ducing a model-level router (Li et al., 2022; Gururangan et al., 2023).
Recurrence Controller A series of works introduce recurrent networks for Neural Architecture Search (NAS) (Zoph & Le, 2016; Ramachandran al., 2017; Pham al., 2018; Liu et al., 2018). They introduce a recurrent controller network that predicts the current layer-i's architecture (like CNN filters' number, size, and stride) based on layer-i's input hidden states and previous recurrent states (Zoph & Le, 2016). While these works use RNN to predict model architecture configurations of each layer for all inputs, RMoE utilizes RNN to help the router select expert combinations for each token, which can be viewed as a dynamic version of NAS.
A.2 EXPERIMENT SETUP
Enwiki8 and WikiText-103 We follow the default configurations in CompeteSMoE (Pham al., 2024). Each model is trained for 80,000 steps with Adam optimizer. The learning rate is 0.0007 with 4000 warmup steps, and the batch size is 48. The main used model is decoder-only transformer- based architecture with 8 layers and a hidden size of 352. It includes 16 experts, where the top 2 are selected during computation, each with an expert size of 352. The model uses 8 attention heads and handles sequences up to 512 tokens in length, with an attention span of 2048 tokens. It incorporates a dropout rate of 0.1 and load balancing factor of 0.01 to ensure an even distribution of expert utilization. Computation Cost Each 8-layer model is trained on one NVIDIA-A100 GPU for approximately 21 hours.
Large Scale Pre-training For model architecture, our 24-layer model employs Rotary Embedding for positional encoding, SwiGLU for activation functions, and RMSNorm to enhance the model's efficiency and performance. Other model configuration includes a hidden size of 1280, 20 atten- tion heads, initialization method standard deviation of 0.02, sequence length of 4096, and maximum positional embedding length of 4096. All dropout rates are set to 0. For the MoE part, 16 experts, with each expert having feedforward network hidden size of 448, following
15

---

Published as a conference paper at ICLR 2025
the fine- fine-grained MoE settings, and each token activating 4 experts. We tokenizer with 96512 vocabulary size, which adds approximately 123M embedding parameters and 123M vocab- ulary projection head parameters. Under this configuration, each model has approximately 664M non-embedding parameters, and every token activates 334M non-embedding parameters. The total parameter is around 910M. For pre-training configurations, we use global batch size of 1 1120, a warmup period of 2000 iterations, learning rate of 4.2c-4, a minimum learning rate of 4.2e-5, cosine learning rate decay, Adam optimizer with 81 0.9 and 32 0.95, a weight decay of 0.1, and gradient clipping at 1.0. Computation Cost Each 24-layer model is trained on 8 NVIDIA-A100 GPUs for approximately 5 days.
Instruction Tuning Data The Alpaca (Taori et al., 2023) dataset is an open-source instruction- following dataset created by Stanford researchers, inspired by OpenAl's CharGPT. The dataset con- sists of 52,000 instruction-response pairs generated using the text-davinci-003 model by providing diverse and comprehensive instructions and recording the corresponding responses. It is designed to facilitate the training and evaluation of models in understanding and generating human-lik text responses to various instructions.
Instruction Tuning Setting We use the codebase² and corresponding default configurations. 6 (bf16) precision to accelerate training while maintaining numer- ical stability. The model is trained for 3 epochs using Adam optimizer with global batch size 128. We set the learning rate to 2e-5 and do not apply weight decay. A warmup ratio of 0.03 is used to gradually increase the learning rate at the beginning of training, and we utilize a cosine learn- ing rate scheduler adjust throughout the training process, promoting smoother convergence. Computation Cost Each is trained on 8 NVIDIA-A 100 GPUs for approximately 2 hours.
More concretely,
Evaluation Tasks
Here we shortly describe our used evaluation datasets:
ARC-Easy is a subset of the A12 Reasoning Challenge (ARC) dataset (Clark et al., 2018). It consists of multiple-choice questions from elementary and middle school science exams that relatively easier than the ARC-Challenge set. These questions require basic reasoning and knowledge appli- cation.
Hellaswag (Zellers et al., 2019) is dataset designed for commonsense reasoning and narrative prediction. It involves choosing the most plausible continuation of a given scenario from multiple options. The task is challenging because it requires understanding and applying common sense knowledge.
PIQA (Bisk et al., 2020) dataset tests model's ability to understand and reason about physical interactions and affordances. The task involves selecting the correct answer to questions about everyday physical activities.
SciQ (Welbl et al., 2017) is a dataset of science questions that includes multiple-choice and direct- answer formats. It aims to test model's ability to understand and reason with scientific concepts typically taught the school level.
LAMBADA (Paperno et al., 2016) is dataset designed for language modeling and comprehension. The task involves predicting the last word of a given passage, which requires deep understanding of the context provided by the preceding text.
A.3 FURTHER PRETRAINING VALIDATION
To further validate the scalability of RMOF, we conduct experiments with larger model sizes and increased pre-training corpus. Both MoE models followed the design principles of DeepSeek- MoF. (Dai et al., 2024), utilizing fine-grad grained experts and shared experts to maintain strong baselines. We evaluated the models on more challenging benchmarks, including Hellaswag, MMLU, GSM8K, and HumanEval, to assess their language capabilities, multi-domain knowledge, mathematical skills, and coding abilities. Additionally, we tested the models' perplexity on multiple domain test datasets and reported the average results.
hhan:://githuuccco//sssaaaaaaana alpaca
16

---

Published as a conference paper at ICLR 2025
Tab. 9 and Tab. 10 present the performance of 15-billion parameter model with 2.7 billion activated experts, trained on 120 billion and 400 billion tokens, respectively. The results show that RMoE consistently delivers improvements
with increased data volumes. The findings indicate that RMoE enhances performance in standard language modeling tasks, such as Hellaswag and PPL., and improves performance more complex reasoning tasks.
Table Performance comparison SMoE, SMoE-MLP and RMoE the model scale of 15B activation 2.7B parameters, training 120B tokens.
Hellaswag
MMLU
GSM8K
Avg PPL
Pretrain 80B Tokens
SMoE 67.69
46.24
24.18
7.406
SMoE-MLP
67.98
46.47
23.58
7.437
RMoE
68.00
47.74
27.14
7.361
Pretrain 100B Tokens
SMoE SMoE-MLP
70.98
50.61
30.78
6.754
70.8
50.6
30.17
6.786
RMoE
71.02
51.74
32.98
6.732
Pretrain 120B Tokens
SMoE
72.03
52.79
34.8
6.447
SMoE-MLP
72.19
52.81
34.57
6.479
RMoE
72.36
54.02
36.13
6.425
Table 10: Performance comparison SMoE, SMoE-MLP and RMoE model scale
15B activation
2.7B parameters, training 400B tokens.
Hellaswag
Avg PPL Pretrain 200B Tokens
MMLU
GSM8K
SMoE
69.48
49.96
33.21
7.718
SMoE-MLP
69.76
50.27
31.77
7.736
RMoE
70.00
52.21
32.98
7.608
Pretrain 280B Tokens
SMoE SMoE-MLP
72.40
54.66
42.61
6.477
72.62
55.33
38.51
6.502
RMoE
73.18
56.06
44.35
6,400
Pretrain 400B Tokens
SMoE
76.39
59.54
52.16
5.685
SMoE-MLP
76.09
59.96
51.71
5.709
RMoE
76.72
60.60
52.99
5.620
A.4 ADDITIONAL OBSERVATIONS
17

---

Published as a conference paper at ICLR 2025
A.4.1 ROUTER GRADIENT NORM AND DROP RATIO
Table 11: Comparison of linear and RNN routers record the router gradient every 10k training steps (20B tokens). We compute the gradient with language modeling (LM) loss and load balance (LB) loss. Drop ratio is the ratio of dropped tokens and all tokens as we assign capacity factor 1.0 for each expert.
of gradients and drop ratios various training steps.
Training steps (k step)
0.1
10
20
30
40
50
60
Linear router
grad from the whole loss
1.058
0.194
0.1911
0.198
0.208
0.217
0.221
grad from LM loss
0.625
0.183
0.184
0.192
0.204
0.215
0.220
grad from LB loss
0.433
0.011
0.008
0.006
0.004
0.002
0.001
drop ratio
35.6
5.43
5.34
5.17
4.89
4.64
4.50
RNN router
grad from the whole loss
0.972
0.160
0.153
0.153
0.155
0.155
0.154
grad from LM loss
0.636
0.146
0.138
0.139
0.144
0.148
0.151
grad from LB loss drop ratio
0.337
0.014
0.015
0.014
0.011
0.007
0.003
38.7
6.35
6.30
5.94
5.32
4.54
4.09
Based on the setting of training 15B models for 120B tokens, we investigate how the gradient norm of the router varies throughout the entire training process. When training an MoE-based model, the gradient of the router has two separate sources: due to (1) the language modeling (LM) loss, and (2) the load balancing (LB) loss that forces the router to assign tokens different experts balanced manner. Therefore, for each router, we compare the gradient from the LM loss only and from the whole training loss. We calculate the average for 100 training steps to estimate the gradient norm.
Furthermore, to better investigate the relation between the router behavior and the router gradient, we calculate the drop ratio for the router. This because during the large-scale MoE pre-training, to ensure the training efficiency, the expert is usually controlled by an hyper-parameter called capacity factor, which determines the total tokens that one expert can process. If the router assigns tokens to some expert that exceeds its capacity, the expert will drop tokens with the lowest scores. And define the drop ratio as tokens dropped / total tokens. The LB loss mentioned before is critical to decreasing the drop ratio.
According to Tab. 11, we have the following observations: 1. The gradient norm of the RNN router is generally smaller than that of the linear router. And for both routers, the drop ratio decreases with the training. 2. According to the drop ratios, we observe the significant behavioral difference between the two routers: during the early training phase (10k steps - 30k steps), the drop ratio of the linear router is noticeably lower than that of the RNN router; the drop ratio of the RNN router archives at the lower value in the end. The trend observed in the drop ratio is consistent with the results of the gradient norm. The grad norm for LB loss is relatively higher in the RNN router until the final training stage (50k 60k), whereas the gradient from LB loss in the linear router high at the beginning and generally low during the later part of training (10k 60k).
These phenomena indicate that the LB loss could dominate the training of the linear router: when the drop ratio is low and stays unchanged, the grad from LB loss will be low because the router is already well-optimized for LB loss. Such early convergence in the LB loss may reach suboptimal solution in the trade-off between optimizing load balance and language modeling, On the contrary, the gradient of the RNN router from LB loss stabilizes in the early training steps (10k 30k), and the gradient from the Im loss keeps decreasing, suggesting that the RNN router is more optimized towards the LM loss.
18

---

Published as a conference paper at ICLR 2025
A.4.2 MUTUAL INFORMATION
import numpy np from sklearn metrics import mutual_info.ss score
def discretize prr. dist dist, bins=100):
Discretize the probability distribution into discrete bins.
discretized np. digitize (prob-dist, bins=np linspace (0, 1, bins)) return discretized
def caaatte_mttall___iiteee
(xl, x2, bins=1000
Calculate mutual information between each pair of distributions in x and x2. xl, x2: numpy arrays of shape (N, 16)
bins: number of bins to use for discretization Returns numpy array of mutual information values.
valuel []
for i in range (xl. (xl. [0]]][[]]]
disciccret = discretize-pp prob distt bins)
x2-discretized
beccretize-probdiig (x2 bins)
mi = mutual info score discretized, x2 discretized) mi values append
return np. array mi values)
A.4.3 EXPERT SIMILARITIES
def
get.similarities (htoh4_0, htoh4_1, h4toh
avg_key-0 htoh4_0 mreandim=1l # (num_experts.. 4h, h)
avg_key_1 htoh4_1 meam (dim=l) # (num_experts: 4h, h)
vaguvalue h4toh mean (dim=2) # (num_experts, h. 4h)
normed key 0 nn. functional nor.aall.oommaieeeaaaa (avg key 0, p=2, dim=l)
normed key
functional normalize ( (avg key 1, p=2, dim=1)
normed_value = mn. functional aoraalloooooeeeeeeeeeeao (avg.value p=2, dim=1)
torch. cat ([normed.key._l normed.key.1, normed_value] dim=1)
# compute the average expert similarity
similarity torch .mm( www.ee-avgeeeeeeeeeeeeeeeee avg-sim normed-similan return avg sim
avg
mean () i)
A.4.4 MORE ROUTER ENTROPY DISTRIBUTIONS
19

---

Published as a conference paper at ICLR 2025
Mean MI of RMoE-NP-rO.5
Mean CosineMoE
12
7
7-
6-
1.044
1.0
6-
0,967
11
5
5
Layer
4
6180
6/27
4
0.119
6.04
0/1
-1A
1.801
1.54
6.55
0.995
0.474
3.000
0.000
0.90
2-
1
61841
6.00
12
1
0.324
3.048
82
0
D/4
01.64
DA12
00:04
K.44
0
1
2
4
5
7
1
2
4
5
Layer
Layer
Figure 6: Mutual information of RMoE-NP-10.3 and CosineMoE settings
Figure 7: Mutual information of SMoE, RMoE, RMOE-NP, and RMoE-NP-10.5 in 24-layer models.
20

---

Published as a conference paper at ICLR 2025
Entropy of Layer
RMoE
Entropy Layer 1/
RMoE
HyperMet
HyperNot
RandonetsE
RandemMMtE
CostneMoE
CostreMoE
5.0
2.5
1.5
1.0
0.5
0.5
0.5
HyperMoE
0.5
HyporMoE
Cata Estmay
RMOE
Gate Entwwyy
RMOE
2.5
SMGE
2.9
2.5
SMdE
Entropy of Layer (7)
Entropy Layer 4 (7)
MMuE
HyperMuE
Hypas MoE
RandomettE
RandomMtt
CosirwMot
CusireMot
1.0
3.5
3.0
3.5
g
2.5
1.5
1.0
0.5
CosinsMaE
CosinaMaE
RandomMta
RandoneMtE
0.5
HyperMoE
0.
HyperMoE
Gete Entiopy
INNEE
Gate Eiitwoy
1.5
2.0
2.5
SMdE
2.0
2.5
SMdE
Entropy of Layer 6/ (7)
IMuE
Entropy Layer (7)
HyperHoE
HyperMot
Ransemtiof
Bandombtot
CosineMoE
CusiruMot
XHoE
2.5
2.0
2.0
155
1.5
1.0
CosinoMaE
CosInaMaE
RandomitiaE
RanconetiE
6.9
0.5
0.5
Cato Entreyy
Gace Entrogy
1.5
2.0
2.5
SMoE
2.0
SMaE
Figure 8: Gate score entropy distribution over Enwiki test different routers 8-layer models.
A.4.5
ROUTER WEIGHTS INFORMATION
A.4.6
EXPERT SELECTION FREQUENCY
A.5
ADDITIONAL RESULTS
21

---

Published as a conference paper at ICLR 2025
Table 12: More SMoE and RMoE variants pre-training costs and evaluation results in selected informative Im-evaluation-haal harness tasks. means supervised fine-tuning on the Alpaca dataset. The names and metrics for short names the table are: 'ARC-e' for ARC-Fasy, acc; 'Hella' for Hellaswag, acc-norm; "Piqa" for PIQA, acc-norm; 'Lamb' for LAMBADA, acc.
Algorithm
Training
ARC-e
Hella
Piqa
Sciq
Lamb
Avg
(5k steps)
47.14
35.51
64.69
76.2
14.61
47.63
+sft
50.93
35.82
65.61
74.7
17.81
48.97
+sft (freeze gate)
50.59
35.78
66.32
74.7
18.18
49.11
SMoE
40B (10k steps)
52.57
40.85
67.74
83.4
26.74
54.26
+sft
53.7
42.07
68.61
83.5
32.8
56.13
+sft (freeze gale)
53.45
41.94
68.88
83.1
32.06
55.89
20B
47.01
35.91
65.23
78.7
19.13
49.20
+sft
48.53
36.9
66.21
79.6
24.74
51.20
RMoE
+sft (freeze router)
48.65
36.88
66.43
80.1
24.55
51.32
+sft (freeze router and GRU)
49.24
36.79
66.16
79.7
24.32
51.24
GRU
40B
51.18
41.38
67.79
83.6
32.58
55.31
p 128
+sft
53.20
43.05
68.55
83.8
37.16
57.15
+sft (freeze router)
53.03
42.96
68.34
83.6
36.68
56.92
+sft (freeze router and GRU)
53.11
43.16
68.77
82.8
37.57
57.08
20B
47.47
35.91
65.78
76.2
20.03
49.08
+sft
48.36
36.49
65.07
77.4
22.86
50.04
RMoE
+sft (Treeze router)
48.27
36.42
65.23
76.9
22.88
49.94
+sft (freeze router GRU)
48.23
36.46
64.94
77.3
22.61
49.91
GRU
40B
53.07
41.15
68.52
84.0
19.17
53.18
256
+sft
54.46
43.06
67.46
84.9
24.57
54.89
+sft (freeze router)
54.45
43.10
67.19
84.1
23.93
54.55
+sft (freeze router and GRU)
54.50
43.13
67.36
83.8
23.62
54.48
20B
47.77
35.39
64.80
79.5
25.00
50.49
+sft
48.27
36.47
65.51
76.6
22.18
49.81
RMoE
+sft (freeze router)
47.73
36.41
65.78
76.6
22.88
49.88
+sft (freeze router and GRU)
48.19
36.22
65.29
76.8
23.5
50.00
GRU
40B
51.64
41.37
66.81
86.0
22.76
53.72
p = 512
+sft
52.82
42.68
68.55
86.0
26.88
55.39
+sft (freeze router)
52.48
42.61
68.44
86.0
27.23
55,35
+sft (freeze router and GRU)
52.74
42.44
68.77
86.3
27.13
55.48
RMoE
20B
46.63
35.7
64.91
76.1
16.24
47.92
+sft
48.40
36.45
65.51
77.3
22.65
50.06
RNN
+sft (freeze router)
48.70
36.29
65.45
77.3
22.60
50.07
= 256
+sft (freeze router and RNN)
49.24
36.48
65.56
77.7
23.20
50.44
22

---

Published as a conference paper at ICLR 2025
SHat
RIKE
Layer
RMvE-nn
RAMaE-no
RWat--p--1.0
RMeE--p-10
RNcE-
TIMot-rg
0.5
0.5
1.5
1.5
Entmpy
Entmpy
2.5
Entropy
Entropy Layer
10
Denste
RVAe 1.0
RMAEE np
no
np
RNCE
RMLE
0.5
0.5
1.5 Gate Eranpy
RMIEE
1.5 Gate Enrroyy
RMAE
SWAE
Entropy of Layer 5/
Entropy of Layer
nD
BiMeE-nn-11
mo-1
RMLE
Gate Ertmay
2.0
SWAE
Gate Sxtmpy
2.9
SMIE
Figure 9: Gate score entropy distribution over Enwiki test set for different information passing settings in 8-layer models.
23

---

Published as a conference paper at ICLR 2025
AM05
35
2.5
MoCA
Extraey
Extragy
5966
Enersyy of Layer
Eneppy Layer
2016
15
35
4000
3.0 Gate Entragy
4045
Gate Entraay
3/5
25
Entropy of Layer
Encrosy of Layor 5
5506
4.0
25
6.0
1.0
8.5
3.8
Gire Enttngy
23
Cite Enbcay
2.5
2.5
0.5
3.5
200E
S1RCE
Figure 10: Gate score entropy distribution over Enwiki test set for different routers. RMoE can combined with XMoE to encourage the exploration of XMoF.
24

---

Published as a conference paper at ICLR 2025
Care Weight Merm ayert
2A0E
320
3.50
2
the horm
SM6C
614
Gane We dhe horm (24, ayeret
Gute We ehr 124 byeen
214
538-
Weicht
Suie Weath 501 laynos
633
015
SPAA
434
Figure 11: Different layers' router weight statistics (left column: norm and right column: standard deviation)
Enwiki8 setting. (1) different layers have different norms and STDs, which inspires us to introduce layerwisc projector in Equ 4 and explains using the shared projector can hurt RMoE's performance (Tab. 6). (2) While
SMoE routers show larger weight norms than RMoE settings, their standard deviations are not the highest. The large router
potentially explain the larger IB and OB in Tab. 8.
25

---

Published as a conference paper at ICLR 2025
Expert Frequency of SMoE
Expert Frequency of RMoE
0.14
0.34
0.12
0.12
0.08
0.08
0.06
0.06
0.02
0.07
0.00
0.00
Layer
Layer
Expert Frequency of HyperMoE
Expert Frequency RandomMaE
0.14
0.12
0.12
0.08
0.08
0.00
0.06
0.04
0.00
0.00
Layer
Layer
Expert Frequency of XMOE
Expert Frequency of CosineMoE
0.32
0.12
mequueccy
0.30
0.20
0.08
0.06
0.06
0.04
0.04
0.00
0.00
Layer
Lay
Figure 12: Different methods' expert selection frequency on medium size models Enwiki8. (1) RMoE slightly increases expert imbalance than SMoE. (2) Methods using a frozen-random cea----randmiiiiiaaie initialize router (Hyper- MoB and RandomMoF) show more imbalance problems.
26

---

Published as a conference paper at ICLR 2025
SM
1.0
XMoE R
smiienty
1.0
0.
0.8
SHot
IMAÉ
0.6
HyperMoE
0.6
StableMuE
0.5
RandomMoE
layer
layer
Expert similarity in Enwiki8 training experiments. RandomMoB shows the highest expert sim- ilarity. XMoE, which introduces down-projected cosine routing resolve representation collapse in SMoE, shows the lowest expert similarity. While RMoE doesn't significantly diversify experts as the large-scale training settings (left), it can be further combined with XMoE, which largely increases expert diversity and brings improvement (right).
Figure 13:
Validation BPC vs Layer
1.22
SMoE
RMOE
1.20
RMOE-NP
Validatio Resuits (BPC)
RMoE- NP-1O-
2.18
2.16
1.14
1.12
110
10
(Layer)
Figure 14: Validation BPC on Enwiki8 with different model sizes (6, 12, 18, 24, 32 layers).
27


Reviews:
[{"review_id": "4zJQ9JbX30", "reviewer": "Reviewer_rPzv", "summary": "This paper proposes the Layerwise Recurrent Router for Mixture-of-Experts, a novel approach to enhance efficiency in MoE architectures within LLM. The authors highlight limitations in traditional MoE routers that operate independently across layers, potentially resulting in suboptimal parameter utilization. To address this, Recurrent-Router MoE introduces a GRU-based mechanism that enables dependencies between routing decisions across layers, which improves the diversity and effectiveness of expert utilization. The proposed method integrates seamlessly with other MoE architectures and demonstrates substantial improvements in performance on various language modeling tasks. Through comprehensive experiments, the authors show that RR-MoE achieves superior performance over several baseline router designs with only marginal computational overhead.", "questions": "The paper emphasizes the importance of gradient flow through the GRU-based recurrence in RMoE but does not mention whether Truncated Backpropagation Through Time (BPTT) or any specific method was used to manage gradients across layers. Did the authors consider using truncated BPTT for controlling gradient flow through the recurrent router? Including details on gradient management would help clarify if additional methods, like truncation, could further optimize training.\n\n Would it be possible to conduct a more direct analysis, such as by visualizing or measuring gradient norms through layers, to further validate the specific impact of recurrent gradients? \n\n Could the authors provide insights or quantitative analysis on how the recurrence depth impacts the RMoE’s routing decisions?\n\nThe current evaluation omits challenging benchmarks like MMLU and GSM8K, mentioning that the models are still under-trained. Can the authors comment on whether they plan to scale the training or incorporate task-specific fine-tuning to gauge RMoE’s performance on these harder tasks? Additional results on these benchmarks could help better gauge RMoE’s robustness and adaptability in real-world scenarios.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper offers a simple but new enhancement to MoE architectures by introducing a layerwise recurrent mechanism, distinguishing it from other MoE router designs that operate independently across layers. The methodology is sound, with a well-motivated design backed by extensive experimental validation and ablation studies. The work is valuable to the research community, as it not only provides a means to improve existing MoE designs but also contributes insights into the impact of cross-layer recurrence on expert diversity and selection.", "weaknesses": "The authors mention that their model is under-trained and does not yet perform well on challenging tasks like MMLU and GSM8K. A more complete evaluation on these benchmarks, perhaps by including more training tokens or longer training duration, would provide a fuller picture of the model's real-world applicability.\n\nAlso the paper presents extensive performance metrics and ablation studies, a more direct analysis of gradient propagation through the recurrent router could strengthen the argument for the importance of recurrent gradients in optimizing routing decisions. Also, do we need truncated backpropagation similar to BPTT?\n\nThe paper could benefit from a deeper comparative analysis with recent MoE models that also use recurrence or other advanced routing strategies (such as XMoE, HyperMoE etc). This would help clarify how RRMoE’s benefits compare to similar innovations in this rapidly evolving field.", "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": "LyzpcrMDys", "reviewer": "Reviewer_bumc", "summary": "Paper introduces a gated recurrent connection for routing decisions in MoEs such that routing decision in a layer depend on shared GRU so that routing decisions between layers are no longer independent. The intuition expressed by authors is that this should help improve model because routing decisions in current layer should be able to leverage routing info from previous layers.", "questions": "Suggestions for improvement are to show results on a larger MoE model. Not all models need to be 100s of billions of parameters, but a 36M parameter MoE model is quite small these days, and it is hard to be convinced that small gains on this model would translate to a more realistic sized model.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Paper clearly expresses main idea, that routing decisions in MoE should have some dependency on previous routing decisions, and how they achieve this and test the results.", "weaknesses": "The results are relatively minor improvements (1-2% relative), the models are very small (on the order of 36M parameters for an MoE model) and the benchmarks are small (enwiki8 & wikitext).", "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": "qE2svBnrJM", "reviewer": "Reviewer_Cw7i", "summary": "The paper presents the Layerwise Recurrent Router for Mixture-of-Experts (RMoE), designed to address inefficiencies where routers at each layer operate independently, leading to suboptimal token-expert assignments. RMoE integrates a GRU to enable cross-layer communication, establishing dependencies between routing decisions across layers. The authors validate RMoE through experiments across various model sizes and tasks, demonstrating consistent improvements over baseline MoE models. Analysis shows that RMoE’s gains are primarily due to effective cross-layer information sharing, resulting in better expert utilization and diversity.", "questions": "- Clarify the architecture of the layer-specific projectors used in Eq. 4; are there structural variations across layers, or have alternative projector architectures been tested for potential performance improvements?\n- In Table 1, improvement over baselines is minor, yet gains are more noticeable in Table 2 (billion-level settings). Could the authors explain the differing results, specifically regarding training sizes and conditions?\n- Column names in Table 8 appear inconsistent with descriptions in the text; could these be corrected or clarified?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The use of a GRU to incorporate layerwise information into MoE is a novel and well-motivated approach to address parameter inefficiencies in SMoE.\n- Extensive ablation studies demonstrate the importance of recurrent gradient flow, validating the primary contributor to RMoE’s improvements.\n- Interesting analysis on the increase in cross-layer mutual information and gating score entropy with RMoE, providing insights into how cross-layer information sharing improves expert selection and model robustness.", "weaknesses": "- Current performance improvement appears marginal; val and test losses alone do not demonstrate statistically significant gains over baselines. Benchmarking on standard language modeling tasks would provide more practical context for the model's effectiveness.\n- Analysis of gradient flow could be more rigorous; while recurrent gradient flow is highlighted as crucial in Section 5, more sophisticated gradient tracking techniques would strengthen this claim.\n- The additional computation stage introduced by the GRU increases complexity, yet the impact on training efficiency is not thoroughly analyzed across different model sizes.", "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": "TbFci4cLm7", "reviewer": "Reviewer_oYLa", "summary": "This paper proposes to use RNN to replace the linear projection layer to route the tokens in MoE models. The proposed model achieves better results than vanilla router across scales (after revision).", "questions": "NA", "limitations": "", "rating": 8, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1) Elegant design. Using previous routing representation to improve the routing decision is a smart idea.\n2) Some experimental results in Section 6 (Observation) are interesting, such as Layer-wise recurrence encourages expert diversity.\n3) Comprehensive ablation studies in Section, which shows the effectiveness of RNN router under this scale.\n4) Unshare part of trainable parameters in RNN also makes sense. I suggest to cite this work (https://arxiv.org/abs/2107.11817), which unshare the layernorm when sharing MoE Transformer blocks.", "weaknesses": "I have only one important concern, i.e. the scale of the experiments. I agree the paper is promising in many aspects, but the experiments are too toy for a model architecture paper in 2024. We need at least 3B/7B activated parameters and 1T tokens to justify the scalability of the proposed method. \n\nHowever, I understand it might be difficult to acquire enough computation resources for some reasons, so I decided to give a positive rating to this paper.", "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)