Mixture-of-Experts with Expert Choice Routing
Authors: Yanqi Zhou
Tao Lei
Hanxiao Liu
Nan Du
Yanping Huang
Vincent Y Zhao
Andrew M. Dai
Zhifeng Chen
Quoc V Le
James Laudon
Venue: NeurIPS 2022 Accept (2022)

Abstract: Sparsely-activated Mixture-of-experts (MoE) models allow the number of parameters to greatly increase while keeping the amount of computation for a given token or a given sample unchanged. However, a poor expert routing strategy (e.g. one resulting in load imbalance) can cause certain experts to be under-trained, leading to an expert being under or over-specialized. Prior work allocates a fixed number of experts to each token using a top-k function regardless of the relative importance of different tokens. To address this, we propose a heterogeneous mixture-of-experts employing an expert choice method. Instead of letting tokens select the top-k experts, we have experts selecting the top-k tokens. As a result, each token can be routed to a variable number of experts and each expert can have a fixed bucket size. We systematically study pre-training speedups using the same computational resources of the Switch Transformer top-1 and GShard top-2 gating of prior work and find that our method improves training convergence time by more than 2×. For the same computational cost, our method demonstrates higher performance in fine-tuning 11 selected tasks in the GLUE and SuperGLUE benchmarks. For a smaller activation cost, our method outperforms the T5 dense model in 7 out of the 11 tasks.

Mixture-of-Experts with Expert Choice Routing
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Y. Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon
Google, Mountain View, CA, USA (yanqiz, taole, hanxiaol, dunan, huangyp, vzhao, adai, zhifenge, qvl, jlaudon)@google..oo
Abstract
Sparsely-activatte ceaturr-of-experree (MoE) models allow the number of param- eters to greatly increase while keeping the amount of computation for given token or given sample unchanged. However, a poor expert routing strategy can cause certain experts to be under-trained, leading to an expert being under over-specialized. Prior work allocates fixed number of experts to each token using top-k function regardless of the relative importance of different tokens. To address this, we propose a heterogeneous mixture-of-expere employing an expert choice method. Instead of letting tokens select the top-k experts, we have experts selecting the top-k tokens. As result, each token can be routed to variable number of experts and each expert have fixed bucket size. We systematically study pre- pre-training speedups using the same computational resources of the Switch Transformer top-1 and GShard top-2 gating of prior work and find that method improves training convergence time by more than 2x. For the same computational cost, our method demonstrates higher performance in fine-tuning 11 selected tasks in the GLUE and SuperGLUE benchmarks. For a smaller activation cost, our method outperforms the T5 dense model 7 out of the 11 tasks.
Introduction
Scaling up model capacity, dataset size, and training time has demonstrated huge success enhancing the performance of computer vision architectures [4] [3] [4] as well neural language models [2]
20 26 27]. The final model quality has been found to have a power-law relationship with the amount of data, model size, and compute time [16 20). However, training efficiency, which is defined as the total amount of computation used to achieve superior model quality than the state the art
system (21), should receive greater attention
increase efforts towards green Al (29
Sparsely gated mifture-of-experts [31 (MoE) provides an effective way scale model capacity given a fixed computational cost, and has recently played an important role in increasing the training efficiency large-scale language models [10] 21 MoE operate by adopting a number experts, each sub-network, and by activating only one a few experts for each input token. A gating network must be chosen and optimized order to route each token to the most suited expert(s). For example, recent work has implemented sparse routing via k k-means clustering [12), linear assignment to maximize token-expert affinities [22], or hashing (8 28). Many the prior work use a routing strategy concerning the token choice, where each token selects the best one or two experts.
We argue that the independent token choice of prior work often leads an imbalanced load of experts, which causes training inefficiency and sub-optimal training of the model. In order to mitigate this
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

FINE
FNVZ
p-00
Figure 1: High-level Comparison Between Conventional MoE and expert choice MoE.
issue, previous sparsely gated networks introduce additional auxiliary losses as regularization to prevent too many tokens being routed to single expert, but the effectiveness is still limited. Recent explore alternative strategies for routing, but they focus on pre-training only and do not demonstrate performance gain on downstream tasks. Moreover, none of the previous methods consider allocating a variable number of experts to each token based on importance, which can be beneficial.
approaches [ 22 28
We propose a very simple yet effective routing method we are calling expert choice. Unlike conven- tional MoE where tokens select one or two top-scoring experts, our method lets each expert pick the top-k tokens. Our method guarantees perfect load balancing, allows variable number of experts for each token, and achieves substantial gains in training efficiency and downstream performance as demonstrated in our experiments. Our major contributions include:
We identify common pitfalls in conventional MoE such as load imbalance as described in Section We then propose heterogeneous, expert choice method to provide a fluid allocation of model parameters based learnt fiken-to-expertt importance. This method intrinsically guarantees load balance without imposing auxiliary loss.
We show our method provides over 2x faster training convergence in a 8B/64E (8 billion activated parameters, 64 experts) model, compared to the top-1 1 and top-2 gating counterparts in Switch Transformer
[10] and GShard [21].
We show our method demonstrates strong scaling when increasing the number of experts from 16 to 128, evaluated in training perplexity.
We show our method demonstrates strong performance on downstream tasks selected from GLUE and SuperGLUE at all the evaluated scales. More specifically, our 8B/64E model outperforms T5 1 IB dense model in out of 11 tasks evaluated.
2
Related Work
Various approaches have been proposed to scale up neural network capacity to improve performance. Recent works have successfully scaled models to billions parameters via various Model parallelism 30 splits weights and tensors across multiple cores while pipeline parallelism [18] 24 splits different layers across devices with micro-batches pipelined to the different layers. To enable continued scaling of neural networks, improving model training and serving efficiency has become a critical research area.
Scaling:
forms of model parallelism (2) 21 26 27 33
Conditional Computation: Computation decisions can be made dynamically based on the input Conditional computation has been proposed as way increase the capacity of deep neural network without increasing the amount of computation, by activating certain parameters and computation on demand, on per-example or per-token basis (3]. Conditional convolution layers with task gating has been used to combat catastrophic forgetting when sequence of learning problems are optimized. The gating decisions may binary or sparse and continuous, stochastic or deterministic.
23
25

---

Mixture of Experts: Sparsely-gated MoE [31 is the first model to demonstrate massive improve- ments in model capacity, training time, or model quality with gating. Switch Transformer [TO] simplifies the gating by selecting only the top expert per token using softmax over the hidden state and demonstrates better scaling than previous work. All the prior work requires an auxiliary loss to explicitly encourage balancing. This loss term has to be carefully weighted to not overwhelm the primary loss. However, auxiliary loss does not guarantee balancing and a hard capacity factor has to be imposed. As result, many tokens can still be unprocessed by the MoE layer. Hard MoE [12] with a single decoding layer can be efficiently trained to good effect on large scale hashtag prediction tasks. Base Layers [22] formulate linear assignment that maximizes token-expert affinities while ensuring each expert receives an equal number of tokens. Hash layers [8] 28 devise hashing techniques on input tokens. However, the evaluations are limited to pre- mre-trainin perplexity. THOR [37] randomly activates experts during training and inference and is trained with a consistency regularization loss. THOR has demonstrated strong performance on translation tasks. Different from these prior works, our method is a learnt method that enables heterogeneous MoE and effectively improves downstream fine-tuning performance.
3 Method
We first identify a few pitfalls in the routing method of conventional mixture-of-experts (MoE) models and then present our method using expert choice to tackle these problems.
3.1 Pitfalls of Token-Choice Routing
MoE can be computationally advantageous compared to a dense model, routing strategy must be used to assign each token to the most-suited experts. Conventional MoE models employ token-choice We argue that this routing which independently selects the top-k experts for each token 10 21 31..
strategy has a few pitfalls that lead to sub-optimal training.
Load Imbalance: Token-choice routing often lead to poor load balancing across experts. That is. some experts may be trained with most tokens, leaving the remaining experts under-utilized. Experts can be under specialized because a lot of model capacity the under-utilized experts are wasted. On the other side, some tokens will not be processed, since over-utilized experts can only take a maximum number of tokens each step in order to avoid running out of memory. Load imbalance can also hurt step latency, thus inference time, as the step latency can be determined by the most loaded expert. Previous methods add an auxiliary loss on load balancing to mitigate the issue. However, this auxiliary loss does not guarantee balanced load, especially during the important early stages of training. Indeed, we empirically observe that the over-capacity ratio can reach 20%-400 % for some experts in token choice routing, indicating that significant portion of the tokens routed to these experts will be dropped.
Under Specialization: Each MoE layer uses a gating network to learn token-to-expe affinity. Ideally, the learnt gating network should produce the affinity such that similar or relevant tokens are routed to the same expert. A sub-optimal strategy can produce redundant experts and/or experts that are not sufficiently specialized. Under specialization may result by imposing an large auxiliary loss which favors more load balanced but less effective routing, Finding the right balance on the auxiliary loss to promote both load balancing and specialization challenging for token-choice routing.
Same Compute for Every Token: Finally, in a token-choice strategy each token receives exactly k experts and therefore occupies the same amount of compute. We hypothesize that this is not necessary nor desired. Instead, a MoE model should flexibly allocate its compute resource based on the complexity of the input. Motivated by the aforementioned observations, we next describe simple yet effective method which produces load balanced assignments based on expert choice.
3.2 Heterogeneous MoE via Expert Choice
Different from conventional routing, an expert choice method independently selects top-k tokens for each expert, where & is fixed expert capacity (i.e. the number of tokens each expert can take). Despite its simplicity, expert choice achieves perfect load balancing by design. It also enables more flexible allocation of model compute since tokens can be received by variable number of experts.
3

---

In our experiments, we set & as
k: 1 X
where n is the total number of tokens in the input batch (such as batch size x sequence length), c is the capacity factor, and the number of experts. The capacity factor denotes on average how many experts are utilized by a token. Given input token representations X E IRad where d is the model hidden dimension, our method produces token-to-expert assignment denoted by three output matrices I, G and P. The matrix / is an index matrix where [(i, j] specifies j-th selected token of the i-th expert. The gating matrix G E Ra denotes the weight of expert for the selected token, and P [ refers an one-hot version of that will be used to gather tokens for each expert. These matrices are computed using gating function,
S = Softmax X Wo), "x
G, 7 = TSp k), P Onehot(TT
where S denotes the token-to-expert affinity scores, Wg Rdxe denotes the expert embeddings, and T)p selects the k largest entries for each row of S
Similar to Switch Transformer [10] and GShard [2]]. we apply mixture of experts and the gating function in the dense wera-f (FFN) layer, it is the most computationally expensive part in a Transformer-based network. The input to the gated FFN, denoted by Xin E gexkxd, is produced using the permutation matrix P. Here [ Rxd denotes the input of the i-th expert. Similarly, let W1 and W2 denote the parameters of gated FFN in which W1 and W2 [i] E Raxd" denote the parameter matrices of the i-th expert. We compute the output of each expert [e[i] [i] as follows,
Xin = P
W: X. = Ge.l(XX..ii] Will) W
(3)
We omit the bias terms here for brevity. The finally output of the gated FFN layer Xour E Raxd can be obtained given Xe, the permutation and gating matrices P and G,
Xoun!!, d = Σ Plij. G[i, j.
(4)
Both X, and Xont can be efficiently computed using Einstein summation (einsum) operations.
3.3 Expert Choice with Additional Constraint
We also consider regularizing expert choice routing by limiting the maximum number of experts for each token. We are interested in whether adding this constraint improves pre-training and fine- tuning results. More importantly, it helps analyzing to what degree using variable number of experts per token affects the model performance.
Let A E Rex be positive matrix where Ali, j] represents whether the i-th expert selects j-th token. We solve the following entropy-regularized linear programming problem
max A
(ST,A) + (A)
s.t.
VE: Σ
A[i,j']=
Vj Σ AR
Vi,j: 0 0 A(i S
where < ST. .A > denotes the inner product, ((A) is the sum of element entropy and b > is an integer that upper bounds the selection for each token. Adding a small entropy term gives near-integer solution while enabling a fast iterative solver we can run on TPUs. Specifically, the solution space is the intersection of three convex sets each satisfying one of the linear constraints. 99) that alternatively projects the intermediate solution onto one of the convex sets After A is computed, the routing indices I is selected using TopK (A, k) instead.
We use Dykstra's 's algorithm
H(AA) Eis -A(i, log Ali,j]
We use À 0.001 and maximum of 100 iterations.

---

Model
Type
Ppprrams
Tact-garams
M
Theads
dheed
O.IB
Dense
130M
130M
0.1B/16E
548M
145M
0.1B/32E
MoE
1.0B
145M
12
768
3,072
12
64
B/64/
MoF
1.98
145M
0.1B/128R
MoF
3.7B
145M
128
8B
Dense
8.7B
8.7B
4.096
16,384
8B/64E
MoE
143B
9.8B
64
Table 1: Sizes and architectures of both MoE and dense models that were trained in our experiments. Models are grouped by the number of activated parameters per token. All trained models share the same learning hyperparameters described in Section
3.4 Model Architecture
At the high level, we adopt the idea of sparsely activated Mixture-of-Expers (MoE) [31] We use Transformer architecture and replace the feed-forward component of every other Transformer layer with a MoE layer, following recent practice 21). Interleaving regular Transformer layers and MoE layers empirically improves model performance and training efficiency, probably because forcing some shared components in between MoE layers can mitigate the negative effects skipping tokens. Several additional modifications adopted in recent work have been applied in our experiments. For example, we replace the standard positional embedding with per-layer relative positional bias [5]. In the non-MoE fred-forward sub-layers (only every other layers MoE layers), replace the first linear projection and the activation function with the Gated I Linear Unit (). which computes the component-wise product of two linear transformation of the input, followed by Gaussian Error Linear Unit [15] activation function.
As described earlier, each MoE layer consists of group of independent feed-forward networks as denoted as "experts". The gating function in Eq. (2) uses a softmax activation function to model probability distribution over these experts. This distribution denotes the preference over experts of each incoming token, which is computed similarly in a conventional gating network [10] 21 31)
During training, each MoE layer's learnable gating network described in Eq. (2) is trained to the input to activate the best subset of experts using a top-k function along the token dimension. An "shuffle" stage and "unshuffle" stage inserted the MoE layer, where the first stage gathers the
tokens to their designated experts while the second stage permutes the tokens back to their original order in the input batch. This step is formulated in Eq. and Eq. (4).
Similar to conventional MoE method, there are more parameters in the MoE layer. However, the activated model size per token can be comparable to dense layer because during training inference, only limited subset of experts is activated for any given token. For instance, Switch Transformer [TO] has only one activated expert while GShard [2]] uses two experts per token. In method, the number activated experts can vary for each token but the overall computation is kept the same as set
the baseline architectures by fixing the capacity factor in Eq. (T. Unless otherwise specified, c= 2 such that our method can be directly compared to the top-2 token-choice gating in GShard.
We train several variants of our architecture at the 100M scale (i.e. 100M expert size) by increasing the number of experts to understand the scaling effects of our method. We also train a 8B scale MoE model. The large MoE model is partitioned with 2D sharding algorithm as presented in GSPMD [36], which fully exploits the 2D topology of the TPU cluster (19]. Across different scales and setups, our method outperforms related work and demonstrates strong downstream task performance on selected tasks in GLUE and SuperGLUE.
Experiments
4.1 Setup
Table summarizes the hyperparameter settings of different MoE models. As reference point, we also include the respective dense model configurations with comparable numbers activated parameters per-token during inference. To study of the effect of scaling the number of experts,
5

---

3.0
G5_ DS_T CF2
3.1
EC EC_128C
TOP2 128
2.8
EC_ 64E
Eval Perplexity
Perplexity
TOP2 64E
EC_32E
2.6
2.9
TOP2, TOP2_32E 32E
EC. 16E
2.4
TOP2_16E
2.8
>2x
2.2
2.7
2.0
0
250
500
750 1000
1250 1500
1750
2.6
100
200
300
400
500
www.
600
10K Steps
10K Steps
(b)
Figure 2: (a) Training convergence is more than 2x faster using method compared to GShard top-2 gating. (b) Training perplexity scales strongly with the number of experts while keeping expert size fixed. EC consistently outperforms GShard top-2 gating.
studied varying the number of experts but fixing the per expert size 100M parameters. For example, O. 1B/44E represents the architecture of approximately 100M parameter dense model with every other layer replaced by a 64-expert MoE layer. The MoE model degenerates into dense transformer architecture when each MoE layer only has one expert. While 7parums is the total number of trainable parameters, Ract params represents the number of activated parameters per token. number of Transformer layers, M is the model dimension, H is the hidden dimension after the projection in each transformer layer, Ahoads the number of attention heads, and dheed is the hidden dimension of each attention head.
the total
Dataset: We use the high-quality dataset from GLaM 7 of 1.6 trillion tokens that are representative of a wide range of natural language use cases. An in-house classifier is trained to classify between collection of curated text and other webpages and estimate the content quality of webpage. A high-quality filtered subset of webpages are combined with books, Wikipedia pages, conversations, forums, and news to create the final dataset. The data and mixture weights can be found in Table in the GLaM paper.
Model Training: Our model training follows the setups of GLaM 7 where a maximum sequence length of 1024 tokens is adopted. We use an Adafactor optimizer
32 with first-moment decay constant for the first 10K
31 = 0 and second-momen decay 83 0.99. We keep the learning
training steps, and then decay it with an inverse square root schedule. Unlike most related works, we do impose any auxiliary loss for load balance, such as described in Switch Transformer (10) and GShard (2I]. We use the SentencePiece subword tokenizer with a vocabulary of size of 256K. The largest model (8B/64E) trained on 512 TPU V4 chips. We use a dropout rate of during training as the number of tokens in the training data corpus is much greater than the total number of tokens during training.
Model Evaluation: We mainly focus on evaluating the finetuning performance on the 11 selected
tasks from GLUE and SuperGLUE benchmarks 34 35
4.2 Training Efficiency
We first study training efficiency and convergence. We use expert choice with a capacity factor of (EC-CF2) to match the activated model size and computational cost per token basis in GShard top-2 gating and run both for fixed number of steps. The results are shown Fig. 2 (a). Comparing to GShard top-2 gating, which showed stronger performance in both perplexity in the evaluation dataset and fine-tuning on downstream tasks compared to Switch Transformer top-1 gating, EC-CF2 converges more than 2x faster during training. More specifically, EC-CF2 reaches the same perplexity as GShard top-2 less than half the steps, and with each GShard top-2 step being 20% slower than method. As explained in Section 3.1 the slower step time in top-2 gating is due to load imbalance

---

100M/128E
100M/64E
Name
Metric
Split
Top-1
Top-2
EC-CF2
Top-1
GS Top-2
EC-CF2
BoolQ
acc
dev
77.4
76.5
76.9
73.2
77.5
79.7
CB
acc
dev
87.5
80.9
89.1
85.9
84.4
89.1
CoLA
acc
dev
78.9
84.0
86.7
64.1
85.2
88.3
MNLI
acc
dev
82.3
83.6
84.9
80.8
86.7
MRPC
acc
82.6
81.0
83.1
81.3
84.4
QNLI
dev
89.5
88.6
89.0
89.4
89.7
QQP
dev
90.6
90.3
90.4
88.9
90.5
91.0
RTE
dev
77.0
78.9
78.5
74.1
79.3
81.6
SST2
dev
92.0
94.5
94.6
91.8
95.1
95.1
WIC
acc
dev
67.8
65.5
68.1
64.4
67.8
65.6
WNLI
aoc
dov
65.6
70.3
67.2
68.8
68.8
71.7
Avg
81.0
81.3
82.6
78.4
82.2
84.0
100M/32E
8B/64E
Name
Metric
Split
ST Top-1
GS Top-2
EC-CF2
Top-1
GS Top-2
EC-CF2
BoolQ
ace
dev
74.5
79.0
79.3
89.1
89.5
89.2
CB
acc
dev
80.6
81.3
92.2
93.8
96.7
100
acc
92.2
93.8
88.3
MNLI
dev
83.1
87.8
88.0
90.7
91.4
91.1
MRPC
dev
82.3
85.2
84.4
89.3
91.7
90.6
QNLI
dev
91.6
91.9
92.5
94.5
94.9
95.0
QQP
dev
90.1
91.5
92.0
92.1
92.5
93.8
RTE
acc
dev
75.0
79.1
78.1
91.0
92.2
95.2
SST2
acc
dev
93.3
94.4
95.4
97.1
98.0
97.7
WIC
acc
dev
62.5
65,9
69.8
74.5
76.4
83.8
WNLI
acc
dev
65.6
64.1
68.8
78.1
82.8
92.8
Avg
83.5
85.0
88.9
Table 2: Expert choice with capacity factor of (EC-CF2) outperforms Top-1 gating in Switch Transformer (ST) and top-2 gating in GShard (GS) on GLUE and SuperGLUE tasks. Note that with an expert size of 100M parameters, 100M/32E works best for our method and Ghard Top-2 while 100M/128E works better for Switch Transformer Top-1. Our method consistently outperforms the others across all the scales.
where some experts can receive a lot more tokens than the desired capacity. As result, the step latency will be bottlenecked by the most loaded expert.
4.3 Scaling the Number of Experts
As presented in T increasing the number of experts effectively increases model capacity without increasing activated model size. We scale the number experts while fixing the expert size to 100M parameters for both expert choice (EC) and GShard (Top-2) methods and find both methods work well in terms perplexity on the evaluation dataset during pre-training. As demonstrated Fig. (b), having more experts consistently improves training perplexity.
4.4 Fine-tuning on GL.E and SuperGLUE
To validate whether improved perplexity directly translates to better performance downstream tasks, we perform fine-tuning on 11 selected tasks from GLUE and SuperGI.UE, We compare three MoE methods including Switch Transformer top-1 gating (ST 1). GShard top-2 gating (GS Top-2) and our method (EC-CF2) that matches the activation memory size and computational cost of GS Top-2. Indicated by the results in Table our EC-CF2 method consistently outperforms the related methods and yields more than 2% average accuracy increase in large 8B/64E setting. Table further 3urther compares our 8B/64E model against its dense counterpart. Again, our method achieves stronger fine-tuning results, increasing the average score by 3.4 point.
Interestingly, we observe the 100M/32E model setting works the best for both GS Top-2 and EC-CF2, even though the effective model capacity is smaller than that of 100M/64E and 100M/128E. This result indicates that a good training perplexity does not always translate better performance of downstream tasks.

---

Model
BoolQ
CB
CoLA
MNLI
MRPC
QNLI
QQP
RTE
SST2
WIC
WNLI
Avg
Dense 8B
88.2
100
86.4
91.3
86.7
94.7
91.2
92.2
97.2
75.6
78.1
89.2
EC-CF2 8B/44E
89.2
100
89.1
91.1
90.6
95.0
93.8
95.2
97.7
83.8
92.8
92.6
Table 3: Comparison between Dense 8B and Expert Choice (EC-CF2) 8B/64E models: Our method
significantly outperforms the dense model in downstream tasks.
Table 4: (a) Limiting the number of experts token in expert choice method affects downstream accuracy. (b) Comparing Hash Layer.
Distribution of the number of experts routed to per token in a 100M/64E model.
Figure 3:
0.5
Fraction of Tokens
0.4
Method
Max # of Experts
Avg acc.
0.3
EC-CAP2
83.2
0.2
EC-CAP3
84.0 + 0.4
0.1
BC-CF2
84.0 + 0.2
0
Hash Layer
81.3 #.1
0
2
4
>4
Number of Experts
4.5 Heterogeneity Matters
We regularized expert choice by limiting the maximum number of experts for each token, using the method described in Section 3.3 Table (ooeorrsspotss the average accuracy the 11 selected datasets. EC-CAP2 the variant of our expert choice method by limiting the number of experts of each token to 2. This decreases the fine-tuning accuracy by 0.8 points average. In addition, EC-CAP3 allows a maximum of 3 experts per token and achieves on par results compared to the vanilla expert choice method. This ablation study confirms that allowing variable number of experts per token is indeed helpful.
Capped Expert Choice:
Variable Experts per Token: We compute statistics on token-to-expert routing, particularly on the ratio tokens that have been routed certain number of experts. According to Fig. 3 a majority of tokens have been routed to one or two experts while 23% have been routed to three or four experts and only about 3% tokens have been routed to more than experts. This plot verifies our hypothesis that our method learns to allocate variable number experts to tokens, which can be beneficial for important tokens.
4.6 Comparison with Hash Layer
In this section, we compare our method with Hash Layers (28). We use mod x to map token ID an expert ID. This ensures load balance and generates specialized experts. The fine-tuning results are presented in the last row in Table 4 Hashing based routing performs worse than expert choice in terms of average scores and variance. This indicates that load balancing alone does not generate all the benefits
4.7 Ablation
We study the capacity factor in our expert choice method and compare the training perplexity with the bascline top-1 gating method used in Switch Transformer. As described in Eq. ([) the capacity factor determines how many experts on average cach token be routed to, thus the bucket size of each expert. In all our previous experiments, we use a capacity factor of 2, which matches the computational footprint of the top-2 gating used in GShard method. To match the computation cost on per-token basis fairly with top-1 gating used in Switch Transformer, we reduce the capacity factor to and plot the training perplexity in Fig. 4 (a). Not surprisingly, using smaller capacity factor yields higher perplexity, but our method still significantly outperforms top-1 gating. We further push the capacity factor down 0.5, and observe that it still outperforms the top-1 gating.
Capacity Factor:
Comparison with Dense Models on Pre-training: We compare our method with dense models pre-training. As shown in Fig. (b), our method consistently outperforms the dense method in

---

3.6
EC_CF2
3.4
10064E
EC_CR
DENSE 100M
3.4
EC CF0.5
3.2
8B64E
Eval Perplexity
TOP1
DENSE 88
3.0
3.2
Lval Ra
2.8
3.0
2.6
2.4
2.8
2.2
2.6
2.0
160
200
300 400 500
600 700 10K Steps
0 100
200
300
400
500
600
700
10K Steps
(a)
(b)
Figure 4: (a) Varying the capacity factor in our expert choice method: Decreasing the capacity factor from two to one degrades the perplexity but still outperforms the top-1 gating. (b) Training perplexity comparison with dense models.
perplexity and convergence time. For small expert size of 100M parameters, the benefit of sparse gating is even more significant. Orthogonal to results presented in Fig. (b), where scaling the number of experts improves model performance, Fig. (b) shows that increasing expert capacity also significantly increases model performance.
Conclusion
We propose new routing method for sparsely activated mixture-of-experts (MoE) models. This method addresses load imbalance and under-utilization of experts conventional MoE methods, and enables selecting different numbers of experts for each token. Our model demonstrates more than 2x training efficiency improvements when compared to the state-of-the-art GShard and Switch Transformer models, and also achieves strong gains when finetuning 11 datasets in the GLUE and SuperGLUE benchmark.
Limitations
The expert choice method might not immediately apply to auto- uuto-reerressvve text generation our current implementation takes the past and future tokens to perform the top-k selection. One possible solution is to collect a large batch of imput sequences, dispatch tokens of the same sequence into separate groups, and perform expert choice routing for each group. Another scenario where the expert choice method does not immediately apply is when the batch size becomes very small during serving or inference. global top-k can selected instead and we can cap the number of times each expert token gets selected. We leave these possible improvements for future work.
Another long-standing issue with MoE has been the large memory footprint. Even though computa- tional cost can be reduced using sparsely gated networks, the total number of parameters increases linearly or sub-linearly with the number of experts. Increasing the number of experts requires reser- vation of a large number of hardware devices. Therefore, dynamic (used) power is saved while static (reserved) power not. Power saving techniques such the ability put hardware devices into low power states while not in use [17 can help with reducing the reserved power requirements.
References
[1] Davide Abati, Jakub Tomezak, Tijmen Blankevoort, Simone Calderara, Rita Cucchiara, and Babak Ehteshami Bejnordi. Conditional channel gated networks for task-aware continual learning. In Proceedings of the IEEF/VV Conference on Computer Vision and Pattern Recognition, pages 3931-3940, 2020.

---

[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:1887-1101,2000
[3] Kyunghyun Cho and Yoshua Bengio. Exponentially increasing the capacity-to-compuratie ratio for conditional computation in deep learning. arXiv preprint arXiv: 1406. 7362, 2014.
[4] Zihang Dai, Hanxiao Liu, Quoe V. Le, and Mingxing Tan. CoAtNet: Marrying convolution and attention for all data sizes. In Advances in Neural Information Processing Systems, 2021.
[5] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.0286/ 2019.
[6] Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, ICML' 17, page 933-941. JMLR.org, 2017.
[7] Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with tixture-of-experre In International Conference on Machine Learning, pages 5547-5569. PMLR, 2022.
[8] Dheeru Dua, Shruti Bhosale, Vedanuj Goswami, James Cross, Mike Lewis, and Angela Fan. Tricks for training sparse translation models. arXiv preprint arXiv:21 2021.
[9] Richard L Dykstra. An iterative procedure for obtaining i-projections onto the intersection of convex sets. The annals of Probability, pages 975-984 1985.
[10] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.
[11] Golnaz Ghiasi, Tsung- Yi Lin, and Quoc Le. Nas-fpn: Learning scalable feature pyramid architecture for object detection. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 7036-7045, 2019.
[12] Sam Gross, Mare' Aurelio Ranzato, and Arthur Szlam. Hard mixtures of experts for large scale weakly supervised vision. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 6865-6873, 2017
[13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[14] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 330-645. Springer, 2016.
[15] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 1rXiv:1606.08415, 2016.
[16] Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Patwary, Mostofa Ali, Yang Yang. and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint arXiv:1712.004000 2017.
[17] Ping Huang, Zuocheng Xing, Tianran Wang, Qiang Wei, Hongyan Wang, and Guitao Fu. A brief survey on power gating design. In 2010 10th IEEE International Conference on Solid-State and Integrated Circuit Technology, pages 788-790 2010.
[18] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances neural information processing systems, 32, 2019.
10

---

[19] Norman P. Jouppi, Doe Hyun Yoon, George Kurian, Sheng Li, Nishant Patil, James Laudon, Cliff Young, and David A. Patterson. A Somain-speciffe supereomputer for training deep neural networks. Commun. ACM, 63(7):67-78, 2020.
[20] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, arXiv preprint 2rXiv:2001.08361, 2020.
[21] Dmitry Lepikhin, HooukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with condi- tional computation and automatic sharding. In International Conference on Learning Represen- tations, 2021.
[22] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large. sparse models. In International Conference on Machine Learning, pages 0265-6274. PMLR, 2021.
[23] Min Lin, Jie Fu, and Yoshua Bengio. Conditional computation for continual learning. arXiv preprint arXiv: 96635, 2019.
[24] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM Symposiun on Operating Systems Principles, pages 1-15, 2019.
[25] Joan Puigcerver, Carlos Riquelme, Basil Mustafa, Cedric Renggli, André Susano Pinto, Sylvain Gelly, Daniel Keysers, and Neil Houlsby. Scalable transfer learning with expert models. arXiv preprint arXiv: 13239, 2020.
[26] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training, 2018.
[27] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res.,
2020.
[28] Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566, 2021.
[29] Roy Schwartz, Jesse Dodge, Noah A Smith, and Oren Etzioni. Green ai. Communications of the ACM, 63(12):54-63, 2020.
[30] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanan- takool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al. Mesh-tensooffoow: Deep learning for supercomputers. Advances in neural information processing systems, 31, 2018.
[31] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- -gtted mixture-of-experts layer. arXiv preprint
2017.
[32] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning. pages 4596-4604. PMLR, 2018
[33] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprini arXiv: 1wX.:::909000053.0 2019.
[34] Alex Wang. Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
[35] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 2018

---

[36] Yuanzhong Xu, Hyoukloong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rabul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, al. Gspmd: general and scalable parallelization for ml computation graphs. arXiv preprint arXiv:2 2021.
[37] 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 arXiv:2110.04266 2021.
7 Checklist
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contribu- tions and scope? Yes
(b) Have you read the ethics review guidelines and ensured that your paper conforms them? Yes
(c) Did you discuss any potential negative societal impacts of your work? N/A. Not any.
(d) Did you describe the limitations of your work? Yes
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results? Yes. We include details in the experiment setup to help reproduce the main results.
(b) Did you specify all the training details? Yes
(c) Did you report error bars? Yes
(d) Did you include the amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? Yes
(a) If your work uses existing assets, did you cite the creators? Yes
(b) Did you mention license assets? No. The used dataset is not released yet.
(c) Did you include any new assets either in the supplemental material or as URL? No. The dataset not released yet.
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/cu- rating? No. Not using persons' data.
(c) Did you discuss whether the data you are using/curating contains personally identifiable in- formation or offensive content? Yes. The dataset does not contain any personally identifiable information or offensive content.
12


Reviews:
[{"review_id": "u1Y8wOtF_cY", "reviewer": "Reviewer_1p81", "summary": "This work proposes a new routing method called Expert Choice Routing for the Mixture-of-Experts (MoE) architectures. In this method instead of routing each token to a specific number of (top-$k$) experts, each expert receives a specific number of (top-$k$) tokens. Naturally using this approach, as mentioned in the manuscript, the following properties will be held:\n- Perfect load balancing will be achieved between all of the experts, since every expert receives the same number of tokens\n- Different tokens can be routed to various number of experts\n\nSimilar to GShard [1] and Switch Transformer [2], the authors apply the MoE and gating logic in the dense feed forward layer (FFN). The authors also use these works as baselines to compare their work against.\nThis comparison shows that using the same amount of computation cost the Expert Choice Routing method achieves superior results on GLUE and SuperGLUE benchmarks. Additionally, it is shown that using the expert choice routing method, the training efficiency is almost doubled. Lastly, it is shown that when using expert choice routing, training perplexity consistently decreases as the number of experts are increased.\n\nReferences\n\n[1] 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 International Conference on Learning Representations, 2021.\n\n[2] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.", "questions": "- How different batch sizes affect the performance of expert choice (EC) routing? Are there regimes of batch sizes where using EC routing may not provide any advantages over token choice (TC) routing?\n- What is the criteria for bolding numbers in Table 2? I assumed that highest number in each row is bolded. However, after inspecting the numbers, I found that this is not the case. For example, in the 100M/128 scenario on QQP, the highest number belongs to ST-Top-1 but the number belonging to EC-CF2 is highlighted.\n- In Figure 2, it is clearly shown that EC routing improves the training efficiency compared to the GShard (GS). Was it also observed that eval perplexity at the end of the training is lower for EC? Does EC converge to a smaller eval perplexity compared to the GS at the end of the training?", "limitations": "The authors adequately addressed the limitations and potential negative impacts of this work in the manuscript.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The manuscript is well written. The method and background are explained clearly and are easy to follow.\n- The paper introduces a novel routing technique which improves the training efficiency and accuracy of Mixture-of-Expert (MoE) architectures.\n- The authors' central claims are sound, convincing, and well supported.\n\nWeaknesses:\n- There are few writing mistakes which should be fixed:\n    - In Table 1, the 0.1B/128E model has 128 experts (not 64).\n    - Incorrect citations (e.g. In [18] the paper name is missing. However, from pages 103 to 112 it can be inferred that the paper is GPipe)\n    - In section 4.7, the term \"We further pushes\" should be changed to \"We further push\"\n- In the last sentence of section 4.2, it is claimed that each GShard top-2 step is 20% slower than expert choice routing. However, this is not supported and explained anywhere within the manuscript. Additional clarification is necessary regarding this claim.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Xiv47yViofk", "reviewer": "Reviewer_49ur", "summary": "The paper proposes a new token routing strategy for the MoE model architecture. The existing methods mostly focus on how each token chooses experts to be processed by, but this paper designs the opposite way which each expert chooses which tokens to process. By doing this, the authors claim that MoE models' intrinsic load balancing issue is resolved naturally. Furthermore, by applying more computation to more important tokens, the paper presents faster convergence and higher performance on a GLUE and SuperGLUE tasks.", "questions": "1. In equation (2), can you clarify which dimension is used to do softmax?\n2. Can you clarify why subsets of tasks from GLUE and SuperGLUE are selected? Are they all sequence classification tasks? If that's the case, I wondered how expert choice works for token classification tasks (ReCORD). Some of less importance tokens might not be fully processed and could potentially cause inaccurate perditions for them.\n3. When experts choose tokens, do they make choices across different sentences in a batch? If so, how would it affect the training dynamics? Can a model prefer some specific sentences? If not, could you explain how a batch is handled?", "limitations": "N/A", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "MoE models are getting more widely used due to their efficient scaling property with higher quality than the vanilla dense models. However, there has been a load balancing issue from the beginning, and it is known to be critical to the model performance. Still, it is hard to solve and prevents models to be effectively trained such as under-trained experts. Given this situation, the paper proposes a very unique and noble approach to solve the problem naturally. This is very timely and makes sense intuitively. The paper is well-organized, clearly written and easy to follow.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "J3yXBRk9UAX", "reviewer": "Reviewer_oGHv", "summary": "The paper presents expert choice routing for mixture-of-experts models instead of the more common token-based routing strategies. The proposed method ensures ideal load balancing, resulting in faster training convergence as well as increased performance on downstream tasks, partly attributed through solving the under-spezialization problem in previous MOE approaches.", "questions": "- If I choose an appropriate hashing function for my token-based routing (cf. Balanced assignment [28]), shouldn't that ensure a balanced expert load and as such contradicts your statement that previous work \"leads to an imbalanced load of experts\" ?  ", "limitations": "Limitations are appropriately addressed. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths:\n\n- Simple but effective method.\n- Expert choice routing demonstrates faster convergence and improved downstream performance.\n- Insightful comparison towards Hashing Layers, confirming that just balancing the load has some performance drawbacks.\n- Results are well presented and and claims are well supported by experimental results.\n\n### Weaknesses:\n\n- **(small)**: I think the paper would still benefit from a Figure similar to Figure 4 in [28] to drive home the point of load balancing.\n- **(small)**: Adding a rough wall clock time for the training time would be helpful.\n\n### Minor points:\n\n- Next time make sure to omit $\\verb+final+$ and $\\verb+preprint+$ in the NeurIPS $\\LaTeX$ package which adds line numbers and hence makes reviewing easier.\n- Try out the `sort&compress` option to natbib which sorts the references and makes things a little cleaner presentation-wise.\n- Clean up references to be consistent in terms of the included information, i.e conference citations should have the same style and make sure to handle preprints consistently and add their arXiv identifier. E.g. [2], [4], [18], [30] all have different formatting. [2] doesn't even contain the title.\n- Table 1 Caption \"Session 4.1\" -> \"Section 4.1\". Generally, I found that you are inconsistent with adding cross reference links to quickly jump to Sections, e.g. in Section 4.5, the reference to Section 3.3 doesn't have one. Switch to `cleveref` and use $\\verb|\\Cref|$ for consistency and to avoid typos.\n- Just below Table 2 there seems to be a lonely (b) with no corresponding (a) that most likely shouldn't be there.\n- Table 2: There are a few errors in the printing of bold numbers and formatting. I always assume that bold number highlight the best performance, even if it's not the presented method. Hence, I'd suggest adding that as well as fixing the errors where one of the other methods actually performs better.  **100M/128E**: BoolQ, QNLI, QQP, WNLI -> ST or GS outperform reducing the outperformed tasks from 10/11 -> 6/11. Here also GS Top-2 should be printed bold for RTE. **100M/64E**: GS Top-2 should be bold for WiC. **100M/32E**: Can we get the number for ST Top-1 on CB for completeness? GS Top-2 should be bold for MRPC, RTE. **8B/64E**: GS Top-2 should be bold for BoolQ, MNLI, MRPC, SST2\n- Appendix Table 2 EC-CF0.5 last column has additional space that shouldn't be there.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "9vWWDFBJGF", "reviewer": "Reviewer_HqBR", "summary": "The paper works on sparse MoE for large-scale Transformer.\n\nKey idea: Rather than letting tokens select the top-k experts, the proposed method selects the top-k tokens for each expert.\n\n", "questions": "- How about the finetuning numbers of dense models in Table 2?\n\n- \"model training follows the setups of GLaM\": is the model pretrained using the GPT task? \n\n- The results are compared with the T5 dense model, but the models use different training data. The numbers are not fairly comparable. How about the results using the same data?\n\n- Are there numbers of training the model on public text data, such as CommonCrawl?\n\n- How about the results for vision Transformers or machine translation instead of just language modeling?", "limitations": "There are two limitations mentioned in Sec 6.\n\n- The proposed method can not immediately apply to auto-regressive text generation.\n\n- Large memory footprint.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The idea of selecting tokens for experts to experts is novel for MoE model training.\n\n2. Better convergence speed compared with previous methods.\n\n3. The method naturally solves the load balance issue of MoE.\n\n\n********************************************\nWeaknesses:\n\n1. It is non-trivial to apply the proposed method to the autoregressive generation of GPT-like models, because all the tokens are not pre-given at the beginning. However, the most successful MoE-based models are GPT applications, which restricts the usage of the proposed method. This limitation is also mentioned in Sec 6.\n\n2. The fine-tuning performance of dense models is not reported in Table 2. Only MoE-based models are evaluated under the transfer setting. It is straightforward to compare the finetuning numbers for dense models.\n\n3. The results are compared with the T5 dense model, but the models use different training data. The numbers are not fairly comparable.\n\n4. The training data are not publicly available. It makes future comparisons difficult for other papers.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept