Original: Mixture-of-Experts with Expert Choice Routing
Revised: Mixture-of-Experts with Expert Choice Routing

Revised: s)

Original: Abstract
Revised: Abstract

Original: 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.
Revised: Sparsely-activated Mixture-of-experts (MoE) models allow the number of parameters to greatly increase while keeping the amount of computation for a giventoken or a given sample unchanged.

Original: However, a poor expert routing strategy can cause certain experts to be under-trained, leading to an expert being under or over-specialized.
Revised: However, a poor expert routing strategy cancause certain experts to be under-trained, leading to an expert being under orover-specialized.

Original: Prior work allocates a ﬁxed number of experts to each token using a top- k function regardless of the relative importance of different tokens.
Revised: Prior work allocates a ﬁxed number of experts to each tokenusing a top- k function regardless of the relative importance of different tokens.

Original: To address this, we propose a heterogeneous mixture-of-experts employing an expert choice method.
Revised: Toaddress this, we propose a heterogeneous mixture-of-experts employing an expertchoice method.

Original: Instead of letting tokens select the top- k experts, we have experts selecting the top- k tokens.
Revised: Instead of letting tokens select the top- k experts, we have expertsselecting the top- k tokens.

Original: As a result, each token can be routed to a variable number of experts and each expert can have a ﬁxed bucket size.
Revised: As a result, each token can be routed to a variablenumber of experts and each expert can have a ﬁxed bucket size.

Original: We systematically study pre-training speedups using the same computational resources of the Switch
Revised: We systematicallystudy pre-training speedups using the same computational resources of the Switch

Original: Transformer top-1 and GShard top-2 gating of prior work and ﬁnd that our method improves training convergence time by more than 2 × .
Revised: Transformer top-1 and GShard top-2 gating of prior work and ﬁnd that our methodimproves training convergence time by more than 2 × .

Original: For the same computational cost, our method demonstrates higher performance in ﬁne-tuning 11 selected tasks in the GLUE and SuperGLUE benchmarks.
Revised: For the same computationalcost, our method demonstrates higher performance in ﬁne-tuning 11 selected tasksin the GLUE and SuperGLUE benchmarks.

Original: For a smaller activation cost, our method outperforms the T5 dense model in 7 out of the 11 tasks.
Revised: For a smaller activation cost, ourmethod outperforms the T5 dense model in 7 out of the 11 tasks.

Original:  Introduction
Revised:  Introduction

Original: Scaling up model capacity, dataset size, and training time has demonstrated huge success in enhancing the performance of computer vision architectures [13, 14, 11, 4] as well as neural language models [26, 2, 20, 27].
Revised: Scaling up model capacity, dataset size, and training time has demonstrated huge success in enhancingthe performance of computer vision architectures [4, 11, 13, 14] as well as neural language models [2,20, 26, 27].

Original: The ﬁnal model quality has been found to have a power-law relationship with the amount of data, model size and compute time [16, 20].
Revised: The ﬁnal model quality has been found to have a power-law relationship with the amountof data, model size and compute time [16, 20].

Original: However, training efﬁciency, which is deﬁned as the total amount of computation used to achieve superior model quality than the state of the art system [21], should receive greater attention as we increase our efforts towards green AI [29].
Revised: However, training efﬁciency, which is deﬁned asthe total amount of computation used to achieve superior model quality than the state of the artsystem [21], should receive greater attention as we increase our efforts towards green AI [29].

Original: Sparsely gated mixture-of-experts [31] (MoE) provides an effective way to scale model capacity given a ﬁxed computational cost, and has recently played an important role in increasing the training efﬁciency of large-scale language models [21, 10].
Revised: Sparsely gated mixture-of-experts [31] (MoE) provides an effective way to scale model capacitygiven a ﬁxed computational cost, and has recently played an important role in increasing the trainingefﬁciency of large-scale language models [10, 21].

Original: MoE operate by adopting a number of experts, each as a sub-network, and by activating only one or a few experts for each input token.
Revised: MoE operate by adopting a number of experts,each as a sub-network, and by activating only one or a few experts for each input token.

Original: A gating network must be chosen and optimized in order to route each token to the most suited expert(s).
Revised: A gatingnetwork must be chosen and optimized in order to route each token to the most suited expert(s).

Original: For example, recent work has implemented sparse routing via k -means clustering [12], linear assignment to maximize token-expert afﬁnities [22], or hashing [28, 8].
Revised: Forexample, recent work has implemented sparse routing via k -means clustering [12], linear assignmentto maximize token-expert afﬁnities [22], or hashing [8, 28].

Original: Many of the prior work use a routing strategy concerning the token choice , where each token selects the best one or two experts.
Revised: Many of the prior work use a routingstrategy concerning the token choice , where each token selects the best one or two experts.

Original: We argue that the independent token choice of prior work often leads to an imbalanced load of experts, which causes training inefﬁciency and sub-optimal training of the model.
Revised: We argue that the independent token choice of prior work often leads to an imbalanced load of experts,which causes training inefﬁciency and sub-optimal training of the model.

Original: In order to mitigate this issue, previous sparsely gated networks introduce additional auxiliary losses as regularization to prevent too many tokens being routed to a single expert, but the effectiveness is still limited.
Revised: In order to mitigate thisissue, previous sparsely gated networks introduce additional auxiliary losses as regularization toprevent too many tokens being routed to a single expert, but the effectiveness is still limited.

Original: Recent approaches [22, 28, 8] explore alternative strategies for routing, but they focus on pre-training only and do not demonstrate performance gain on downstream tasks.
Revised: Recentapproaches [8, 22, 28] explore alternative strategies for routing, but they focus on pre-training only and do not demonstrate performance gain on downstream tasks.

Original: Moreover, none of the previousmethods consider allocating a variable number of experts to each token based on importance, which can be beneﬁcial.
Revised: Moreover, none of the previous

Original: [Figure]
Revised: [Figure]

Original: We
Revised: We

Revised: Like p = 0.

Revised: [Figure]

Revised: Router

Revised: Router

Revised: Token 1 Token

Original:  Like
Revised: We Like

Original: To
Revised: To

Original: [Figure] In
Revised: Play Soccer In

Original: The
Revised: The

Original: Field
Revised: Field

Revised: We Like To Play

Original: {+Token 
+}
Revised: {+Token 
+}

Original: {+Token 
+}
Revised: {+Token 
+}

Original: {+Token 5 Token 
+}
Revised: {+Token 5 Token 
+}

Revised: Token 7 Token

Original: [Figure]
Revised: We Like Soccer Field FFN FFN Top-k

Revised: Top-k methods consider allocating a variable number of experts to each token based on importance, whichcan be beneﬁcial.

Original: We propose a very simple yet effective routing method we are calling expert choice .
Revised: We propose a very simple yet effective routing method we are calling expert choice .

Original: Unlike conventional MoE where tokens select one or two top-scoring experts, our method lets each expert pick the top- k tokens.
Revised: Unlike conventional MoE where tokens select one or two top-scoring experts, our method lets each expert pick thetop- k tokens.

Original: Our method guarantees perfect load balancing, allows a variable number of experts for each token, and achieves substantial gains in training efﬁciency and downstream performance as demonstrated in our experiments.
Revised: Our method guarantees perfect load balancing, allows a variable number of expertsfor each token, and achieves substantial gains in training efﬁciency and downstream performance asdemonstrated in our experiments.

Original: Our major contributions include:
Revised: Our major contributions include:

Original: • We identify common pitfalls in conventional MoE such as load imbalance as described in Section 3.1.
Revised: • We identify common pitfalls in conventional MoE such as load imbalance as describedin Section 3.1.

Original: We then propose a heterogeneous, expert choice method to provide a ﬂuid allocation of model parameters based on a learnt token-to-expert importance.
Revised: We then propose a heterogeneous, expert choice method to provide a ﬂuidallocation of model parameters based on a learnt token-to-expert importance.

Original: This method intrinsically guarantees load balance without imposing an auxiliary loss.
Revised: This method intrinsically guarantees load balance without imposing an auxiliary loss.

Original: • We show our method provides over 2 × faster training convergence in a 8B/64E (8 billion activated parameters, 64 experts) model, compared to the top-1 and top-2 gating counterparts in Switch Transformer [10] and GShard [21].
Revised: • We show our method provides over 2 × faster training convergence in a 8B/64E (8 billion activated parameters, 64 experts) model, compared to the top-1 and top-2 gating counterpartsin Switch Transformer [10] and GShard [21].

Original: • We show our method demonstrates strong scaling when increasing the number of experts from 16 to 128, evaluated in training perplexity.
Revised: • We show our method demonstrates strong scaling when increasing the number of experts from 16 to 128, evaluated in training perplexity.

Original: • We show our method demonstrates strong performance on downstream tasks selected from GLUE and SuperGLUE at all the evaluated scales.
Revised: • We show our method demonstrates strong performance on downstream tasks selected from GLUE and SuperGLUE at all the evaluated scales.

Original: More speciﬁcally, our 8B/64E model outperforms a T5 11B dense model in 7 out of 11 tasks evaluated.
Revised: More speciﬁcally, our 8B/64E model outperforms a T5 11B dense model in 7 out of 11 tasks evaluated.

Original:  Related Work
Revised:  Related Work

Original: Scaling: Various approaches have been proposed to scale up neural network capacity to improve performance.
Revised: Scaling: Various approaches have been proposed to scale up neural network capacity to improveperformance.

Original: Recent works have successfully scaled models to billions of parameters via various forms of model parallelism [21, 33, 27, 26, 2].
Revised: Recent works have successfully scaled models to billions of parameters via variousforms of model parallelism [2, 21, 26, 27, 33].

Original: 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.
Revised: Model parallelism [30] splits weights and tensorsacross multiple cores while pipeline parallelism [18, 24] splits different layers across devices withmicro-batches pipelined to the different layers.

Original: To enable continued scaling of neural networks, improving model training and serving efﬁciency has become a critical research area.
Revised: To enable continued scaling of neural networks,improving model training and serving efﬁciency has become a critical research area.

Original: Conditional Computation: Computation decisions can be made dynamically based on the input [25, 23].
Revised: Conditional Computation: Computation decisions can be made dynamically based on the input [23,25].

Original: Conditional computation has been proposed as a way to increase the capacity of a deep neural network without increasing the amount of computation, by activating certain parameters and computation on demand, on a per-example or per-token basis [3].
Revised: Conditional computation has been proposed as a way to increase the capacity of a deepneural network without increasing the amount of computation, by activating certain parameters andcomputation on demand, on a per-example or per-token basis [3].

Original: Conditional convolution layers [1] with task-speciﬁc gating has been used to combat catastrophic forgetting when a sequence of learning problems are optimized.
Revised: Conditional convolution layers [1]with task-speciﬁc gating has been used to combat catastrophic forgetting when a sequence of learningproblems are optimized.

Original: The gating decisions may be binary or sparse and continuous, stochastic or deterministic.
Revised: The gating decisions may be binary or sparse and continuous, stochastic ordeterministic.

Original: Mixture of Experts: Sparsely-gated MoE [31] is the ﬁrst model to demonstrate massive improvements in model capacity, training time, or model quality with gating.
Revised: Mixture of Experts: Sparsely-gated MoE [31] is the ﬁrst model to demonstrate massive improvements in model capacity, training time, or model quality with gating.

Original: Switch Transformer [10] simpliﬁes the gating by selecting only the top expert per token using a softmax over the hidden state and demonstrates better scaling than previous work.
Revised: Switch Transformer [10]simpliﬁes the gating by selecting only the top expert per token using a softmax over the hidden stateand demonstrates better scaling than previous work.

Original: All the prior work requires an auxiliary loss to explicitly encourage balancing.
Revised: All the prior work requires an auxiliary loss toexplicitly encourage balancing.

Original: This loss term has to be carefully weighted to not overwhelm the primary loss.
Revised: This loss term has to be carefully weighted to not overwhelm the primary loss.

Original: However, auxiliary loss does not guarantee balancing and a hard capacity factor has to be imposed.
Revised: However, auxiliary loss does not guarantee balancing and a hard capacity factor has tobe imposed.

Original: As a result, many tokens can still be unprocessed by the MoE layer.
Revised: As a result, many tokens can still be unprocessed by the MoE layer.

Original: Hard MoE [12] with a single decoding layer can be efﬁciently trained to good effect on large scale hashtag prediction tasks.
Revised: Hard MoE [12] witha single decoding layer can be efﬁciently trained to good effect on large scale hashtag prediction tasks.

Original: Base Layers [22] formulate a linear assignment that maximizes token-expert afﬁnities while ensuring each expert receives an equal number of tokens.
Revised: Base Layers [22] formulate a linear assignment that maximizes token-expert afﬁnities while ensuringeach expert receives an equal number of tokens.

Original: Hash layers [28, 8] devise hashing techniques on input tokens.
Revised: Hash layers [8, 28] devise hashing techniques oninput tokens.

Original: However, the evaluations are limited to pre-training perplexity.
Revised: However, the evaluations are limited to pre-training perplexity.

Original: THOR [37] randomly activates experts during training and inference and is trained with a consistency regularization loss.
Revised: THOR [37] randomlyactivates experts during training and inference and is trained with a consistency regularization loss.

Original: THOR has demonstrated strong performance on translation tasks.
Revised: THOR has demonstrated strong performance on translation tasks.

Original: Different from these prior works, our method is a learnt method that enables heterogeneous MoE and effectively improves downstream ﬁne-tuning performance.
Revised: Different from these prior works,our method is a learnt method that enables heterogeneous MoE and effectively improves downstreamﬁne-tuning performance.

Original:  Method
Revised:  Method

Original: We ﬁrst 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.
Revised: We ﬁrst identify a few pitfalls in the routing method of conventional mixture-of-experts (MoE) modelsand then present our method using expert choice to tackle these problems.

Original: 3.1 Pitfalls of Token-Choice Routing
Revised: 3.1 Pitfalls of Token-Choice Routing

Original: MoE can be computationally advantageous compared to a dense model, a routing strategy must be used to assign each token to the most-suited experts.
Revised: MoE can be computationally advantageous compared to a dense model, a routing strategy must beused to assign each token to the most-suited experts.

Original: Conventional MoE models employ token-choice routing which independently selects the top- k experts for each token [31, 21, 10].
Revised: Conventional MoE models employ token-choicerouting which independently selects the top- k experts for each token [10, 21, 31].

Original: We argue that this strategy has a few pitfalls that lead to sub-optimal training.
Revised: We argue that thisstrategy has a few pitfalls that lead to sub-optimal training.

Original: Load Imbalance: Token-choice routing often lead to poor load balancing across experts.
Revised: Load Imbalance: Token-choice routing often lead to poor load balancing across experts.

Original: That is, some experts may be trained with most tokens, leaving the remaining experts under-utilized.
Revised: That is,some experts may be trained with most tokens, leaving the remaining experts under-utilized.

Original: Experts can be under specialized because a lot of model capacity in the under-utilized experts are wasted.
Revised: Expertscan be under specialized because a lot of model capacity in the under-utilized experts are wasted.

Original: On the other side, some tokens will not be processed, since over-utilized experts can only take a maximum number of tokens at each step in order to avoid running out of memory.
Revised: On the other side, some tokens will not be processed, since over-utilized experts can only take amaximum number of tokens at each step in order to avoid running out of memory.

Original: Load imbalance can also hurt step latency, thus inference time, as the step latency can be determined by the most loaded expert.
Revised: Load imbalance canalso hurt step latency, thus inference time, as the step latency can be determined by the most loadedexpert.

Original: Previous methods add an auxiliary loss on load balancing to mitigate the issue.
Revised: Previous methods add an auxiliary loss on load balancing to mitigate the issue.

Original: However, this auxiliary loss does not guarantee a balanced load, especially during the important early stages of training.
Revised: However, thisauxiliary loss does not guarantee a balanced load, especially during the important early stages oftraining.

Original: Indeed, we empirically observe that the over-capacity ratio can reach 20%–40% for some experts in token choice routing , indicating that a signiﬁcant portion of the tokens routed to these experts will be dropped.
Revised: Indeed, we empirically observe that the over-capacity ratio can reach 20%–40% forsome experts in token choice routing , indicating that a signiﬁcant portion of the tokens routed tothese experts will be dropped.

Original: Under Specialization: Each MoE layer uses a gating network to learn token-to-expert afﬁnity.
Revised: Under Specialization: Each MoE layer uses a gating network to learn token-to-expert afﬁnity.

Original: Ideally, the learnt gating network should produce the afﬁnity such that similar or relevant tokens are routed to the same expert.
Revised: Ideally, the learnt gating network should produce the afﬁnity such that similar or relevant tokens arerouted to the same expert.

Original: A sub-optimal strategy can produce redundant experts and/or experts that are not sufﬁciently specialized.
Revised: A sub-optimal strategy can produce redundant experts and/or experts thatare not sufﬁciently specialized.

Original: Under specialization may result by imposing an large auxiliary loss which favors more load balanced but less effective routing.
Revised: Under specialization may result by imposing an large auxiliary losswhich favors more load balanced but less effective routing.

Original: Finding the right balance on the auxiliary loss to promote both load balancing and specialization is challenging for token-choice routing.
Revised: Finding the right balance on the auxiliaryloss to promote both load balancing and specialization is challenging for token-choice routing.

Original: 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.
Revised: Same Compute for Every Token: Finally, in a token-choice strategy each token receives exactlyk experts and therefore occupies the same amount of compute.

Original: We hypothesize that this is not necessary nor desired.
Revised: We hypothesize that this is notnecessary nor desired.

Original: Instead, a MoE model should ﬂexibly allocate its compute resource based on the complexity of the input.
Revised: Instead, a MoE model should ﬂexibly allocate its compute resource basedon the complexity of the input.

Original: Motivated by the aforementioned observations, we next describe a simple yet effective method which produces load balanced assignments based on expert choice .
Revised: Motivated by the aforementioned observations, we next describe asimple yet effective method which produces load balanced assignments based on expert choice .

Original: 3.2 Heterogeneous MoE via Expert Choice
Revised: 3.2 Heterogeneous MoE via Expert Choice

Original: Different from conventional routing, an expert choice method independently selects top-
Revised: Different from conventional routing, an expert choice method independently selects top-

Original: k tokens for each expert, where k is a ﬁxed expert capacity (i.e. the number of tokens each expert can take).
Revised: k tokensfor each expert, where k is a ﬁxed expert capacity (i.e. the number of tokens each expert can take).

Original: Despite its simplicity, expert choice achieves perfect load balancing by design.
Revised: Despite its simplicity, expert choice achieves perfect load balancing by design.

Original: It also enables more ﬂexible allocation of model compute since tokens can be received by a variable number of experts.
Revised: It also enables moreﬂexible allocation of model compute since tokens can be received by a variable number of experts.

Original: In our experiments, we set k as
Revised: In our experiments, we set k as

Original: [Equation]
Revised: [Equation]

Original: where n is the total number of tokens in the input batch (such as batch size × sequence length), c is the capacity factor, and e is the number of experts.
Revised: where n is the total number of tokens in the input batch (such as batch size × sequence length), c isthe capacity factor, and e is the number of experts.

Original: The capacity factor c denotes on average how many experts are utilized by a token.
Revised: The capacity factor c denotes on average how many experts are utilized by a token.

Original: Given input token representations X ∈ R n × d where d is the model hidden dimension, our method produces a token-to-expert assignment denoted by three output matrices I , G and P .
Revised: Given input token representations X ∈ R n × d where d is themodel hidden dimension, our method produces a token-to-expert assignment denoted by three outputmatrices I , G and P .

Original: The matrix I is an index matrix where I [ i, j ] speciﬁes j -th selected token of the i -th expert.
Revised: The matrix I is an index matrix where I [ i, j ] speciﬁes j -th selected token ofthe i -th expert.

Original: The gating matrix G ∈ R e × k denotes the weight of expert for the selected token, and P ∈ R
Revised: The gating matrix G ∈ R e × k denotes the weight of expert for the selected token,and P ∈ R

Original: e
Revised: e

Original: × k × n refers to an one-hot version of I that will be used to gather tokens for each expert.
Revised: × k × n refers to an one-hot version of I that will be used to gather tokens for each expert.

Original: These matrices are computed using a gating function,
Revised: These matrices are computed using a gating function,

Original: [Equation]
Revised: [Equation]

Original: where S denotes the token-to-expert afﬁnity scores, W g ∈ R d × e denotes the expert embeddings, and
Revised: where S denotes the token-to-expert afﬁnity scores, W g ∈ R d × e denotes the expert embeddings, and

Original: TopK () selects the k largest entries for each row of S (cid:62) .
Revised: TopK () selects the k largest entries for each row of S (cid:62) .

Original: Similar to Switch Transformer [10] and GShard [21], we apply mixture of experts and the gating function in the dense feed-forward (FFN) layer, as it is the most computationally expensive part in a Transformer-based network.
Revised: Similar to Switch Transformer [10] and GShard [21], we apply mixture of experts and the gatingfunction in the dense feed-forward (FFN) layer, as it is the most computationally expensive part ina Transformer-based network.

Original: The input to the gated FFN, denoted by X in ∈ R e × k × d , is produced using the permutation matrix P .
Revised: The input to the gated FFN, denoted by X in ∈ R e × k × d , is producedusing the permutation matrix P .

Original: Here X in [ i
Revised: Here X in [ i

Original: ] ∈
Revised: ] ∈

Original: R k × d denotes the input of the i -th expert.
Revised: R k × d denotes the input of the i -th expert.

Original: Similarly, let W 1 and W 2 denote the parameters of gated FFN in which W 1 [ i ] and W 2 [ i ] ∈ R d × d (cid:48) denote the parameter matrices of the i -th expert.
Revised: Similarly,let W 1 and W 2 denote the parameters of gated FFN in which W 1 [ i ] and W 2 [ i ] ∈ R d × d (cid:48) denote theparameter matrices of the i -th expert.

Original: We compute the output of each expert X e
Revised: We compute the output of each expert X e

Original: [ i ] as follows,
Revised: [ i ] as follows,

Original: [Equation]
Revised: [Equation]

Original: We omit the bias terms here for brevity.
Revised: We omit the bias terms here for brevity.

Original: The ﬁnally output of the gated FFN layer X out ∈ R n × d can be obtained given X e , the permutation and gating matrices P and G ,
Revised: The ﬁnally output of the gated FFN layer X out ∈ R n × d canbe obtained given X e , the permutation and gating matrices P and G ,

Original: [Equation]
Revised: [Equation]

Original: Both X e and X out can be efﬁciently computed using Einstein summation (einsum) operations.
Revised: Both X e and X out can be efﬁciently computed using Einstein summation (einsum) operations.

Original: 3.3 Expert Choice with Additional Constraint
Revised: 3.3 Expert Choice with Additional Constraint

Original: We also consider regularizing our expert choice routing by limiting the maximum number of experts for each token.
Revised: We also consider regularizing our expert choice routing by limiting the maximum number of expertsfor each token.

Original: We are interested in whether adding this constraint improves pre-training and ﬁnetuning results.
Revised: We are interested in whether adding this constraint improves pre-training and ﬁnetuning results.

Original: More importantly, it helps analyzing to what degree using a variable number of experts per token affects the model performance.
Revised: More importantly, it helps analyzing to what degree using a variable number of expertsper token affects the model performance.

Original: Let A ∈ R e × n be a positive matrix
Revised: Let A ∈ R e × n be a positive matrix

Original: where A [ i, j ] represents whether the i-th expert selects j-th token.
Revised: where A [ i, j ] represents whether the i-th expert selects j-th token.

Original: We solve the following entropy-regularized linear programming problem
Revised: We solve the following entropy-regularized linear programming problem

Original: [Equation]
Revised: [Equation]

Original: where < S (cid:62) , A > denotes the inner product, H ( A ) is the sum of element-wise entropy 1 , and b >is an integer that upper bounds the selection for each token.
Revised: where < S (cid:62) , A > denotes the inner product, H ( A ) is the sum of element-wise entropy 1 , and b > 0is an integer that upper bounds the selection for each token.

Original: Adding a small entropy term gives a near-integer solution while enabling a fast iterative solver we can run on TPUs.
Revised: Adding a small entropy term gives anear-integer solution while enabling a fast iterative solver we can run on TPUs.

Original: Speciﬁcally, the solution space is the intersection of three convex sets each satisfying one of the linear constraints.
Revised: Speciﬁcally, thesolution space is the intersection of three convex sets each satisfying one of the linear constraints.

Original: We use Dykstra’s algorithm [9] that alternatively projects the intermediate solution onto one of the convex sets.
Revised: We use Dykstra’s algorithm [9] that alternatively projects the intermediate solution onto one of theconvex sets.

Original: 2 After A is computed, the routing indices I is selected using TopK ( A, k ) instead.
Revised: 2 After A is computed, the routing indices I is selected using TopK ( A, k ) instead.

Original: 3.4 Model Architecture
Revised: 3.4 Model Architecture

Original: At the high level, we adopt the idea of sparsely activated Mixture-of-Experts (MoE) [31].
Revised: At the high level, we adopt the idea of sparsely activated Mixture-of-Experts (MoE) [31].

Original: We use a Transformer architecture and replace the feed-forward component of every other Transformer  layer with a MoE layer, following recent practice [21, 10].
Revised: We usea Transformer architecture and replace the feed-forward component of every other Transformer [Table] layer with a MoE layer, following recent practice [10, 21].

Original: Interleaving regular Transformer layers and MoE layers empirically improves model performance and training efﬁciency, probably because [Equation] [Table] forcing some shared components in between MoE layers can mitigate the negative effects of skipping tokens.
Revised: Interleaving regular Transformer layersand MoE layers empirically improves model performance and training efﬁciency, probably becauseforcing some shared components in between MoE layers can mitigate the negative effects of skippingtokens.

Original: Several additional modiﬁcations adopted in recent work have been applied in our experiments.
Revised: Several additional modiﬁcations adopted in recent work have been applied in our experiments.

Original: For example, we replace the standard positional embedding with per-layer relative positional bias [5].
Revised: For example, we replace the standard positional embedding with per-layer relative positional bias [5].

Original: In the non-MoE feed-forward sub-layers (only every other layers are MoE layers), we replace the ﬁrst linear projection and the activation function with the Gated Linear Unit [6], which computes the component-wise product of two linear transformation of the input, followed by a Gaussian Error
Revised: In the non-MoE feed-forward sub-layers (only every other layers are MoE layers), we replace theﬁrst linear projection and the activation function with the Gated Linear Unit [6], which computesthe component-wise product of two linear transformation of the input, followed by a Gaussian Error

Original: Linear Unit [15] activation function.
Revised: Linear Unit [15] activation function.

Original: As described earlier, each MoE layer consists of a group of independent feed-forward networks as denoted as “experts”.
Revised: As described earlier, each MoE layer consists of a group of independent feed-forward networks asdenoted as “experts”.

Original: The gating function in Equation 2 uses a softmax activation function to model a probability distribution over these experts.
Revised: The gating function in Eq. (2) uses a softmax activation function to model aprobability distribution over these experts.

Original: This distribution denotes the preference over experts of each incoming token, which is computed similarly in a conventional gating network [31, 10, 21].
Revised: This distribution denotes the preference over experts ofeach incoming token, which is computed similarly in a conventional gating network [10, 21, 31].

Original: During training, each MoE layer’s learnable gating network described in Equation 2 is trained to use the input to activate the best subset of experts using a top- k function along the token dimension.
Revised: During training, each MoE layer’s learnable gating network described in Eq. (2) is trained to usethe input to activate the best subset of experts using a top- k function along the token dimension.

Original: An “shufﬂe” stage and an “unshufﬂe” stage are inserted to the MoE layer, where the ﬁrst stage gathers the tokens to their designated experts while the second stage permutes the tokens back to their original order in the input batch.
Revised: An“shufﬂe” stage and an “unshufﬂe” stage are inserted to the MoE layer, where the ﬁrst stage gathers thetokens to their designated experts while the second stage permutes the tokens back to their originalorder in the input batch.

Original: This step is formulated in Equation 3 and 4.
Revised: This step is formulated in Eq. (3) and Eq.

Original: Similar to conventional MoE method, there are more parameters in the MoE layer.
Revised: Similar to conventional MoE method, there are more parameters in the MoE layer.

Original: However, the activated model size per token can be comparable to a dense layer because during training or inference, only a limited subset of experts is activated for any given token.
Revised: However, theactivated model size per token can be comparable to a dense layer because during training or inference,only a limited subset of experts is activated for any given token.

Original: For instance, Switch Transformer [10] has only one activated expert while GShard [21] uses two experts per token.
Revised: For instance, Switch Transformer [10]has only one activated expert while GShard [21] uses two experts per token.

Original: In our method, the number of activated experts can vary for each token but the overall computation is kept the same as the baseline architectures by ﬁxing the capacity factor c in Equation 1.
Revised: In our method, thenumber of activated experts can vary for each token but the overall computation is kept the same asthe baseline architectures by ﬁxing the capacity factor c in Eq.

Original: Unless otherwise speciﬁed, we set c = 2 such that our method can be directly compared to the top-2 token-choice gating in GShard.
Revised: Unless otherwise speciﬁed, we setc = 2 such that our method can be directly compared to the top-2 token-choice gating in GShard.

Original: 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.
Revised: We train several variants of our architecture at the 100M scale (i.e. 100M expert size) by increasingthe number of experts to understand the scaling effects of our method.

Original: We also train a 8B scale
Revised: We also train a 8B scale

Original: MoE model.
Revised: MoE model.

Original: The large MoE model is partitioned with a 2D sharding algorithm as presented in
Revised: The large MoE model is partitioned with a 2D sharding algorithm as presented in

Original: GSPMD [36], which fully exploits the 2D topology of the TPU cluster [19].
Revised: GSPMD [36], which fully exploits the 2D topology of the TPU cluster [19].

Original: Across different scales and setups, our method outperforms related work and demonstrates strong downstream task performance on selected tasks in GLUE and SuperGLUE.
Revised: Across differentscales and setups, our method outperforms related work and demonstrates strong downstream taskperformance on selected tasks in GLUE and SuperGLUE.

Original:  Experiments4.1 Setup
Revised:  Experiments

Revised: 4.1 Setup

Original: Table 1 summarizes the hyperparameter settings of different MoE models.
Revised: Table 1 summarizes the hyperparameter settings of different MoE models.

Original: As a reference point, we also include the respective dense model conﬁgurations with comparable numbers of activated parameters per-token during inference.
Revised: As a reference point,we also include the respective dense model conﬁgurations with comparable numbers of activatedparameters per-token during inference.

Original: To study of the effect of scaling the number of experts, we studied varying the number of experts but ﬁxing the per expert size to 100M parameters.
Revised: To study of the effect of scaling the number of experts, westudied varying the number of experts but ﬁxing the per expert size to 100M parameters.

Original: For example, 0.1B/64E represents the architecture of an approximately 100M parameter dense model with every other layer replaced by a 64-expert MoE layer.
Revised: For example,0.1B/64E represents the architecture of an approximately 100M parameter dense model with everyother layer replaced by a 64-expert MoE layer.

Original: The MoE model degenerates into a dense transformer architecture when each MoE layer only has one expert.
Revised: The MoE model degenerates into a dense transformerarchitecture when each MoE layer only has one expert.

Original: While n params is the total number of trainable parameters, n act − params represents the number of activated parameters per token.
Revised: While n params is the total number of trainableparameters, n act − params represents the number of activated parameters per token.

Original: L is the total number of Transformer layers, M is the model dimension, H is the hidden dimension after the projection in each transformer layer, n heads is the number of attention heads, and d head is the hidden dimension of each attention head.
Revised: L is the total number of Transformer layers, M is the model dimension, H is the hidden dimension after theprojection in each transformer layer, n heads is the number of attention heads, and d head is the hiddendimension of each attention head.

Original: 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.
Revised: Dataset: We use the high-quality dataset from GLaM [7] of 1.6 trillion tokens that are representativeof a wide range of natural language use cases.

Original: An in-house classiﬁer is trained to classify between a collection of curated text and other webpages and estimate the content quality of a webpage.
Revised: An in-house classiﬁer is trained to classify betweena collection of curated text and other webpages and estimate the content quality of a webpage.

Original: A high-quality ﬁltered subset of webpages are combined with books, Wikipedia pages, conversations, forums, and news to create the ﬁnal dataset.
Revised: Ahigh-quality ﬁltered subset of webpages are combined with books, Wikipedia pages, conversations,forums, and news to create the ﬁnal dataset.

Original: The data and mixture weights can be found in Table 3 in the GLaM paper.
Revised: The data and mixture weights can be found in Table 3 inthe GLaM paper.

Original: Model Training: Our model training follows the setups of GLaM [7] where a maximum sequence length of 1024 tokens is adopted.
Revised: Model Training: Our model training follows the setups of GLaM [7] where a maximum sequencelength of 1024 tokens is adopted.

Original: We use an Adafactor optimizer [32] with ﬁrst-moment decay β 1 = 0 and second-moment decay β 2 = 0 .
Revised: We use an Adafactor optimizer [32] with ﬁrst-moment decayβ 1 = 0 and second-moment decay β 2 = 0 .

Original: We keep the learning rate constant for the ﬁrst 10K training steps, and then decay it with an inverse square root schedule.
Revised: We keep the learning rate constant for the ﬁrst 10Ktraining steps, and then decay it with an inverse square root schedule.

Original: Unlike most related works, we do not impose any auxiliary loss for load balance, such as described in Switch Transformer [10] and
Revised: Unlike most related works, wedo not impose any auxiliary loss for load balance, such as described in Switch Transformer [10] and

Original: GShard [21].
Revised: GShard [21].

Original: We use the SentencePiece subword tokenizer with a vocabulary of size of 256K. The largest model (8B/64E) is trained on 512 TPU V4 chips.
Revised: We use the SentencePiece subword tokenizer with a vocabulary of size of 256K. Thelargest model (8B/64E) is trained on 512 TPU V4 chips.

Original: We use a dropout rate of 0 during training as the number of tokens in the training data corpus is much greater than the total number of tokens during training.
Revised: We use a dropout rate of 0 during trainingas the number of tokens in the training data corpus is much greater than the total number of tokensduring training.

Original: Model Evaluation: We mainly focus on evaluating the ﬁnetuning performance on the 11 selected tasks from GLUE and SuperGLUE benchmarks [35, 34].
Revised: Model Evaluation: We mainly focus on evaluating the ﬁnetuning performance on the 11 selectedtasks from GLUE and SuperGLUE benchmarks [34, 35].

Original: 4.2 Training Efﬁciency
Revised: 4.2 Training Efﬁciency

Original: [Figure]
Revised: [Figure]

Original: Training convergence is more than 2x faster using our method compared to GShard top-2 gating.

Original: (b) Training perplexity scales strongly with the number of experts while keeping the expert size ﬁxed.

Original: EC consistently outperforms GShard top-2 gating.

Original: We ﬁrst study training efﬁciency and convergence.
Revised: We ﬁrst study training efﬁciency and convergence.

Original: We use expert choice with a capacity factor of 2 (EC-CF2) to match the activated model size and computational cost on a per token basis in GShard top-2 gating and run both for a ﬁxed number of steps.
Revised: We use expert choice with a capacity factor of 2(EC-CF2) to match the activated model size and computational cost on a per token basis in GShardtop-2 gating and run both for a ﬁxed number of steps.

Original: The results are shown in Figure 2 (a).
Revised: The results are shown in Fig. 2 (a).

Original: Comparing to GShard top-2 gating, which showed stronger performance in both perplexity in the evaluation dataset and ﬁnetuning on downstream tasks compared to Switch Transformer top-1 gating, EC-CF2 converges more than 2x faster during training.
Revised: Comparingto GShard top-2 gating, which showed stronger performance in both perplexity in the evaluationdataset and ﬁnetuning on downstream tasks compared to Switch Transformer top-1 gating, EC-CF2converges more than 2x faster during training.

Original: More speciﬁcally, EC-CF2 reaches the same perplexity as GShard top-2 in less than half the steps, and with each GShard top-2 step being 20% slower than our method.
Revised: More speciﬁcally, EC-CF2 reaches the same perplexityas GShard top-2 in less than half the steps, and with each GShard top-2 step being 20% slower thanour method.

Revised: As explained in Section 3.1, the slower step time in top-2 gating can be explained by thelarge over capacity factor (e.g. can be as large as 8 in 8B64E setting) used due to load imbalance.

Revised: The step latency can be determined by the most loaded expert.

Original: 4.3 Scaling the Number of Experts
Revised: 4.3 Scaling the Number of Experts

Original: As presented in Table 1, increasing the number of experts effectively increases model capacity without increasing activated model size.
Revised: As presented in Table 1, increasing the number of experts effectively increases model capacity withoutincreasing activated model size.

Original: We scale the number of experts while ﬁxing the expert size to 100Mparameters for both expert choice (EC) and GShard (Top-2) methods and ﬁnd both methods scale well in terms of perplexity on the evaluation dataset during pre-training, as demonstrated in Figure
Revised: We scale the number of experts while ﬁxing the expert size to 100M

Original: [Table]
Revised: [Table]

Original: (b).
Revised: parameters for both expert choice (EC) and GShard (Top-2) methods and ﬁnd both methods scalewell in terms of perplexity on the evaluation dataset during pre-training, as demonstrated in Fig. 2(b).

Original: Having more experts consistently improves training perplexity, which often translates directly to better performance in downstream tasks.
Revised: Having more experts consistently improves training perplexity, which often translates directly tobetter performance in downstream tasks.

Original: 4.4 Fine-tuning on GLUE and SuperGLUE
Revised: 4.4 Fine-tuning on GLUE and SuperGLUE

Original: We perform ﬁne-tuning on 11 selected tasks from GLUE and SuperGLUE on the pretrained models of three MoE methods, including Switch Transformer top-1 gating (ST Top-1), GShard top-2 gating (GS Top-2) and our method with a capacity factor of 2 (EC-CF2) to match the activation memory size and thus computational cost of GS Top-2.
Revised: We perform ﬁne-tuning on 11 selected tasks from GLUE and SuperGLUE on the pretrained modelsof three MoE methods, including Switch Transformer top-1 gating (ST Top-1), GShard top-2 gating(GS Top-2) and our method with a capacity factor of 2 (EC-CF2) to match the activation memory sizeand thus computational cost of GS Top-2.

Original: Indicated by the results in Table 2, our EC-CF2 method consistently outperforms the related methods and yields more than 2% average accuracy increase in a large 8B/64E setting.
Revised: Indicated by the results in Table 2, our EC-CF2 methodconsistently outperforms the related methods and yields more than 2% average accuracy increase in alarge 8B/64E setting.

Original: Compared to the state-of-the-art T5 dense model [27], our 8B/64E model outperforms a T5 11B dense model  in 7 out of the 11 tasks, including CB (1.8%), COLA (15.3%),
Revised: Comparing to the 8B dense counterpart, our 8B MoE model achieves strongerpre-training perplexity and stronger ﬁne-tuning performance presented in Table 3.

Original: QQP (0.1%), RTE (1%), SST2 (0.3%), WiC (4.3%), WNLI (0.7%). Impressively, our 8B/64E model only has 9.8B activated parameters, and is more computationally efﬁcient than the T5 11B model.
Revised: Our best sparsemodel outperforms the 8B dense model by 3.4 points on average.

Original: 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 to better performance of downstream tasks.
Revised: 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. Thisresult indicates that a good training perplexity does not always translate to better performance ofdownstream tasks.

Original: 4.5 Heterogeneity Matters
Revised: 4.5 Heterogeneity Matters

Original: Capped Expert Choice: We regularized expert choice by limiting the maximum number of experts for each token, using the method described in Section 3.3.
Revised: Capped Expert Choice: We regularized expert choice by limiting the maximum number of expertsfor each token, using the method described in Section 3.3.

Original: Table 3 reports the average accuracy on the 11 selected datasets.
Revised: Table 4 reports the average accuracy onthe 11 selected datasets.

Original: EC-CAP2 is the variant of our expert choice method by limiting the number of experts of each token to 2.
Revised: EC-CAP2 is the variant of our expert choice method by limiting the numberof experts of each token to 2.

Original: This decreases the ﬁne-tuning accuracy by 0.8 points on average.
Revised: This decreases the ﬁne-tuning accuracy by 0.8 points on average.

Original: In addition, EC-CAP3 allows a maximum of 3 experts per token and achieves on par results compared to the vanilla expert choice method.
Revised: Inaddition, EC-CAP3 allows a maximum of 3 experts per token and achieves on par results compared [Table] [Table] [Table]

Revised: to the vanilla expert choice method.

Original: This ablation study conﬁrms that allowing variable number of experts per token is indeed helpful.
Revised: This ablation study conﬁrms that allowing variable number ofexperts per token is indeed helpful.

Original: Additional results can be found in Supplementary Section C.
Revised: Additional results can be found in Supplementary Section C.

Original: [Table] [Table] Variable Experts per Token: We compute statistics on token-to-expert routing, particularly on the ratio of tokens that have been routed to a certain number of experts.
Revised: Variable Experts per Token: We compute statistics on token-to-expert routing, particularly on theratio of tokens that have been routed to a certain number of experts.

Original: According to Figure 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 4 experts.
Revised: According to Fig. 3, a majorityof tokens have been routed to one or two experts while 23% have been routed to three or four expertsand only about 3% tokens have been routed to more than 4 experts.

Original: This plot veriﬁes our hypothesis that our method learns to allocate a variable number experts to tokens, which can be beneﬁcial for important tokens.
Revised: This plot veriﬁes our hypothesisthat our method learns to allocate a variable number experts to tokens, which can be beneﬁcial forimportant tokens.

Original: 4.6 Comparison with Hash Layer
Revised: 4.6 Comparison with Hash Layer

Original: In this section, we compare our method with Hash Layers [28].
Revised: In this section, we compare our method with Hash Layers [28].

Original: We use mod x to map a token ID to an expert ID.
Revised: We use mod x to map a token IDto an expert ID.

Original: This ensures load balance and generates specialized experts.
Revised: This ensures load balance and generates specialized experts.

Original: The ﬁne-tuning results are presented in the last row in Table 3 and more detailed scores can be found in Supplementary
Revised: The ﬁne-tuning resultsare presented in the last row in Table 4 and more detailed scores can be found in Supplementary

Original: Table 2.
Revised: Table 2.

Original: Hashing based routing performs much worse than expert choice in terms of average scores and variance.
Revised: Hashing based routing performs much worse than expert choice in terms of average scoresand variance.

Original: This indicates that load balancing alone does not generate all the beneﬁts .
Revised: This indicates that load balancing alone does not generate all the beneﬁts .

Original: 4.7 Ablation

Original: [Figure]
Revised: [Figure]

Original: E v a l P e [Figure] p l e
Revised: Ablation

Original: [Figure]
Revised: [Figure]

Original: y

Original: [-EC_10064E DENSE_
-]
Revised: [-EC_10064E DENSE_
-]

Original: M

Original: EC_8B64E DENSE_8B [Figure] Capacity Factor: We study the capacity factor in our expert choice method and compare the training perplexity with the baseline top-1 gating method used in Switch Transformer.
Revised: Capacity Factor: We study the capacity factor in our expert choice method and compare the trainingperplexity with the baseline top-1 gating method used in Switch Transformer.

Original: As described in 
Revised: As described in Eq.

Original: Equation 1, the capacity factor determines how many experts on average each token can be routed to, thus the bucket size k of each expert.
Revised: (1),the capacity factor determines how many experts on average each token can be routed to, thusthe bucket size k of each expert.

Original: 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.
Revised: 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.

Original: To match the computation cost on a per-token basis fairly with top-1 gating used in Switch Transformer, we reduce the capacity factor to 1 and plot the training perplexity.
Revised: To matchthe computation cost on a per-token basis fairly with top-1 gating used in Switch Transformer, wereduce the capacity factor to 1 and plot the training perplexity.

Original: Figure 4 (a) demonstrates that our method with a capacity factor of 1 ( EC _ CF 1 ) still signiﬁcantly outperforms top-1 gating.
Revised: Fig. 4 (a) demonstrates that ourmethod with a capacity factor of 1 ( EC _ CF 1 ) still signiﬁcantly outperforms top-1 gating.

Original: We further pushes the capacity factor further down to 0.5, and ﬁnd it only marginally hurts training perplexity, outperforming top-1 gating.
Revised: We furtherpush the capacity factor further down to 0.5, and ﬁnd it only marginally hurts training perplexity,outperforming top-1 gating.

Original: Not surprisingly, using a larger capacity factor in expert choice yields better performance.
Revised: Not surprisingly, using a larger capacity factor in expert choice yieldsbetter performance.

Original: More details on downstream tasks are presented in Supplementary Section B.
Revised: More details on downstream tasks are presented in Supplementary Section B.

Original: Comparison with Dense Models: We compare our method with dense models where no MoE are used.
Revised: Comparison with Dense Models: We compare our method with dense models where no MoE areused.

Original: As shown in Figure 4 (b), our method consistently outperforms the dense method in training perplexity and training convergence time.
Revised: As shown in Fig. 4 (b), our method consistently outperforms the dense method in trainingperplexity and training convergence time.

Original: For a small expert size of 100M parameters, the beneﬁt of sparse gating is even more signiﬁcant.
Revised: For a small expert size of 100M parameters, the beneﬁt ofsparse gating is even more signiﬁcant.

Original: Orthogonal to results presented in Figure 2 (b), where scaling the number of experts improves model performance, Figure 4 (b) shows that increasing expert capacity also signiﬁcantly increases model performance.
Revised: Orthogonal to results presented in Fig. 2 (b), where scaling thenumber of experts improves model performance, Fig. 4 (b) shows that increasing expert capacityalso signiﬁcantly increases model performance.

Original: We conduct ﬁne-tuning on the selectedtasks described in Table 2 on a dense 8B model and report the average score of 89.2, which is 3.4% lower than 8B/64E EC-CF2 in our method.
Revised: We conduct ﬁne-tuning on the selected 11 tasksdescribed in Table 2 on a dense 8B model and report the average score of 89.2, which is 3.4% lowerthan 8B/64E EC-CF2 in our method.

Original: More details are in Supp.
Revised: More details are in Supp.

Original: Section A.
Revised: Section A.

Original:  Conclusion
Revised:  Conclusion

Original: We propose a new routing method for sparsely activated mixture-of-experts (MoE) models.
Revised: We propose a new routing method for sparsely activated mixture-of-experts (MoE) models.

Original: This method addresses load imbalance and under-utilization of experts in conventional MoE methods, and enables selecting different numbers of experts for each token.
Revised: Thismethod addresses load imbalance and under-utilization of experts in conventional MoE methods,and enables selecting different numbers of experts for each token.

Original: Our model demonstrates more than 2x training efﬁciency improvements when compared to the state-of-the-art GShard and Switch
Revised: Our model demonstrates morethan 2x training efﬁciency improvements when compared to the state-of-the-art GShard and Switch

Original: Transformer models, and also achieves strong gains when ﬁnetuning on 11 datasets in the GLUE and
Revised: Transformer models, and also achieves strong gains when ﬁnetuning on 11 datasets in the GLUE and

Original: SuperGLUE benchmark.
Revised: SuperGLUE benchmark.

Original:  Limitations
Revised:  Limitations

Original: The expert choice method might not immediately apply to auto-regressive text generation as our current implementation takes in the past and future tokens to perform the top- k selection.
Revised: The expert choice method might not immediately apply to auto-regressive text generation as ourcurrent implementation takes in the past and future tokens to perform the top- k selection.

Original: One possible solution is to collect a large batch of input sequences, dispatch tokens of the same sequence into separate groups, and perform expert choice routing for each group.
Revised: Onepossible solution is to collect a large batch of input sequences, dispatch tokens of the same sequenceinto separate groups, and perform expert choice routing for each group.

Original: Another scenario where the expert choice method does not immediately apply is when the batch size becomes very small during serving or inference.
Revised: Another scenario where theexpert choice method does not immediately apply is when the batch size becomes very small duringserving or inference.

Original: A global top- k can be selected instead and we can cap the number of times each expert or token gets selected.
Revised: A global top- k can be selected instead and we can cap the number of times eachexpert or token gets selected.

Original: We leave these possible improvements for future work.
Revised: We leave these possible improvements for future work.

Original: Another long-standing issue with MoE has been the large memory footprint.
Revised: Another long-standing issue with MoE has been the large memory footprint.

Original: Even though computational cost can be reduced using sparsely gated networks, the total number of parameters increases linearly or sub-linearly with the number of experts.
Revised: Even though computational cost can be reduced using sparsely gated networks, the total number of parameters increaseslinearly or sub-linearly with the number of experts.

Original: Increasing the number of experts requires reservation of a large number of hardware devices.
Revised: Increasing the number of experts requires reservation of a large number of hardware devices.

Original: Therefore, dynamic (used) power is saved while static (reserved) power is not.
Revised: Therefore, dynamic (used) power is saved while static(reserved) power is not.

Original: Power saving techniques such as the ability to put hardware devices into low power states while not in use [17] can help with reducing the reserved power requirements.
Revised: Power saving techniques such as the ability to put hardware devices into lowpower states while not in use [17] can help with reducing the reserved power requirements.

Original:  Checklist
Revised:  Checklist

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Original: Yes (b) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: Yes (b) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: Yes (c) Did you discuss any potential negative societal impacts of your work?
Revised: Yes(c) Did you discuss any potential negative societal impacts of your work?

Original: N/A. Not any.
Revised: N/A. Not any.(d)

Original: (d) Did you describe the limitations of your work?
Revised: Did you describe the limitations of your work?

Original: Yes (a) Did you include the code, data, and instructions needed to reproduce the main experimental results?
Revised: Yes(a) Did you include the code, data, and instructions needed to reproduce the main experimentalresults?

Original: Yes.
Revised: Yes.

Original: We include details in the experiment setup to help reproduce the main results.
Revised: We include details in the experiment setup to help reproduce the main results.(b)

Original: (b) Did you specify all the training details?
Revised: Did you specify all the training details?

Original: Yes (c) Did you report error bars?
Revised: Yes(c) Did you report error bars?

Original: Yes (d) Did you include the amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: Yes(d) Did you include the amount of compute and the type of resources used (e g., type of GPUs,internal cluster, or cloud provider)?

Original: Yes (a) If your work uses existing assets, did you cite the creators?
Revised: Yes (a) If your work uses existing assets, did you cite the creators?

Original: Yes (b) Did you mention the license of the assets?
Revised: Yes (b) Did you mention the license of the assets?

Original: No.
Revised: No.

Original: The used dataset is not released yet.
Revised: The used dataset is not released yet.(c)

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: Did you include any new assets either in the supplemental material or as a URL?

Original: No.
Revised: No.

Original: The dataset is not released yet.
Revised: The datasetis not released yet.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Original: No.
Revised: No.

Original: Not using persons’ data.
Revised: Not using persons’ data.(e)

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
Revised: Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

Original: Yes.
Revised: Yes.

Original: The dataset does not contain any personally identiﬁable information or offensive content.
Revised: The dataset does not contain any personally identiﬁableinformation or offensive content.