Towards Being Parameter-Efficient: A Stratified Sparsely Activated Transformer with Dynamic Capacity
Authors: Haoran Xu
Maha Elbayad
Kenton Murray
Jean Maillard
Vedanuj Goswami
Venue: EMNLP 2023 Findings (2023)
Keywords: Stratified Mixture-of-Experts
Parameter-efficiency
Dynamic capacity
TL;DR: We propose Stratified Mixture of Experts (SMoE) models, characterized by a stratified structure that allows for dynamic allocation of capacity to input tokens, taking into consideration the varying capacity requirements of different tokens.

Abstract: Mixture-of-experts (MoE) models that employ sparse activation have demonstrated effectiveness in significantly increasing the number of parameters while maintaining low computational requirements per token.  However, recent studies have established that MoE models are inherently parameter-inefficient as the improvement in performance diminishes with an increasing number of experts. We hypothesize this parameter inefficiency is a result of all experts having equal capacity, which may not adequately meet the varying complexity requirements of different tokens or tasks. In light of this, we propose Stratified Mixture of Experts (SMoE) models, which feature a stratified structure and can assign dynamic capacity to different tokens. We demonstrate the effectiveness of SMoE on three multilingual machine translation benchmarks, containing 4, 15, and 94 language pairs, respectively. We show that SMoE outperforms multiple state-of-the-art MoE models with the same or fewer parameters.

'Code is released at https://github..com ixxu/ Miwieeddiittuttoo___eee Towards Being Parameter-Eficient: A Stratified Sparsely Activated Transformer with Dynamic Capacity
Haoran Xu Maha Elbayad Kenton Murray Jean Maillard Vedanuj Goswami
Johns Hopkins University, Meta Al
{hxu64, kenton @eeeeeeheeeeeeeee edu {elbayadm, jeanam,,mjjjjjmm....m. vedanuj @ tamett.co com
Abstract
The efficiency comes from sparsely activating a subset of the neural network weights for each incoming sample. However, MoE is reported to be parameter-inefficint (Hoffmann et al., 2022; Zuo et al.. 2021; Gao et al., 2022) i.e., there are diminishing improvement returns from adding more experts. For example, Switch Transformer (Fedus et al., 2021) only outperforms T5 (Raffel et al., 2020) by an average of 0.7 on the GLUE benchmark (Wang al., 2018) despite being 35x larger. Similarly, in the translation task, a MoE model with 20 times more parameters only offers an average improvement of 0.3 BLEU on its ablation dataset (MoE-64 vs. 1.3B dense) (NLLB Team et al., 2022).
Mixtere-of-experts (MoE) models that employ sparse activation have demonstrated effectiveness in significantly increasing the number of parameters while maintaining low computational requirements per token. However, recent studies (Hoffmann et al., 2022; 7uo et al., 2021; Gao et al., 2022) have established that MoE models are inherently marameter-ineffcient as the improvement in performance diminishes with an increasing number of experts. We hypothesize this parameter inefficiency is result of all experts having equal capacity, which may not adequately meet the varying complexity requirements of different tokens tasks. In light of this, we propose Stratified Mixture of Experts (SMoE) models, which feature stratified structure and can assign dynamic capacity to different tokens. We demonstrate the effectiveness of SMoE three multilingual machine translation benchmarks, containing 4, 15, and 94 language pairs, respectively. We show that SMoE outperforms multiple state-of-thc-art MoE models with the same or fewer parameters.'
We hypothesize that this parameter inefficiency stems from the equal capacity assignment, where particularly define 'capacity' as the number of parameters used for the incoming token. For the current MoE models, the capacity of experts are same used for serving all tokens. However, diffcrent tokens may demand varying capacities. For instance, in the context of multilingual machine translation, certain translation directions may necessitate greater capacity to prevent overfitting, while others only require smaller capacity. To address this limitation, our hypothesis posits that the dynamic allocation of capacity to tokens results in more efficient utilization of parameters. Thus, we propose Stratified Mixture of Experts (SMoE) models, characterized by a stratified structure, which allows for the dynamic assignment of capacity to incoming tokens.
1 Introduction
Scaling up the model and data size has shown tremendous success in enhancing model performance across a large number of NL.P tasks (Devlin et al., 2019; Conneau et al., 2020; Kaplan et al.. 2020; Brown et al., 2020). Sparsely gated mixture of experts (MoE) (Shazeer et al., 2017; Lepikhin et al., 2021) provides an effective way to greatly scale the model size under the same computational cost and achieves state-of- the-art performances on various tasks including natural language understanding (Fedus et al., 2021), machine translation (NLLB Team et al., 2022), language modeling (Du et al., 2022), etc.
A high-level comparison of vanilla MoE and SMoE. is presented in Figures la and 1b. In vanilla MoE, single routing gate connects to all E experts and sends tokens the top-k experts. Here, we take E-5 an example. In SMoE, the experts are divided into two strata. Each stratum has its own routing gate that connects to all experts in the current stratum as well as all experts

---

PFN
PFN2
PFN3
FFN
FFNs
FFN,
FFNg
Gate
Gate
stratum
(a) Vanilla MoF
LayerNorm
FFN
FFN
Gates
FFN
FFN2
FFN&
stratum
FFN
FFN2
FFNg
Gate
stratum
Gate,
LayerNorm
stratum
(b) SMoE
Architectural design of SMoE block
Figure I: A high-level illustration of the vanilla MoE and SMoE. a) Vanilla MoE: gate is connected to all experts and sends tokens to top-k selection. b) SMoE: Experts are stratified into strata (L=2 in this example). Each stratum has gate that is connected all subsequent experts. Tokens be directly sent the last stratum only experience one expert, or be sent to both strata and have more capacity. Hence, the dynamic capacity of a token depends on how many experts needs to pass. c) detailed architectural design, where a comprehensive explanation of the design components will be presented in Section 3.
the subsequent strata. If Gate; assigns tokens to Expert 4 or 5, the tokens will only need to pass through a single expert (an FFN layer). However, if tokens are sent to experts in the first stratum (Experts to 3), they will need to go through the next stratum well, meaning that another expert will be assigned by Gatez before exiting the SMoE block. This allows SMoE dynamically assign capacity different tokens. In addition, a comprehensive illustration of the architectural design of the SMoE model provided in Figure Ic. A thorough explanation of the design elements will be provided in Section 3. Our main contributions are summarized as follows:
parameters to achieve performance on- par with naive MoE (Lepikhin et al., 2021). Furthermore, we carry out an in- depth analysis to probe the factors that impact dynamic capacity assignment, including the language of tokens and the position of the SMoE block within the model's architecture.
Background and Related Work
Massively multilingual machine translation models have been developed to handle several translation directions simultancously in a single model (Aharoni et al., 2019). However, the use of shared parameters for different languages often leads to negative transfer and decreased performance (Conneau et al., 2020; Fan et al., 2020). In contrast to dense MMT models, sparsely gated mixture- of-experts (MoE) models, which activate a subset of parameters for each input, have been shown to significantly improve translation performance (Kim et al., 2021; NLLB Team et al., 2022). Shazeer et al. (2017) first demonstrated the benefit of adding MoE layers to scale RNN models for improved translation performance, and Lepikhin et al. (2021) extended this work transformer architectures (Vaswani et al., 2017). MoE. layers in the transformer model replace a single feedforward network (FFN) layer with FFN layers, denoted
We introduce the concept of dynamic capacity for MoE models and propose mixture-offexpprts model with a stratified structure, namely SMoE, which can automatically assign dynamic capacity to differentt incoming tokens to make experts become more parameter-efficient.
We focus the task of multilingual machine translation (MMT) and show that SMoE substantially outperforms numerous strong baselines with fewer than or the same number of parameters. For instance, we demonstrate that SMoE only needs half the number of

---

with Each FFN layer is an expert. Given an imput token I, we have
3 Stratified Mixture of Experts
3.1 Architectural Design
The guiding design principle for Stratified Mixture of Experts (SMoE) is to assign dynamic capacity to tokens. Given E experts in MoE block, partition them into strata. The 7th stratum has gate Gate; which routes tokens an expert in the current stratum as well as the subsequent ones. This means that tokens can never be sent back to the previous strata. Tokens keep getting routed in the SMoE block until they reach the final stratum. Different tokens will pass through varying number of experts, resulting in different capacities, according to the assignment of gates. In vanilla MoE, however, the capacity through which every token goes that of single FFN layer. The workflow of how a token passes a SMoE block is shown in Figure 1c. For example, some tokens the 1s stratum may be assigned to experts in the 2nd stratum while others are sent to the 3rd stratum. After several rounds of assignments, tokens finally exit the block after reaching the last stratum.
Ve {I....;F)) FFN)(x) w| ReLU(W! .)),
(1)
where w/(h and w:(o are the weights of FFN,. A trainable routing gate with weights W predicts scores for these experts to use for the input z, in the form of a routing vector G :
= softmax(W, xx).
(2)
We select the set of top-K experts, denoted with E C , and compute the output of the
MoE layer as follows:
Tout EGe- (F). 330
MoE models suffer from the notorious load imbalance issue, where the gate weights could collapse and send most tokens to the same expert. As a result, recent research has focused on designing better auxiliary load balancing loss functions encourage tokens be evenly distributed across experts, e.g., Lewis et al. (2021) formulated token-to-exper ex-ext allocation as a linear assignment problem, Roller et al. (2021) modified the feedforward layer to hash to different sets of weights depending on the current token, Zoph et al. (2022) proposed router z-loss that resolves instability issues, and Zhou al. (2022) reversely design an expert-to-token allocation algorithm. Other lines of investigation in MoE include regularization techniques such as gating dropout (Liu et al., 2022) and output masking (EOM and FOM) (NLLB Team et al., 2022), as well as novel MoE architectures, such as conditional MoF routing (CMR) that add an extra branch beside MoE layer (NLLB Team et al., 2022), Pyramid- Residual MoE (Rajbhandari al., 2022), hybrid dense and MoE model with more experts in the last layers.
In SMoE, the successive application of multiple FFN layers to a token can result in training instability. Therefore, following the approach of Vaswani et al. (2017); Wang et al. (2019); Xiong et al. (2020), we incorporate layer normalization (LayerNorm, Ba et al. (2016)) before dispatching tokens to experts and a residual connection after the tokens have passed through the experts. See Figure le for overview of the design of our stratified experts.
Formally, given T tokens in denote the d-dimensional representation of the ghh token in the ith stratum of the current SMoE block with Titt- Let & be the set of experts visible to the current gate (current stratum plus subsequent strata) and let F; [& be its cardinality. Before being dispatched FFN layers, tokens are firstly normalized with LayenNorm.
we
1 LayerNorm(a..)).
(4)
Then, Gate; with weights W; predicts a probability distribution G3 E R, scoring all visible E, experts at that stratum:
However, all previous work default to equal capacity for all tokens regardless of their language. frequency, or any other property. In the subsequent sections, we present Stratified Mixture of Experts (SMoE) model that automatically assigns dynamic capacities to different types of tokens.
G. Gate,(a(,)) softmax - (5)
Following Lepikhin et al. (2021), we dispatch each token to at most k=2 experts. If E is the set of selected top-k experts and the expert with

---

the highest score in the 5th (j > i) layer, æi, will be assigned to the jth layer and the output computation on the token is:
4.1 Datasets
In this study, we consider three datasets comprising 4, 15, and 94 languages each. The initial two datasets are extracted from the primary bitexts of the NLLB-200 training dataset, and adopt their resource-level categorizations: high-resource 2 IM), very low-resource (S 100K), and low- resource (the remaining). These two datasets are developed and evaluated using the Flores-200 dataset. The third dataset is OPUS- PP0 (Zhang et al., 2020). We follow Zhang et al. (2020) and divide directions into high-resource, medium- resource, and low-resource categories.
J = Σ (C.. eCE
(6)
where Gt.e is the gate score for the expert e. Finally, We employ residual connection after the FFN layer:
Zj Ij.t Tipt-
(7)
Tokens gradually pass experts deeper layers until they finish passing the final layer.
3.2 Load Balancing
NLLB M4 dataset From NLLB, we pick 4 languages from 4 different linguistic families, of which are high-resource and the other 2 are low-resource: Northern Sotho (nso, 526K parallel sentences), Malay (msa, 1M), Tagalog (tgl, 1M), Catalan (cat, 634K), totaling 3.2M training examples.
Similar Lepikhin
(2021), we encourage tokens be uniformly distributed across all visible experts. Each gate has loss term to balance the load. For Gate, the loss is:
=
Σ SaPes
(8)
BEE;
NLLB M15 dataset Taking into account linguistic diversity and larger data size, we expand the M4 dataset cover a of diverse 15 languages. M15 covers linguistic families and a balanced number of high- low-resource. and very low-resource languages (each category has 5 languages). We show detailed listing and information on the MI5 dataset in Appendix A.
where Se is the fraction of tokens dispatched to expert e, as their first choice, through top-k-gatte g-gting::
1 Σ 1 (arg max Gt,e = e]: fe T
and is the average routing probability to that expert over the T tokens in the mini-batch:
OPUS-100 addition to the datasets derived from NLLB, we also utilize OPUS- to examine scenario involving a larger number of languages. OPUS-100 encompasses a total of 100 languages, which supports 94 development/test language pairs.
T Σ Gie Pe = T
(10)
The auxiliary loss for the current SMoE block is computed by taking the average of the loss over all gates within the block.
Evaluation During inference, we use beam search with a beam size of 5 and a length penalty of 1.0. We report BLEU scores (Papineni et al., 2002) for models trained on NLLB dataset and sacrebleu (Post, 2018) with flores200 tokenizer for models trained on OPUS-100.
C a L
Ci
(11)
where a is a hyperparameter to control the strength of the load balancing loss. We average the loss over all SMoE blocks in the architecture the final auxiliary loss appended to the original task loss.
4.2 Baselines
We use five strong baselines to evaluate the effectiveness of SMoE. All baselines are own implementation following the settings from the original papers. Note that the total number of experts i.e., the full model capacity is kept constant
4 Experiments
Contrarr to NLLB Team et al. (2022). in our study, very low-resource is not included in the lowe-nsource category but considered as an independent sot.
We evaluate the proposed SMoE on a many-to- many multilingual neural machine translation task.

---

for all models in order to ensure fair comparison (8 experts for M4 and 16 experts for M15).
In MoE models, every other FFN layer of the encoder and decoder are substituted with an MoE layer. For SMoE, in the Sdh stratum, we enforce that each expert processes, at most, 2 x T:/ E, tokens, where 7; is the number of tokens in the mini-batch sent to the layer and E; is the number of visible experts. For the other MoE baselines, it is 2 x T/E, where T is the number of tokens in the mini- batch and E is the total number of experts. The multiplicative coefficient a for the auxiliary load balance loss is set to 0.01. vocabulary of size 32k for both M4 and M15 and 64K for OPUS-100 with SentencePiece (Kudo and Richardson, 2018). For a fair comparison, All models are trained for the same number of updates. More details can be found in Appendix B.
Vanilla MoE. An MoE model with Top-2 gating (Lepikhin et al., 2021).
Switch Transformer. An MoE model with Top-1 gating (Fedus et al., 2021). Switch Transformer was introduced to mitigate training instabilities and improve the efficiency of Top-2 MoE models. Note that switch transformer uses fewer FLOPs per token due to the top-1 gating approach.
MoE + EOM. A vanilla MoE model regularized with Expert Output Masking (EOM) (NLLB Team et al., 2022). EOM masks the expert output for a random fraction (Peom) of outputs. We set Peom = 0.1 following the suggestion of NLLB Team et al.
(2022)
4.4 SMoE configurations
We use series of numbers separated by hyphens to describe the SMoE configuration. For instance, SMoE-4-4-8 indicates that all MoE blocks have 3 strata, where the 1st stratum has experts, the 2nd has 4 and the 3rd has 8.
Conditional
MoE Routing (CMR).
CMR (NLLB Team et al., 2022) augments MoE with a binary gate that sends tokens to one of two branches: (1) a shared FFN layer and (2) an vanilla MoE layer. Note that this method requires extra parameters due to the added shared FFN layer. The CMR budget constraint is set to 0.8.
4.5 Results
M4 Results The results are in Table 1. We consider two configurations for SMoE: SMoE-4-4 and SMoE-2-2-2-2. The better SMoE settings for M4 is SMoE-2-2-2-2. In M4, SMoE-2-2-2-2 outperforms Switch Transformer by +0.9 BLEU on average and vanilla MoE by +0.74 BLEU, Out of 5 MoE baselines, CMR achieves the best performance (+0.4 BLEU over Switch Transformer and +0.3 BLEU over vanilla MoE), however, CMR models have more parameters well as more FLOPs per token.
Stacking MoE Layers. A similar model architecture to SMoE where we simply stack multiple MoE layers, e.g., stacking 2 MoE layers with experts each vs. one block of SMoE with 2 strata where each stratum has 4 experts. Unlike SMoE where each expert is surrounded with residual skip connection and preceded with a LayerNorm, here the stacking is naive without any addition.
4.3 Training Details
It worth noting that simply stacking MoE layers degenerates the model performance, which indicates the importance and effectiveness the specific design of SMoE.
Following Johnson et al. (2017), we prepend source sentences with special language token <2xxx> to indicate the target language. We use data sampling temperature of T-1 suggested by NLLB Team al. (2022) to train on NLLB datasets, and 5 suggested by Zhang et al. (2020) to train on OPUS-100
M15 results We show results in Table 2. We consider multitude of settings splitting the 16 experts per layer over 2, or 4 strata. The best SMoE settings for the larger M15 dataset SMoE-4-12. This configuration demonstrated an average improvement of +1.04 BLEU over the Switch Transformer and +0.93 BLEU over the vanilla MoE across the 15 languages evaluated. However, MoB+EMM and CMR only improve vanilla MoE by +0.52 and
The dense model architecture, backbone for all trained models, is Transformer model (Vaswani et al., 2017) with 12 layers (6 on encoder and 6 on decoder). We transformermase and E-8 experts for M4, and transformenbig and 16 experts for M15 and 100S
'transformenhas? FFN dimension of 2048, 8 heads, and embedding dimension of 512; transformenig:: FFN dimension
4096, 16 heads, and embedding dimension 1024.

---

XXX-
cng-tXX
Avg.
Paarameteers
FLOPs/tok
Methods experts)
Dense Model (Vuswani ul., 201 7)
28.87
37.01
41.76
40.35
26.68
44.30 30
38.20
45.50
37.83
60M
167M
Vanilla MoE (Lepikhin et al., 2021)
29.45
41.26
44.23
42.91
27.59
47.12
40.10
47.66
40.04
148M
192M
Switch Transformer (Fedus et al., 2021)
29.69
40.35
44.50
43.19
27.73
46.44
39,88
47.26
39.88
[48M
167M
McE+EOM (NLLB Team et al., 2022)
30.52
40.50
44.52
43.48
27.71
46.85
40.09
47.14
40.10
148M
192M
McE+CCM (NLLBB Team ul., 2022)
30.54
41.29
44.65
43.71
27.76
46.87
40.36
47.65
40.35
161M
217M
Stacking 2 2-4expert MoE
28.65
39.24
42.02
41.36
27.58
46.27
40.06
47.50
39.09
148M
242M
SMoE-4-4
30.14
41.05
45.13
43.97
27.86
47.48
39.85
47.21
40.40
148M
217M
Stacking 2-expert MoE
25.53
35.53
34
37.43
24.99
41.35
36.85
42.15
35.27
148M
327M
S96E-2-2-2-2
30.96
41.88
86
43.37
27.83
47.96
40.95
48.39
40.78
247M
Table Overall BLEU results on the M4 dataset. The best values are bold and the second-best values underlined.
The number of experts is 8 for all methods. The two SMoE models attain the two best performances across all languages.
eng
XXX-teng
Avg.
#Parameters
FLOPs/tok
Methods (detsult 16 experts)
alll
ligh
low
very low
all
high
low
Very low
Dease Model (Vaswani et ol., 2017)
29.27
38.68
32.44
16.09
29.94
37.55
32.37
19.89
29.61
200M
506M
Vanilla MoE (Lepikhin et al., 2021)
32.01
40.41
34.77
20.84
32.00
39.38
34.73
21.89
32,00
963M
606M
Switch Transformer (Fedus et al., 2021)
31.67
40.44
34.72
19.87
32.09
39.85
34.75
21.68
31.89
963M
506M
McE+EOM (NLLB
2022)
40.63
20.25
40.42
35.33
606M
McE+CMR (NLLB Team al., 2022)
31.83
40.22
34.81
20.46
33.02
40.41
35.42
23.21
32:42
1,01B
707M
SMoE-4-12
33,00
41.14
35.77
22.10
32,86
40.63
35.14
22.80
32,93
963M
656M
2-layer SMoE
SMoE- 12-4
32.52
41.11
35.56
20.89
32.94
40.03
35.21
23.60
32.73
963M
757M
SMoE-8-8
32.28
41.24
35.33
20.27
32.72
40.30
34,80
23.08
32.50
963M
707M
SMoE-4-4-8
32.83
41.61
35.82
21.07
32.24
39.40
34.53
22.80
32.54
963M
724M
3-layer SMoE
SMoE-8-4-4
32.24
40.99
20.61
33,06
40.36
35.12
23.71
32.65
963M
805M
4-layer SMoF
SMoE-4-4-4-4
32.87
41.50
35.68
21.42
32.63
40.03
34,50
23.35
32.75
963M
825M
Vanilla MoE. experts
32,98
41.05
35.69
22.19
32,90
40.79
35.49
22.42
32.94
1.77B
606M
Table 2: Overall BLEU results on the M15 dataset. The best values are bold and the second-best values are underlined. Unless otherwise mentioned, the number of experts is 16. All SMoB models outperform the baselines. The best setting is SMoE-4-12 which outperforms vanilla MoE by +0.93 BLEU. Vanilla MoE would require double its parameters to achieve similar performance to SMoE-4-12.
+0.53 BLEU, respectively. Note that if vanilla MoE wants to achieve similar performance to SM (32.94 vs. 32.93 BLEU on average), has to increase experts from 16 to 32, which almost doubles the total number of parameters from 963M to 1.77B (last row of Table 1), which means our model is much more parameter-efffcient.
94 languages OPUS-100 does not support
the remaining languages ) We select the two optimal SMoE configurations in M15: SMoE-4-12 and SMoE-4-4-4-4. The SMoE-4-12 configuration consistently demonstrates superior performance, achieving larger margin compared to our best baselines, EOM and CMR. The SMoE-4-12 outperforms naive MoE model and Switch Transformer by +1.01 and +1.63 BLEU, significantly outperforming the gains achieved by EOM (+0.11 and +0.73) and CMR (+0.14 and 0.76).
SMoE models with more strata, allowing for more depth, do not guarantee better performance. A clear example is SMoE-4-12 and SM0E-4-4-4-4 in M15 (32.93 vs. 32.75 averaged BLEU). However, for 'balanced' SMoE (equal number of experts per stratum), fewer experts per stratum achieves better performance: SMoE-2-2-2-2 outperforms SMoE-4, 4 (40.78 vs. 40.40 BLEU) on M4 and SMoE-4-4-4-4 outperforms SMoE-8, (32.75 vs. 32.50 BLEU) on M15.
Overall, SMoE outperforms all baselines with the same fewer number of parameters.
4.6 Computational Cost
As token may pass through multiple strata in any given SMoE layer, the average computational cost higher than other MoE models. In the last column of Tables and 2, we record FLOPs per token for all models, and Table shares the
OPUS-100 Results: A Larger Performance Gap. Table 3 presents the comprehensive results. Notably, the performance disparity becomes more pronounced we scale our experiments to

---

eng +XXX
XXX -xx-
Avg.
Methods
high
medium
high
medium
low
Dense Model (Vaswani et al., 2017)
27.37
23.89
31.17
29.76
30.60
29.40
31,85
31.49
28.99
Vanilla MoE (Lepikhin al., 2021)
30.34
26.16
34.78
33.38
32.38
31.20
33.67
33.21
31.36
Switch Transformer (Fedus et al., 2021
30.03
25.76
34.46
33.27
31.44
30.58
33.00
31.18
30.74
MoE+EOM (NLLB Team et al., 2022)
30.48
26.24
34.76
33.86
32.45
31.23
33.69
33.44
31.47
McE+CMR (NLLB Team et al., 2022)
30.56
26.32
34.98
33.77
32.43
31.33
33.69
33.11
31.50
SMGE-4-12 SMoE-4-4-4-4
32.15
27.99
36.67
35.06
32.58
32.42
33.76
31.37
32.37
31.84
27.31
36.52
35.33
32.71
31.87
33.80
33.04
32.28
Table 3: Overall BLEU results on the OPUS- 100 dataset. The best values bold and the second-best values are underlined. The number of experts is 16. We consider the two best settings in M15 dataset, SMoE-4-12 and
SMoE-4-4-4-4, Both of them substantially outperform all baselines. The number of parameters and FLOPs/tok for MoE models are the same Table 2.
22
2.25
Requesied Capacity (RC)
2.0
Requested Capacit (RC)
2.00
1.8
1.50
1.25
rigees
e00 no
09
ugaare
ergan
19 00
wonh
erc
erean
(a) Decoder
(b) Encoder
Figure Average requested capacity (RC) all tokens each translation direction. blue bars are for xxx xxx-eng directions and the purple bars for eng-xxx directions. Directions in each subset are sorted from higu-resource to low-resource. On the decoder side, the average RC of eng tokens is similar regardless of the source language, but averaged RC has large variance if the target language is different. On the encoder side, RC is always different even though the source language is the same.
same information with Table 2. Although SMoE requires more FLOPs per token, the additional cost only a marginal increase over vanilla MoE models. For example, in M15 and OPUS-100, our best setting SMoE-4-12 merely uses 8% more FLOPs/tok than other top-2-gating MoE models, but significantly outperforms all of them.
we examine three potential influencing factors: the language of the input token, the frequency of the token, and the depth of the SMoE block. All analysis is conducted using the SMoE-4-4-4-4 model trained on the M15 dataset.
5.1 The Language of The Input Token
5 Analysis
Here, we investigate whether different languages have different RCs. We begin with collecting the average RC in all translation directions for all tokens in the training and development (Flores- 200 dev) sets. We investigate SMoE blocks in the encoder and decoder separately as they process different tokens (source tokens the encoder and target tokens for the decoder). The average RC is then averaged across all SMoE blocks in either encoder or decoder.
The advantage of SMoE is that it can assign dynamic capacity to different tokens, with some tokens passing through only one expert and others passing through multiple experts. Here, we define the Requested Capacity (RC) as the average number of experts that token need to pass in one SMoE block. RC of a token is dictated by how the SMoE gates route it throughout the different strata. To understand what may affect the RC of tokens,
Figure 2a shows the average RC in the decoder for each translation direction in M15. When translating into English (xxx-eng, blue bars), we
'FLOPss are calculated for the forward pass as done in
Kuplan et al. (2020).

---

observe that the target English tokens have similar RC in the decoder side (x= 1.85 experts) irrespective of the source language. When translating from English (eng-xxx), the RC varies lot with respect to the target language.
403320
Than Phecuaty Bark
033220
Unlike in the decoder where only the target language matters, Figure 2b shows variability in RC with respect to both source and target languages i.e. not only the language of the tokens tes (source), but also the target language we will be translating into once move to the decoder. We hypothesize the special symbol at the beginning of the source sequence (<2xxx>) can affect the capacity assignment. In conclusion, capacity assignment sensitive to the target language in the decoder and to the translation direction in the
High Roquole Sipecit
Requbled
Figure 3: Violin plots of the token frequency in high- RC (left) and low-RC (right) tokens. Unlike high-RC tokens, low-RC RC tokens tend to be highly frequent ones.
2.1
2.02
2.0
Averoge Capardly
1.9
1.8
encoder.
1.69
1.7
1.53
5.2 Token Frequency
1.8
As in the previous section, record the average RC for all tokens in training and development data, and in all translation directions. To avoid looking at all 32K tokens in our vocabulary, we select the top-25 tokens with the highest RC each SMoE block and in each translation direction, totaling 4500 t We similarly collect the bottom- 25 tokens with the lowest RC. After removing tokens repeatedly selected by different directions end up with
1.5
Figure 4: Average RC of SMoE blocks in different locations of the architecture.
5.3 Location of The SMoE Block
We analyze in this section the average RC in relation the location of the SMoE block in the transformer architecture.
by different SMoE blocks, 2580 unique high-RC tokens and 3208 unique low- RC tokens. We draw in Figure 3 a violin plot show the distribution of tokens in these two groups in terms of their frequency in the training data. We rank the frequencies on the y-axis so that lower rank means more frequent tokens, e.g.. rank 0 corresponds to the most frequent token in training data. The results show that there is strong correlation between frequency and RC for tokens with the highest RC. On the other end of the spectrum, tokens with the lowest RC tend to high-frequency tokens, as indicated by the right violin plot being wider at the bottom part (rank < 10k). Many of these high-frequency tokens are basic subword units (like _li, _el, _pa) or punctuation marks. One can interpret RC as metric for 'difficulty in processing given token'. The model was overly exposed to these frequent tokens, and as such, does not require lot of capacity to process them.
depicted Figure 4, RC varies depending on the location of the SMoE block. Early encoder layers (encoder 2nd layer is the first SMoE block in the model) request more capacity than the subsequent encoder SMoE blocks. We hypothesize that this first block takes the task of mapping tokens coming from different languages and different scripts shared space.
Conclusion
This work presents Stratified Mixture of Experts (SMoE) models, a novel design for Molss that is capable of dynamically assigning capacity to input tokens. Through experimental evaluation at three scales (M4, M15, and OPUS- we have demonstrated that the proposed SMoE model surpasses the performance of many current state-of-the-art MoE methods. This proves that dynamically assigning capacity to tokens in MoE models viable solution address the MoE's parameter inefficiency. Additionally, we conduct a thorough analysis to investigate the factors that
$25 6 (#SMoE blocks) x 30 (#directions) 4500..

---

influence dynamic capacity assignment, including the language of the tokens and the location of the SMoE block within the model architecture.
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary. Guillaume Wenzek. Francisco Guzmán, Edouard Grave, Myle Out, Luke Zettlemoyer. and Veselin Stoyanov. 2020. Unsupervised cross-lingual representation learning at scale. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 8440 8451, Online. Association for Computational Linguistics.
Limitations
Stratified Mixture of Experts (SMoE) aims to improve the performance Mixture of Experts (MoE) models by assigning dynamic capacity to different tokens. While SMoE has demonstrated performance improvements over many state-of- the-art baselines, it also comes with additional computational cost compared traditional MoE models. However, the cost is small and the benefits SMoE in terms of improved performance often outweigh this added computational cost, especially in tasks where performance is critical. For example, in OPUS- with 8% FLOPs/tok, SMoE-4-12 achives + +1.01 BLEU compared with traditional MoE (Lepikhin et al., 2021).
Jacob Devlin, Ming-Wai Wei Chang. Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long and Short Pupers), pages 4171-4186, Minneapolis, Minnesota. Association for Computational Linguistics.
Nan Du, Yanping Huang. Andrew M Dai, Simon Tong, Dmitry Lepikhin. Yuanzhong Xu. Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. 2022. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning, pages 5547-5569 PMLR.
Acknowledgements
Angela Fan, Shruti Bhosale, Holger Schwenk. Zhiyi Ma. Ahmed Ei-Kishky, Siddharth Goyal, Mandeep Baines, Onur Celebi, Guillaume Wenzek, Vishrav Chaudhary, et al. 2020. Beyond english-centric multilingual machine translation. arXiv preprint arXiv:2010.111 11125.
We thank anonymous reviewers for their insightful feedback. We also extend gratitude to Steven Tan and Yunmo Chen for their valuable suggestions.
William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.
References
Roee Aharoni, Melvin Johnson, and Orhan Firat. 2019. Massively multilingual neural machine translation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 3874-3884. Minneapolis, Minnesota. Association for Computational Linguistics.
Ze-Feng Gao, Peiyu Liu, Wayne Xin Zhao, Zhong-Y Yi Lu, and Ji-Rong Wen. 2022. Parameter-efficien mixture-offexperts architecture for pre-trained language models. arXiv preprint arXiv:2203.0 104.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford. Diego de Las Casas, Lisa Anne Hendricks. Johannes Welbl, Aidan Clark, et al. 2022. An empirical analysis of compute-optimal large language model training.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. 2016. Layer normalization. arXiv preprint 10X.v:6007.06450.
Advances Neural Information Processing Systems.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter. Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc.
Melvin Johnson, Mike Schuster, Quoc V. Le, Maxim Krikun, Yonghui Wu, Zhifeng Chen, Nikhil Thorat, Fernanda Viégas, Martin Wattenberg, Greg Corrado, Macduff Hughes, and Jeffrey Dean. 2017. Google's multilingual neural machine translation system: Enabling zero-shot translation. Transactions of the Association for Computational Linguistics, 5:339 351.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. arXiv preprint 2001:000000800

---

Young Jin Kim, Ammar Ahmad Awan, Alexandre Muzio, Andres Felipe Cruz Salinas, Liyang Lu, Amr Hendy, Samyam Rajbhandari, Yuxiong He, and Hany Hassan Awadalla. 2021. Scalable and efficient moe training for multitask multilingual models. arXiv preprint 20Xiv:219.1000 104.3
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi. Ammar Ahmad Awan, Jeff Rasley, and Yuxiong Advancing mixture-
He. 2022. Deepspeed-most of-experts inference and training to power next-generation ai scale. In ICML.
Stephen Roller. Sainbayar Sukhbaatar, Jason Weston, et al. 2021. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566.
Diederik P Kingma and Jimmy Ba. 2014. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.690..
Taku Kudo and John Richardson. 2018. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 66-71, Brussels, Belgium. Association for Computational Linguistics.
Noam Shazeer, *Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference Learning Representations.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances neural information processing systems, 30.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2021. hard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2018. GLUB: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNIP Workshop BlackboxNL.P. Analyzing and Interpreting Neural Nenworks for NLP, pages 353-355, Brussels, Belgium. Association for Computational Linguisties.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. 2021. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, pages 6265 6274. PMLR.
Qiang Wang, Bei Li. Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. 2019. Learning deep transformer models for machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 1810-1822.
Rui Liu, Young Jin Kim. Alexandre Muzio, and Hany Hassan. 2022. Gating dropout: Communication- efficient regularization for sparsely activated transformers. In International Conference on Machine Learning, pages
PMLR.
Marta R NLLB Team, Costa-jussà, James Cross, Onur Çelebi, Maha Elbayad, Kenneth Heafield, Kevin Heffernan, Elahe Kalbassi, Janice Lam, Daniel Licht, Jean Maillard, et al. 2022. No language left behind: Scaling human-centered machine translation. arXiv preprint 2207:020000467
Ruibin Xiong, Yunchang Yang. Di He, Kai Zheng. Shuxin Zheng, Chen Xing. Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. 2020. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei- Jing Zhu. 2002. Bleu: method for automatic evaluation machine translation. In Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics. pages 311-318, Philadelphia, Pennsylvania, USA. Association for Computational Linguistics.
Biao Zhang, Philip Williams, Ivan Titov, and Rico Sennrich. 2020. Improving massively multilingual neural machine translation and zero-shot translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages (628-1639, Online. Association for Computational Linguistics.
Yanqi 7hou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen. Quoe Le, and James Laudon. 2022. Mf experts with expert choice routing. arXiv preprint arXiv:220..09360
Matt Post. 2018. A call for clarity in reporting bleu scores, WMT 2018, page 186.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. 2020. Exploring the limits of transfer learning with unified text-to-text transfomme. Journal of Machine Learning Research, 21:1-67.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022. Designing effective sparse expert models. arXiv preprint axXiv:2202.08900

---

Simiao Zuo. Xiaodong Liu, Jian Jiao, Young Jin Kim. Hany Hassan, Ruofei Zhang. Jianfeng Gao, and Tuo Zhao. 2021. Taming sparsely activated transformer with stochastic experts. In International Conference on Learning Representations.

---

A M15 Information
We show the detailed information of the M15 dataset in Table 4.
B Additional Training Details
We employ the transformerease model (with an FFN dimension of 2048 and an embedding dimension of 512) for the M4 dataset, and the transformerbig model (with an FFN dimension of 4096 and an embedding dimension of 1024) for M15 and OPUS-100 dataset. The maximum learning rate is 0.0008 for M4 and M15, and 0.0005 for the OPUS-100 dataset. The optimizer is Adam (Kingma and Ba, 2014) with inver inverse_sqrt learning rate scheduler and weight decay of 0. The total number of training steps is 100K with 8K warm-up steps. The batch size is 13K tokens for M4 and M15, and 65K tokens for OPUS-100.

---

Language
Language code
Parallel Data Size
Resource Level
Language family
Northern Sotho
nso
526 K
Low
Central Narrow Bantu
Rundi
run
454 K
Low
Central Narrow Bantu
Swati
SSw
94 K
Very Low
Central Narrow Bantu
Indonesiam
ind
6.5 M
High
Malayo-Polynesian
Malay
msa
M
High
Pollyssaon
Tagalog
tgl
1 M
High
Malayo-Polinesii Polinesian
Bokmäl (Norwegian)
nob
238 K
Low
North Germanic
Icelandic
isl
M
High
North Germanic
Faroese
fao
4
Very Low
North Germanic
Slovene
slv
15M
High
Southwestern Slavic
Luxembourgish
1tz
8K
Very Low
Western Germanic
Limburgish
lim
5K
Very Low
Western Germanic
Catalan
cat
634 K
Low
Western Romance
Galician
glg
195 K
Low
Western Romance
Friulian
fur
6 K
Very Low
Western Romance
Table 4: Information about the 15 languages in the M15 dataset.


Reviews:
[{"review_id": "TsTYGYGLfY", "reviewer": "Reviewer_T5W3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We genuinely appreciate the valuable feedback provided by the reviewer and have addressed them in a point-by-point manner below. We are more than willing to engage in further discussions with the reviewers should any follow-up questions arise.\n\n### **Regarding your concern about incremental contribution**\n>The technical contribution is incremental. The main novelty comes from restructuring the MoE system to dynamically assign model capacities for different inputs\n\nThanks for your comments! Regarding your concern about the incremental contribution, we are not claiming a fundamental change to an underlying approach. Moreover, we also provided complete experiments and strong baselines to support the effectiveness of our methods. Indeed, our hope is that the relative ease in understanding and reproducing our proposals is why they are most likely to be picked up by the community. We would be happy if the common reaction by the reader was, 'Hunh, that seems easy!' we would say that a key point of this article is the relative straightforwardness of the proposed.\n\n### **Regarding your concern about experiments**\n>The experiment is limited to machine translation tasks \n\nThank you for your insightful comments. We wish to emphasize that multilingual machine translation is particularly apt for evaluating the efficacy of MoE models. This is because it mandates MoE models to handle tokens from diverse languages and resources, in order to mitigate negative language interference. Additionally, the parameter inefficiency is also usually observed under task that needs a large amount (millions) of datasets, where only machine translation data has this scale. As highlighted in line 102, our goal centers on the multilingual machine translation task, and to deepen our study, we conduct comprehensive experiments from 3 different types of MMT datasets (ranging from 4 to 95 languages). Nonetheless, we concur with your viewpoint that expanding the scope to encompass more tasks could further enrich and fortify our paper.\n\n>and some efficient MoE baselines are not compared [1].\n\nWe'd like to clarify that the methodology presented in [1] follows a distinct trajectory in studying MoE models. Their approach leverages pre-trained architectures like T5 or GPT2, whereas our model is designed and trained from scratch. Crucially, when selecting our baselines, we opted for MoE model architectures that shared our primary objective of enhancing multilingual machine translation. This ensures that our comparisons are both robust and meaningful. Furthermore, our experiments on MMT are comprehensive, spanning from 4 to 95 languages, which provides a holistic perspective on language distribution and diversity within the task.\n\nReference:\n[1] Parameter-Efficient Mixture-of-Experts Architecture for Pre-trained Language Models. Gao et al. COLING22.\n\n### **Regarding your concern about comparing with other SOTA non-MoE  MMT methods**\n>The comparison between SMoE and other Sota methods for MT is missing, hence it is not clear whether it is necessary to have multiple experts.\n\nThank you for your constructive feedback! It is hard to define what is the Sota non-MMT model because of the difference in focused languages, model size, and settings (like different vocabulary). Moreover, it has been widely proved that MoE model outperforms the non-MoE models under the same FLOPs/tok. For example, NLLB-EOM and CMR [1] are better than their dense model which is trained on an extensive dataset, utilizing back translation and curriculum learning.\n\nHowever, we value and follow the suggestion of the reviewer and conduct a non-MoE method on OPUS-100. The additional baseline we pick is CLSR [2], which is not only widely recognized as a strong MMT model but also shares a similar spirit/motivation with our study in the context of multilingual neural machine translation. It introduces a gating mechanism that controls whether a given token should be processed by a language-specific parameter or shared parameter (but their capacity for tokens is fixed).  As illustrated below, our approach markedly surpasses CLSR in average performance:\n\n| Methods   | en-xx all | en-xx high | en-xx med | en-xx low | xx-en all | xx-en high | xx-en med | xx-en low | avg.      |\n|-----------|:---------:|:----------:|:---------:|:---------:|-----------|------------|-----------|-----------|-----------|\n| CLSR      | 27.39     | 23.91      | 31.17     | 29.77     | 32.43     | 30.73      | **33.85** | **34.20** | 29.91     |\n| SMoE-4-12 | **32.15** | **27.99**  | **36.67** | **35.06** | **32.58** | **32.42**  | 33.76     | 31.37     | **32.37** |\n\nReference:\n[1] NLLB Team, Costa-jussà MR, Cross J, Çelebi O, Elbayad M, Heafield K, Heffernan K, Kalbassi E, Lam J, Licht D, Maillard J, Sun A. No language left behind: Scaling human-centered machine translation. arXiv preprint arXiv:2207.04672. 2022 Jul 11.\n[2] Zhang B, Bapna A, Sennrich R, Firat O. Share or not? learning to schedule language-specific capacity for multilingual translation. International Conference on Learning Representations 2020 Oct 2.", "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": "1. The paper develops a Stratified Mixture of Experts (SMoE) to address the parameter inefficiency in Mixture-of-experts (MoE) models, which is interesting and reasonable.\n\n2. The SMoE outperforms multiple MoE models on multiple machine translation tasks impressively.", "reasons_to_reject": "1. The technical contribution is incremental. The main novelty comes from restructuring the MoE system to dynamically assign model capacities for different inputs.  \n\n2. The experiment is limited to machine translation tasks and some efficient MoE baselines are not compared [1]. \n\n3. The comparison between SMoE and other Sota methods for MT is missing, hence it is not clear whether it is necessary to have multiple experts.\n\n[1] Parameter-Efficient Mixture-of-Experts Architecture for Pre-trained Language Models. Gao et al. COLING22.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This work proposes a new model called Stratified Mixture of Experts (SMoE) to address the issue of parameter inefficiency in Mixture-of-experts (MoE) models. The authors hypothesize this issue is a result of all experts having equal model capacity, which cannot provide adaptive complexity requirements of different tokens or tasks. They propose to build an MoE system with a stratified structure that can assign dynamic capacity to different tokens. They conducted experiments on two multilingual machine translation benchmarks, where it outperforms multiple MoE models with on-par or fewer model parameters.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "line 263 'extracte' --> 'extracted'.", "Missing_References": "I suggest adding references (paper and resources) for the training data NLLB into the paper. "}, "extra_text": {}}, {"review_id": "IKxgItOvsP", "reviewer": "Reviewer_3H55", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We genuinely appreciate the valuable feedback provided by the reviewer and have addressed them in a point-by-point manner below. We are more than willing to engage in further discussions with the reviewers should any follow-up questions arise.\n\n### **Regarding your concern about expert parallelism**\n>How compatible is the SMoE with expert parallelism? \n\nThanks for your detailed question! We always consider expert parallelism and use ALLTOALL communication, following the algorithm in Gshard [1]. We evenly distribute experts to our GPUs. If we have 16 experts and 4 GPUs, we group every 4 experts into one GPU regardless of the stratum of where they are located. We will clarify this in the paper.\n\nReference:\n[1] Lepikhin D, Lee H, Xu Y, Chen D, Firat O, Huang Y, Krikun M, Shazeer N, Chen Z. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668. 2020 Jun 30.\n\n### **Regarding your concern about overclaiming Sota**\n>The baselines are indeed popular but it is a bit overclaim to say these MoE models are SoTA\n\nThanks for your comments! We would like to clarify that our MoE model is designed to focus on the multilingual machine translation task. Consequently, we selected models that align with our goal (given that expert-choice MoE isn't primarily designed for MMT). Notably, NLLB-EOM/CMR stands out as one of the most outstanding MoE models tailored for MMT. This is the main reason for adopting them as baselines — making the baselines more relevant and meaningful. We acknowledge that Brainformer, being a concurrent work, wasn't available when our study was finalized. Nevertheless, we concur with the reviewer's perspective on the value of incorporating a broader range of baselines, and we appreciate the suggestion to gauge SMoE's standing amidst other robust MoE models.", "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": "1) Stratified MoE structure is very novel and interesting. Different tokens can also have an adaptive computation budget in this case, which may make many impossible tasks possible. (See [1])\n2) The paper is well-written and easy to read.\n3) Section 5 is great. It conveys a more comprehensive understanding about how the model works.\n\n[1] https://arxiv.org/abs/1807.03819", "reasons_to_reject": "1) How compatible is the SMoE with expert parallelism? It seems that the stratified structure requires more carefully implemented expert parallelism. Or authors actually ignored the expert parallelism in this paper?\n2) The baselines are indeed popular but it is a bit overclaim to say these MoE models are SoTA. We have many better models like Brainformer and Expert-choice MoE. It is okay to compare these popular models since some of these advanced models are not open-sourced, but please do not call them SoTA.\n\n\nMinor: Fig 1b and Fig 1c are almost the same. It is okay to keep 1b only.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This paper devotes to making MoE-based transformers more parameter efficient. To this end, the authors propose a Stratified Mixture of Experts (SMoE). That is, authors use a stratified MoE structure that enables different tokens to have different capacities. The SMoE model outperforms many popular MoE baselines.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "Another Parameter Efficient MoE:\nhttps://arxiv.org/abs/2107.11817"}, "extra_text": {}}, {"review_id": "fclV8jSlKC", "reviewer": "Reviewer_zL3X", "summary": "", "questions": "Could you give me explanation or external information for my reasons to reject?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We genuinely appreciate the valuable feedback provided by the reviewer and have addressed them in a point-by-point manner below. We are more than willing to engage in further discussions with the reviewers should any follow-up questions arise.\n\n### **Regarding your concern 1**\n>However, the resulting performance improvement of the architecture is marginal (MoE+CMR 40.35 vs. SMoE-2-2-2-2 40.78).\n\nThank you for the insightful feedback. It's imperative to underscore that while the absolute difference may appear marginal in smaller datasets, these results are averaged over 8, 30, and 190 directions. The intricacies of the dataset can potentially reduce the improvements for all MoE methodologies. Taking the example from the reviewer, indeed, SMoE outperforms CMR on M4 by +0.43 BLEU. However, SMoE-2-2-2-2 outperforms vanilla MoE by +0.74 BLEU but CMR only achieves  +0.31 BLEU improvement. \n\nFurthermore, our methods consistently outperform all strong baselines across all three MMT datasets (M4, M15, and OPU-100), encapsulating varying data size distributions and linguistic diversity. Notably, in OPUS-100, our SMoE registers a commendable +1.01 BLEU advantage over vanilla MOE, whereas CMR, despite demanding more parameters and FLOPs/tok, provides just a +0.14 BLEU increment.\n\n### **Regarding your concern 2**\n>The paper focuses solely on machine translation tasks.\n\nThank you for your insightful comments. We wish to emphasize that multilingual machine translation is particularly apt for evaluating the efficacy of MoE models. This is because it mandates MoE models to handle tokens from diverse languages and resources, in order to mitigate negative language interference. Additionally, the parameter inefficiency is also usually observed under task that needs a large amount (millions) of datasets, where only machine translation data has this scale\n\nAs highlighted in line 102, our goal centers on the multilingual machine translation task. Nonetheless, we concur with your viewpoint that expanding the scope to encompass more tasks could further enrich and fortify our paper.\n\n### **Regarding your concern 3**\n>The paper focuses only on token-level of MoE. Testing on task-level and sequence-level would strengthen the paper.\n\nThank you for your insightful feedback. Our exploration of SMoE models stems from the hypothesis that tokens from distinct languages may require varying capacities. Consequently, we begin with token-level MoE, upon which many strong MoE systems for MMT are built. This approach also benefits from expert parallelism due to the (almost) even distribution of tokens across experts by using load balance loss. Nonetheless, we concur with your suggestion that incorporating task-level or sequence-level perspectives would be beneficial.\n\n### **Regarding your concern 4**\n>Comparing to state-of-the-art non-MoE models on machine translation could better highlight if the improvements are coming from the stratified architecture specifically or just having more parameters/capacity in general.\n\nThank you for your constructive feedback. We would like to first highlight that the improvements have been proved coming from the stratified architecture because our SMoE is compared with other MoE models under the same number of parameters.\n\nSecondly, it is hard to define what is the sota non-MMT model now because of the difference in focused languages, model size, and settings (like different vocabulary). Moreover, it has been widely proved that MoE model outperforms the non-MoE models under the same FLOPs/tok. For example, NLLB-EOM and CMR is better than its dense model which is trained on an extensive dataset, utilizing back translation and curriculum learning.\n\nHowever, we value and follow the suggestion of the reviewer and conduct a non-MoE method on OPUS-100. The method we pick is CLSR [1], which is not only widely recognized as a strong MMT model but also shares a similar spirit/motivation with our study in the context of multilingual neural machine translation. It introduces a gating mechanism that controls whether a given token should be processed by a language-specific parameter or shared parameter (but their capacity for tokens is fixed).  As illustrated below, our approach markedly surpasses CLSR in average performance:\n\n| Methods   | en-xx all | en-xx high | en-xx med | en-xx low | xx-en all | xx-en high | xx-en med | xx-en low | avg.      |\n|-----------|:---------:|:----------:|:---------:|:---------:|-----------|------------|-----------|-----------|-----------|\n| CLSR      | 27.39     | 23.91      | 31.17     | 29.77     | 32.43     | 30.73      | **33.85** | **34.20** | 29.91     |\n| SMoE-4-12 | **32.15** | **27.99**  | **36.67** | **35.06** | **32.58** | **32.42**  | 33.76     | 31.37     | **32.37** |\n\nReference:\n[1] Zhang B, Bapna A, Sennrich R, Firat O. Share or not? learning to schedule language-specific capacity for multilingual translation. International Conference on Learning Representations 2020 Oct 2.\n\n### **Regarding your concern 5 and concern 6**\n>The increased computational cost of SMoE compared to standard MoE is not thoroughly analyzed. The trade-off between improved performance and higher FLOPs should be clarified.\n\nThanks for your insightful feedback!  To understand the optimal configuration of SMoE, we have a comprehensive hyperparameter search in Table 2 including varying the depths and number of experts, and we report all their detailed performance and the FLOPs/tok they cost. While no straightforward correlation emerged between improved performance and higher FLOPs/tok, our findings did spotlight SMoE-4,12 as the most efficient design, incurring only an 8% overhead in FLOPs/tok.\n\n>The SMoE model requires carefully tuning the number of strata and experts per strata for good performance. The paper does not provide much insight into how to set these hyperparameters.\n\nWe identified a pattern where a shallow SMoE with 1/3 experts in lower strata and 2/3 experts in higher strata performs best. This trend held true across experiments on both M15 and OPUS-100. Based on these observations, we would recommend following a similar strategy of stratified layers on other benchmarks. We will make sure to mention this in our conclusion.", "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": "1. It tackles an important problem in improving the parameter efficiency of mixture-of-experts (MoE) models. Recent work has shown that MoE models tend to be parameter inefficient, so addressing this issue could lead to more efficient large-scale models.\n\n2. The proposed Stratified Mixture of Experts (SMoE) model introduces a novel stratified architecture that allows dynamic capacity assignment to tokens. This concept of adaptive capacity seems promising for improving model performance.\nThe paper demonstrates strong empirical results, substantially outperforming several state-of-the-art MoE models on multilingual machine translation benchmarks while using the same or fewer parameters. The gains are particularly pronounced on the largest dataset with 94 languages.\n\n3. The analysis provides useful insights into what factors affect the capacity assignment, shedding light on how the model adapts capacity based on token language, frequency, and position in the architecture.", "reasons_to_reject": "1. You propose that this parameter inefficiency arises from all experts having equal capacity, and then introduce SMoE as a solution, assigning dynamic capacity to different tokens. However, the resulting performance improvement of the architecture is marginal (MoE+CMR 40.35 vs. SMoE-2-2-2-2 40.78). Similar observations are evident in other experiments. Consequently, this evidence is insufficient to validate your hypothesis.\n\n2. The paper focuses solely on machine translation tasks. The benefits of the proposed SMoE model for other NLP tasks are unclear. Testing on a broader set of tasks would strengthen the paper.\n\n3. The paper focuses only on token-level of MoE. Testing on task-level and sequence-level would strengthen the paper.\n\n4. The comparison is only to other Mixture-of-Experts models. Comparing to state-of-the-art non-MoE models on machine translation could better highlight if the improvements are coming from the stratified architecture specifically or just having more parameters/capacity in general.\n\n5. The increased computational cost of SMoE compared to standard MoE is not thoroughly analyzed. The trade-off between improved performance and higher FLOPs should be clarified.\n\n6. The SMoE model requires carefully tuning the number of strata and experts per strata for good performance. The paper does not provide much insight into how to set these hyperparameters.", "questions_to_authors": "Could you give me explanation or external information for my reasons to reject?", "paper_topic_and_main_contributions": "This paper proposes a new architecture for mixture-of-experts (MoE) models called Stratified Mixture of Experts (SMoE). The key ideas and contributions are:\n\n1. MoE models suffer from parameter inefficiency, where adding more experts does not improve performance proportionally. The authors hypothesize this is due to all experts having equal capacity, while different tokens may require different amounts of capacity.\n\n2. SMoE assigns dynamic capacity to tokens by having multiple strata of experts. Tokens can pass through more experts and get more capacity if routed across multiple strata. This allows more efficient use of parameters.\n\n3. SMoE is evaluated on multilingual machine translation across 3 datasets with 4, 15 and 94 languages. It substantially outperforms strong MoE baselines like Switch Transformer, MoE+EOM and MoE+CMR with the same or fewer parameters.\n\n4. Analysis shows the capacity assignment in SMoE depends on factors like target language, source language, token frequency and location of the SMoE block.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a dynamic capacity approach for MoE models to mitigate parameter inefficiency. The reviews expressed concerns about (i) the necessity of applying the approach to tasks beyond MT, (ii) the marginal improvements in MT, (iii) the time requirements and computational efficiency of the approach. The rebuttal has (partly) addressed the concerns. Stronger results on MT and applying the approach n tasks beyond MT certainly strengthen the paper in future.
Decision: Accept-Findings