Exploring Routing Strategies for Multilingual Mixture-of-Experts Models
Authors: Sneha Kudugunta
Yanping Huang
Ankur Bapna
Maxim Krikun
Dmitry Lepikhin
Thang Luong
Orhan Firat
Venue:  (2021)
Keywords: Mixture-of-Experts
Neural Machine Translation
Multilingual
Multi-Task Learning
Conditional Computation
Natural Language Processing

Abstract: Sparsely-Gated Mixture-of-Experts (MoE) has been a successful approach for scaling multilingual translation models to billions of parameters without a proportional increase in training computation. These models, however, are prohibitively large for serving deployment and there is no easy way to extract a sub-network to decode for a particular language pair. This work proposes improved strategies to route MoE models by tasks instead of tokens, thus enabling separation of network structures at decoding time while enjoying the benefits of scale and task sharing at training time.
We compare routing strategies at multiple levels (token, sentence, task) in both, the encoder and the decoder, and conduct extensive experiments on two benchmarks: the public WMT dataset of 30 language pairs and an in-house web-scale dataset of 200 language pairs. On WMT, with a Transformer base model with 32 experts, our task-level MoE outperforms the best performing token-level MoE model by +1.0 BLEU on average over all language pairs. When scaling up to Transformer big model with 128 experts on the large-scale massively multilingual benchmark, our task-level MoE is competitive with token-level MoE while being able to reduce the decoder model size by a factor of $32.34$ and increase peak throughput by 2.6 times at inference.

Under review as a conference paper at ICLR 2021
EXPLORING ROUTING STRATEGIES FOR MULTILIN- GUAL MIXTURE-OF-EXPEE MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
of-urerss- Sparsely-Gated Mixture-of Experts (MoE) has been successful approach for scaling multilingual translation models to billions of parameters without a propor- tional increase in training computation. These models, however, are prohibitívely large for serving deployment and there no easy way to extract a sub-network to decode for particular language pair. This work proposes improved strate- gies to route MoE models by tasks instead tokens, thus enabling separation network structures at decoding time while enjoying the benefits of scale and task sharing at training time. We compare routing strategies at multiple levels (token, sentence, task) in both, the encoder and the decoder, and conduct extensive exper- iments on two benchmarks: the public WMT dataset of 30 language pairs and an in-house web-scale dataset of 200 language pairs. On WMT, with Transformer base model with 32 experts. our task-level MoE outperforms the best performing token-level MoE model by +1.0 BLEU on average over all language pairs. When scaling up to Transformer big model with 128 experts on the large-scale massively multilingual benchmark, our tese- MoE is competitive with token-level MoE while being able to reduce the decoder model size by factor of 32.34 and increase peak throughput by 2.6 times at inference.
INTRODUCTION
Scaling up neural network models has recently received great attention, given the significant quality improvements in variety of areas such as natural language understanding (Raffel et al., 2019 Brown et al. 2020) and multilingual machine translation (Huang al., 2019 Lepikhin al. 2020).
While training massive models on large amounts of data can almost guarantee improved quality, there are two factors affecting their practicality and applicability: (1) training efficiency and (2) ference efficiency. Large dense models are often prohibitively compute-intensive to train, with some models requiring TFlops-days sparsely-gate gated Mfxture-of Experts (MoE) layers
compute (Brown et al.) 2020). A recent line of work has proposed
efficient alternative to dense models (Shazeer) ct al. 2017; Lepikhin et al.] 2020: Riabinin & Gusev] 2020) in order to address training efficiency limitations.
Moł model each token of the input sequence activates dif- ferent subset of the experts, hence the computation cost per token becomes only proportional to the size of the activated sub-network. However, they fail meet requirements on inference efficiency.
vanilla sparsely-gas
Consider a long sequence where each token of the sequence activates a disjoint subset of available experts. From practical standpoint, the inference trace of the full sequence spans several experts independently for every token, resulting in an independent pathway for each token. Although this desired property adding flexibility to the model and increasing its capacity, becomes pro- hibitive for inference for the following reasons: The model parameters in these large models are beyond the memory limit of a single accelerator, and require model parallelism shard them across a cluster of devices during inference. For models with MoE Layers, the input token would be dy- namically routed to different experts allocated to different devices. This further adds communication cost across devices to the overall serving cost. Moreover, due to the sequential nature of the auto- regressive decoding (Kasai et al.) 2020 Chen et al. 2018), the added communication cost from model parallel decoders gets multiplied by the number of decoding steps. To add to this, serving MoE models efficiently requires batching large number of input tokens together, otherwise only a subset of the MoE network will be activated leading device under-utilization.

---

Under review as a conference paper at ICLR 2021
In this work, we study the inference efficiency of sparsely gated MoE models while taking into account the characteristics of the intended application, Multilingual Neural Machine Translation (MNMT). MNMT is an inherently multi-task learning problem, aimed at building single neural network for translating multiple language pairs simultaneously. In a MNMT model, the extent to which parameters are shared across languages determines the magnitude of positive transfer (Bald- win & Ford 1988) and conversely task interference due to the capacity bottleneck (Arivazhagan al., 2019). In ideal scenario, we would want to efficiently train single large MNMT model maximizing transfer while expanding the capacity bottleneck; at the same time, we would like to enjoy the benefits of sparsely activated sub-networks per-task inference time, i.e. extracting out sub-network from the model to decode for particular language pair to actualize inference efficiency.
We propose routing algorithms for MoE models with affordable serving costs. While vanilla MoEs route each sub-word token in the input to its preferred experts, we explore alternative routing strate- gies that leverage global task level information to route all tokens corresponding to particular task collectively the same set of experts. While this strategy could be perceived to be restrictive for parameter sharing across tasks, we empirically demonstrate that routing based on task boundaries performs better when applied to MNMT. During training, mix the inputs from different tasks in the same batch in order to learn the routing network and encourage positive transfer among the tasks. During inference, decode different tasks separately and only load the subset of experts associated with the corresponding task.
We compare our method with multilingual baselines and find that achieve significant gains two benchmarks: a multilingual WMT task with comparable inference cost (43.59 BLEU), described in Section 4 and a large internal dataset (43.6 BLEU), in Section 4.3.2) We see that the gains are comparable with conventional position-wise Mixture-of-Exper models while utilizing decoders with only a fraction (6.25% and 1.56%) of their serving cost. We discuss the trade-offs of these different methods in Section 3.2. In Section 4.3.4 we analyze the routing decisions made in MoE models and motivate our method.
SCALING TRANSFORMERS WITH MIXTURE-OF-EXPERT
The Transformer Vaswani et al. 2017) architecture is popular model used for neural machine translation and other natural language understanding problems. In eqquence-t--sqquenee (of which neural machine translation is one example), the model consists of a separate encoder and decoder, each of which contains multiple Transformer layers. For further details on Transformers, refer the reader to the original paper (Vaswani et al., 2017).
problems
We use the Mixture-of-Experts Transformer models used by Lepikhin et al. (2020), where the MoE layers for the Transformers consist of E feed-forward networks (FFN), such that (FFN, F...)..
FFN ww ReLUU(ie xx))
E
ys FFN.
Here, T, is the input token at position to the MoE layer and each FFN, is a two layer neural network using ReLU activation function. wi, and wo, are the input and output projection weights of the e-th expert. Finally, Ge.E is vector computed by the gating network. For each expert, most values of this vector are zeros, select few experts. The entries chosen from G. determine how much the expert contributes to the final output ys- Note that, in this work choose the top 2 weight experts for each example to be comparable with the prior work.
value being positive. We use this vector to route the token to
The gating network Gs must be considered carefully for efficiency purposes: (1) the utilization of experts must be balanced and (2) the function must be efficient to implement at scale. For a more thorough discussion of MoE transformer, we direct the reader to Lepikhin al. (2020).
3 METHODS
In this section we describe our candidate routing strategies the context of MNMT and discuss their trade-offs from the perspective of the training and inference efficiency. It is known that multi-
2

---

Under review as a conference paper at ICLR 2021
lingual models learn different overlapping representations depending the language this is true for both dense (Wu & Dredze 2019 Tiedemann, 2018: Tan et al. 2019 Östling & Tiedemann, 2016 Kudugunta et al. [2019) and sparse models (Section 4.3.4). Therefore propose changing the routing algorithm GATH( (s)) of MoEs to choose different experts using more natural separations.
3.1 ROUTING STRATEGIES
Given the sequential nature of the multilingual machine translation task. the routing decisions can be made at three different granularities, from bottom up (i) token-level, (ii) sentence-level and (iii) task-level, as detailed below.
Token-level Routing: This is the baseline discussed in Section where each token is routed independently.
Each sequence (sentence), and all tokens that form the sequence, are routed to the same expert. We change the routing algorithm to select experts by sentence representation, calculated by taking the average token representations in given sentence.
Sentence-level Routing:
Task-level Routing: We select experts by task boundaries as opposed to making input- level decisions. In the context of MNMT, these task boundaries can either be defined by the target language (French-to- -Frenih-to-Englishh and language pair (French-to-Englisl and German-to-English are different tasks).
are the same task) or the
S
gz.E = GATE
Σ
x.)) (Sentence-level routing)
S
G..E = GATE GATE(task-idd, (d.)
(Task-level routing)
(2)
We further illustrate the difference in Figure in token-based MoE models (Figure la tokens from each example are routed to different experts, whereas in task-level MoE models (Figure Tb tokens may be routed to the same expert based on task.
3.2 INFERENCE IMPLICATIONS OF ROUTING STRATEGIES
While the MoE models discussed in Shazeer et al. 2017 Lepikhin et al. 2020) train quickly relative to the number of parameters in terms of the wall clock time, they are expensive to serve.
Consider MoE with 512 experts and 50B parameters (Lepikhin al. 2020). When employing token-level routing, each token can be independently routed to different set of experts during single accelerator, the two potential solutions to utilize this model for inference are: (i) Loading experts dynamically from host device depending on routing decisions, (ii) Utilizing marell-parallelim over multiple accelerators for serving. While the first solution incurs heavy host-device communication costs, the second introduces significantly inter-device communication overhead.
inference. Given that the entire model is too large load into memory
Another practical approach to serve large MoE would require model compression via quantization, pruning distillation (Cheng et al. 2017). While the first two strategies haven't 't been explored the context of conditional computation, distillation (Hinton et al.) 2013: Kim & Rush 2016) has been found introduce undesirable artifacts into the student model (Freitag et al.) [019]] Bogoychev & Sennrich, 2019) in the context of NMT. On the other hand, if we limit the number of experts available to every task the model to a small fraction of the total available capacity, it is possible to extract task-specific models for serving, alleviating the need for complex serving strategies or compression. Since decoding time complexity for rugo-reegeesive seq2seq models is dominated by the decoder (Kasai et al. (2020), can also pursue hybrid strategy where the encoder utilizes more expensive routing strategies while the decoder of the model utilizes simpler and efficient routing.
We do note, however, that MoE models that route purely by task boundaries slower to train due to load balancing considerations. All examples in the input batch belonging to the same task would route to the same set of experts, possibly leading some experts bearing significant amount of the load. Balancing between these inference and training time trade-off' merits further exploration.
Summarizing the effective decoding cost of the MoB models utilizing different routing strategies:
3

---

Under review as a conference paper at ICLR 2021
Ippt Sentanee: WAl
Ippu Sentance (11
69654
Eipatr
(a) MoE
(b) Task MoF.
Figure I: Different tokens are routed to different experts in token-based MoE models(a), whereas they may be routed to the same expert based on task some other prior in task-base MoE (b).
Token/Sentence level routing: The routing decisions are made dynamically. Assuming each token/sentence makes disjoint choices, the server needs to load all E experts.
Task-level routing: Tokens corresponding to each input sentence are routed the same experts statically. The server only needs to pre-load K experts (assuming top-K routing).
4
EXPERIMENTS
We compare routing strategies at multiple levels in both, the encoder and the decoder, by conducting extensive experiments two benchmarks: the public WMT dataset with 30 language pairs (Section and an in-house web-scale dataset with 200 language pairs (Section 4.3)
4.1 SETUP FOR WMT EXPERIMENTS
For our experiments, use parallel training and evaluation data from the WMT corpus and adopt the setup used by (Siddhant et al. (2020) with 15 languages, to and from English. Full training data details may be found in Table in the Appendix. The amount of data ranges from more than 60 million sentence pairs en-cs translation direction (en-cs) to roughly 150k sentence pairs for en-gu.
We temperature based data sampling strategy train models, similar the strategy used to train the multilingual models in AArvazhagan et al. (2019): if pz. is the probability that a sentence in the corpus belongs to language pair L,, we sample from distribution where the probability of sampling from L is proportional to PL +. All the experiments in this paper are performed model trained with sampling temperature T = 5.
We use the 142M Transformer Base (Vaswani al. 2017) architecture (or enhanced versions of with MoE layers) for all of our experiments with WMT. Our models are optimized using Adafactor (Shazeer & Stern 2018) with momentum factorization and per-parameter norm clipping threshold of 1.0. We followed a learning rate of 3.0, with 40K warm-up steps for the schedule, which is decayed with the inverse square root the number of training steps after warm-up. BLEU scores presented in this paper are calculated using SacreBLEU (2018) the WMT test sets.
Multilingual baseline: We train Transformer Base model and Transformer Big on this dataset multilingual dense baselines. We share all parameters across language pairs, including the softmax layer and input/output word embeddings. We use a 64k token Sentence Piece vocabulary (Kudo & Richardson 2018). The vocabulary is shared both the encoder and decoder side. Each sentence pair has a <2xx> token pre-pended to the source sentence to indicate the target language, following Johnson et al. (2017).
Mixture of Experts Models: For MoE models, we replace the feed forward network (FFN) of
alternate layers of the Transformer with of identical FFN experts depicted in Figure Ta
4

---

Under review as a conference paper at ICLR 2021
System
Encoder
Routing Granularity
Decader
Fran All
Servin All
Parumeters
Serving Dec.
Average
KX-len
Fligh
Low
Biingual Bavelines
142M
142M
25M
21:01
21.81
15.9
25.15
11.31
Multilinuual Transforma--ua
142M
142M
25M
20.03
23.69
17.5
23.25
15.88
Muliilinuual Trandermer Big
473M
473M
151M
23.84
26.09
22,03
27.69
18.89
Tokes-level experts
Token
533M
22.58
27.49
Sentence lelience--leellaee Moll 32 expect
Sensence
Sentence
533M
533M
221M
19.87
24.05
16,83
22.56
16.14
Language Pair
Languige Fair
155M
32M
21,40
25.21
16.34
23.37
17.34
Target
Target
155M
32M
22.86
25.62
20.19
27.21
17.3
MoE experts
Language
338M
221M
Target
221M
22.33
Token
Language Pair
338M
32M
23.03
26.16
20.28
27.23
17.62
Token
Target
338M
32M
23.62
25.95
21,09
28.48
17:37
Table 1: Routing strategies for ceatturr-of-xxeerte (MoE) models We compare routing experts by either tokens, sentence representations, or tasks (using either language pairs or target languages). For task-level MoE, routing can also be different between encoder and decoder. For results, Average is the average results of all language pairs, whereas xx2en and en2xx are the averages of transla- tions into and from English respectively. High indicates high-resource language pairs (> million sentence pairs) while Low is for low-resource language pairs (< 1 million sentence pairs).
We provide the full individual BLEU scores in the Appendix [A.3] along with bilingual baselines. In addition, we provide the number of parameters for different components of our models in Appendix A.4
For brevity, we provide aggregate BLEU scores in Section 4.2
4.2 COMPARISON DIFFERENT ROUTING STRATEGIES ON WMT
We compare the teken- sentence-level and task-level routing strategies discussed in Section identical network size (32 experts, 533M parameters). The results are presented in Table [ In general, we find that all types of task-level routing performs better than token-level routing. We see that using sentence representations to route examples (Sentence-level MoE 32 experts) performs much worse, so do not conduct further experiments this setting.
When we use Task MoE on both the encoder and the decoder (Task-level MoE 32 experts: Tar- get/Target), we consistent gains across the board. To investigate this further, trained model that has (a) Token MoE on the encoder and Task MoE on the decoder (Task-level MoE - 32 experts: Token/Target Token/Language Pair) and (b) Task MoE the encoder and Token MoE the de- coder (Task-level MoB 32 experts: Target/Tokes or Language Pair/Token). In Table [ we see that using strategy (a) works the best, whether choose to route by the target language or the language pair. In Section 4.3.4. we discuss these observations further.
Overall we find that using Task MoE only on the decoder (Task-level MoE 32 experts: Token/Target works the best, with gains of BLEU over Token MoE. These gains are consistent across xx2en language pairs, en2xx language pairs, high resource languages (more than million sentence pairs), low resource languages and the 2 zero shot pairs.
While the MoE models considered outperform bilingual and multilingual Transformer- Base base- lines with comparable inference cost, they are slight outperformed by the multilingual Transformer- Big by 0.2 BLEU on average. Note that Transformer-Bi Big incurs much higher decoding cost. We measured our task-level MoE achieved 8.4x (338k vs 40.3k tokens/sec) higher peak decoding throughput. However, we reiterate that the motivation behind scaling sparsely is to increase capacity with little overhead while remaining competitive with dense models for example, while it is feasi- ble to train a 473M parameter model (with 8x inference cost), training a much larger dense models to say, 13B model (the size of our scaled up MoE model), is prohibitively slow and expensive.
4.3 SCALING TO MASSIVELY MULTILINGUAL, MASSIVE MT (M4)
We now scale our results up to a larger internal dataset with over 200 language pairs, while also scaling the number of parameters to beyond 10 billion weights. In addition, we look more closely at the gating decisions made by these sparse models and discuss their implications.
4.3.1 EXPERIMENTAL. SETUP
Data: We use an in-h training corpus generated by crawling and extracting parallel sentences from the web (Uszkoreit et al. 2010). This dataset has 204 direct language pairs (102 languages
5

---

Under review as a conference paper at ICLR 2021
Bilingual Baselines 100M
TranstormerMoE(111, 128 Experts) 138
TranstomnerTaskWoEE1222 128 Experts, Route by LP on Deceder) 138 TranstomeerTaskMGE122 128 Experts, Route Task Decoder) 138 Transformer122L Bk wide, heads) 473M
ARIFU
Languages
(a) Performance of different routing strategics on Xx-En language pairs.
Bilingual Baselines 100M
TranstormerMoE11ZL. 128 Experts) 13B
TransfermerTassMMoELLLL 128 Experts, Route LP Decoder) 13B Transfomerraask0o01122,,, 128 Experts, Task Decoder) 13B Iransformer(12L. 8k wide, heads) 473M
AIFUU
Languages
(b) Performance of different routing strategies En-Xx language pairs.
Figure 2: Comparing the performance of different routing strategies for Mixture-of-Exxe -xperts (MoE) models massively multilingual dataset We compare routing experts by tokens, and tasks (using either language pairs or target languages). Given that routing by token on the encoder and routing by task the decoder performed the best WMT (Table [T), for the scaled up 128 expert models we compare. We split the comparison of results into (a) Xx-En language pairs and (b) En-Xx language pairs. The languages on the x-axis sorted left-to-right in descending order of resource size. Best seen in color. Note that the token-level MoE has 6.5B parameters the decoders while our task-level MoE has only 200M.
those settings
to and from English), with a total of 25 billion sentence pairs. This dataset covers diverse range of domains and languages, and quite noisy. There is also a heavy imbalance when it comes to language pair, ranging between 10a and 10 sentence pairs.
the number of examples available In order to record gating decisions while controlling for semantics, we created multi-way aligned evaluation set containing nearly 3k sentence pairs for all languages
Each sentence in evaluation set is semantically identical all other languages.

---

Under review as a conference paper at ICLR 2021
Model: We use the 473M Transformer Big Vassaan et al., 2017) architecture (or modified versions of it in the case of sparse models) as described by Chen et al. (2018) for this set of experiments. Similar to Section 4.1 we (I) share all parameters across language pairs including softmax layer and input/output word embeddings, (2) pre-pend a <2xx> token to the source sentence to indicate
the target language and (3) use Sentence Piece ModeI|Kudo & Richardson (2018) with 64k tokens vocabulary shared on both the encoder and decoder side. We followed the training and architecture as shown in Lepikhin et al. (222
4.3.2 RESULTS
We compare Task-level MoEs and Token-level MoEs to their bilingual and multilingual baselines in Figure 2. We train 128 expert MoE models with routing in these settings: (1) Routing by token on both the encoder and decoder, (2) Routing by token on the encoder and by target language on the decoder and (3) Routing by token on the encoder and by language pair on the decoder.
We find that these scaled up sparse models perform better than their dense baselines, with hybrid task-level routing performing slightly better on En-Xx language pairs and pure token-level routing performing slightly better on Xx-En language pairs. We hypothesize that for the Xx-En tasks, not explicitly dividing expert parameters by tasks on the decoder results in better transfer, thus explain- ing the better performance of token-level routing. This suggests that a hybrid strategy that partially restricts access to experts based on task-boundaries, while still permitting routing by tokens, might provide the right balance between efficiency and quality.
We also note that while both forms of routing have 13B parameters (6.5B decoder) at train time, token level routing only on the decoder uses only 200M parameters at inference time, addition to the practical considerations discussed in Section 3.1 We provide aggregate BLEU scores in Appendix A.6 and parameter count breakdowns in Appendix A.5
4.3.3 COMPARISON OF THROUGHPUT ON MASSIVE MODELS
We further compare Task-level MoEs with Token-level MoEs in terms of throughput across different batch sizes in Figure We measure this by decoding the WMT14 English-German test set with our TaskMoE model and with the bascline TokenMoE model on 128 Cloud TPU V3 cores. We find that our Task-MoE model has 2.6 times higher peak throughput while using 32.34 times less decoder parameters (201M vs 6.5B). Moreover, our Task-MoE model has minimal communication overhead compared to decoding with Token-MoEE (0.2% versus 36% of step time).
We measured that the inference time of the token- MoE model is dominated by the decoder, with the decoders taking 49x the time per step than the encoders. Therefore, the inference cost of task-level routing on decoder only is roughly equivalent to that on both the encoder and decoder.
4.3.4 A CLOSER LOOK AT ROUTING DECISIONS
Now, we analyze the routing decisions made in token-level MoE models to further motivate our investigation. We take token-level MoE model trained on the massively multilingual dataset and decode these models on the multiway tests sets, while logging the routing decisions for every token. We plot the top expert distributions of several tasks with different scripts and language families in Figure For clarity, and because these two groups of languages behave differently in a multilingual setting, we split the gating decisions into those for Xx-En and En-Xx language pairs.
In the encoder (Figure 3a), tokens from all tasks (Xx-En) seem to prefer the same set of few experts slightly over the others. On the other hand, in the decoder (Figure 3b) each task seems to have a slight preference for few experts over the others. Moreover, the set of experts appears to be similar for related languages. For example, English-Spanish and English- Catalan (two Romance Languages) have similar expert distributions and so do English-Russ
and Einlish-Ukranie (two
2 As opposed to displaying BLEU scores for each language pair, we place the baselines on the z-axis zero and report the ABLEU trendline each model we consider. In order to set these bilingual baselines, we train Neural Machine Translation models for each language pair (e.g. a single model for German- to tongglish)), tuned depending on the available training data for that given language We tuned batch-size and different values of regularizann regularization methods (e.g. dropout) in a Transformer-Big or Transformer Banse layout, for high or low- resourced languages respectively.
7

---

Under review as a conference paper at ICLR 2021
Slavic Common
Romenic
South-Central Dravidien
Benue-Congo
Indo-Iranian
South Dravidian
0.030
Pum
uc.m
ouan
0.018
97,0
0.012
9,.1
0.000
vrtn
53.00
90.m0
N2 A 8RR5888 2
$ - RE
Expert
(a) Gating decisions last layer the encoder for language pairs.
Slavic
Benue-Congo
Common Turkic
Indo-Iranian
South Dravidian
Khmer
41,m
0.060
25.00
0.045
0,7
0.030
0.015
eur
0.000
Expert
(b) Gating decisions of the last layer of the decoder for En-Xx language pairs.
We record the gating decisions of our MoE model trained on internal data on multiway parallel dataset. The darker a cell, corresponding to, say en-sr and the 37th expert, the more the expert is used. In (a) the encoder, tokens from all tasks (Xx-En) seem to prefer the same set of few experts slightly over the others; while (b) the decoder each task (En-Xx) seems to slightly prefer few experts over the other. Moreover, the set of experts appears to be similar for related languages. For example, English-Spanish and English-Cate Catalan (two Romance Languages) have similar expert distributions and so do English-Rusee Russian and Aaglish-Urranian (two Slavic Languages).
Figure 3:
Slavic Languages). In the Appendix A.7 we provide expert distribution plots for other layers of this model. In addition, we provide expert distributions the MoE model that routes tokens by target language discussed in Section
Our analysis suggest that, when using token-level routing, task-level decisions emerge naturally in the decoder, providing additional motivation for our proposed routing strategies.
RELATED WORK
Conditional Computation: Conditional computation Bengio et al. (2015). or routing examples
through the neural network by activating only sub-network of the network depending the input has seen success in large scale natural language processing (NLP) (Shazeer et al. (2017); (2017):Lepikkhn Lepikhin
et al. 2020y; [aapa al. (2019) and computer vision Yang al. (201)) tasks. A variety of
8

---

Under review as a conference paper at ICLR 2021
4020
$4e-such decoder Roratios
1.0mmm
13905
2006
3000
Setch Size
Figure 4: Inference cost analysis. Left, encoders inference cost is small (32ms 1565ms) compared to decoders for a 128 Token-MoE model with sequence lenght 64. Right, We measure the throughput of our Task-MoB model and baseline Token-MoE model across batch sizes and that the peak throughput of Task-MoE is 2.6 times higher. In comparison, the peak throughput of Transformer- Big with 473M parameters is 40.3k tokens/sec.
strategies can be used to route examples such as learning a function on the input Shazeer et al. (217 Lepikhin et al. (2020), computational budget or Bapna et al. (2019); Elbayad et al. (2019).
Multi-task Learning: Multi-tase learning Caruana 1997) can improve model performance across all tasks trained on due to regularization and positive transfer between related tasks. Here, sub- networks are be activated depending on the task which the input belongs some of these param- eters may be shared. This approach has seen success in variety of domains such classification, recommender systems and NL.P (Ma et al. (2019) 2018); [Clark et al. (20T9); Collobert Weston (2008); Ruder et al. (2019); Tan et al. (2019)). Like work, some of these models have been designed with inference benefits in mind (Ma et al. (2019)). this work we focus on multi-task learning in the case of multlingual NMT.
Multi-task learning for Multilingual NMT Models: Multi-task learning in multilingual mod- els has been well-studied: while complete parameter sharing is simple and works well (Johnson et al. 2017)), an optimal strategy for sharing parameters and possibly having languages-specific parameters would maximize transfer while minimizing interference [Hokamp al. (2019). Strate- gies involve allocating language specific hidden states, attention modules, decoders or additional specialized layers (Hokamp et al. (2019); Wang al. (0018); Gu et al. (2018); Bapna et al. (2019)) In addition some strategies involve grouping parameters by language group Fan et al. (2020); Tan al. (2019). Compared to these works, our approach parameter sharing designed scale mod- els without impacting inference efficiency (as opposed to simply adding linguage-speciffe capacity)
while still enjoying the benefits of scaling.
CONCLUSIONS
In this work discussed more inference friendly algorithms for routing tokens in Sparse Mixture- of-Fxperts models by making use of task boundaries. We empirically demonstrated that this new algorithm performs well as, better than, conventional token-based routing algorithms on two different datasets: the multilingual WMT setup covering 30 language pairs and a large internal dataset covering 200 language pairs. We discussed the trade-off's of these methods in terms of train- time and serving considerations. In addition, we looked more closely at large MoE models and how their gating decisions differ by task.
We conclude by highlighting that the algorithms that are more inference friendly while retaining the training speed advantages of Mixture-of-EExee of-Exxstss models are a promising direction for future exploration, motivating research on inference efficiency for large models.
REFERENCES
Naveen Arivazhagan, Ankur Bapna, Orhan Firat, Dmitry Lepikhin, Melvin Johnson, Maxim Krikun, Mia Xu Chen, Yuan Cao, George Foster, Colin Cherry, Wolfgang Macherey, Zhifeng Chen, and

---

Under review as a conference paper at ICLR 2021
Yonghui Wu. Massively multilingual neural machine translation in the wild: Findings and chal- lenges, 2019.
Timothy Baldwin and Kevin Ford. Transfer of training: A review and directions for future research. Personnel psychology, 41(1):63-105, 1988.
Ankur Bapna, Naveen Arivazhagan, and Orhan Firat. Simple, scalable adaptation for neural machine translation. arXiv preprint arXiv: (90kv::9990000 (990900878, 2019.
Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv: 2015.
Nikolay Bogoychev and Rico Sennrich. Domain, translationese and noise in synthetic data for neural machine translation, 2019.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14 165, 2020.
Rich Caruana. Multitask learning. Machine learning, 28(1):41-75, 1997.
Mia Xu Chen, Orhan Firat, Ankur Bapna, Melvin Johnson, Wolfgang Macherey, George Fos- ter, Llion Jones, Mike Schuster, Noam Shazeer, Niki Parmar, Ashish Vaswani, Jakob Uszko- reit, Lukasz Kaiser, Zhifeng Chen, Yonghui Wu, and Macduff Hughes. The best of both worlds: Combining recent advances in neural machine translation. Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 76-86, Melbourne, Australia, July 2018. Association for Computational Linguistics. URL http::/wwst ac web. org antho wttp://annooooooooooleee1 ogt/
Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang. A survey of model compression and acceleration for deep neural networks. arXiv preprint arXiv:1710.09282 2017.
Kevin Clark, Minh-Thang Luong. Urvashi Khandelwal, Christopher D Manning, and Quoc V Le. Bam! born-ggain multi-task networks for natural language understanding. arXiv preprint 2019.
Ronan Collobert and Jason Weston. A unified architecture for natural language processing: Deep neural networks with multitask learning. In Proceedings of the 25th international conference Machine learning, pp. 160-167, 2008.
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. arXiv preprint arXiv: 1910.
2019.
Angela Fan. Shruti Bhosale, Holger Schwenk, Zhiyi Ma. Ahmed El-Kishky, Siddharth Goyal, Man- deep Baines, Onur Celebi, Guillaume Wenzek, Vishrav Chaudhary, et al. Beyond english-centric multilingual machine translation. arXiv preprint arXiv:2010. ///00/00////........// 2020.
Markus Freitag, Isaac Caswell, and Scott Roy. APE at scale and its implications on MT evaluation biases. Proceedings of the Fourth Conference on Machine Translation (Volume 1: Research Papers), pp. 34 44, Florence, Italy, August 2019, Association for Computational I. Linguistics. doi: 10.18653/v1/W11 19////299....... URL htt.hhttp://// www a wwwb..wwbe..... wwtp:/tto/aaaoooooooooooooo /w1 5204
Jiatao Gu, Hany Hassan, Jacob Devlin, and Victor OK Li. Universal neural machine translation for extremely low resource languages. arXiv preprint arXiv: 08.88000063,, 2018.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint -xXi:15503002311 2015.
Chris Hokamp, John Glover, and Demian Gholipour. Evaluating the supervised and zero-shot per- formance of multi-lingual translation models. arXiv preprint arXiv: 1906.00655. 2019.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in neural information processing systems, pp. 103-112, 2019.
10

---

Under review as a conference paper at ICLR 2021
Melvin Johnson, Mike Schuster, Quoe V Le, Maxim Krikun, Yonghui Wu, Zhifeng Chen, Nikhil Thorat, Fernanda Viégas, Martin Wattenberg, Greg Corrado, et al. Google's ' multilingual neural machine translation system: Enabling zero-shot translation. Transactions of the Association for Computational Linguisties, 5:339-351, 2017.
Jungo Kasai, Nikolaos Pappas, Hao Peng, James Cross, and Noah A Smith. Deep encoder, shal- low decoder: Reevaluating the speed-quality tradeoff machine translation. arXiv preprint arXiv:2006. 10369, 2020.
Yoon Kim and Alexander M. Rush. Sequence-level knowledge distillation. In Proceedings of the 2016 Conference Empirical Methods in Natural Language Processing, pp. 1317-1327, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/11/D11-113.
URL https: www. acwweb. org/ ozgganthology/Diatt. 1139
Taku Kudo and John Richardson. Sentencepiece: simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint / 2018.
Sneha Reddy Kudugunta, Ankur Bapna, Isaac Caswell, Naveen Arivazhagan, and Orhan Firat. In- vestigating multilingual nmt representations scale. arXiv preprint arXiv: 19090099999999 2019.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
Jiaqi Ma. Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H Chi. Modeling task relation- ships in multi-task learning with multi-gate mixture-of-expert In Proceedings of the 24th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pp. 1930-1939 2018.
Jiaqi Ma, Zhe Zhao, Jilin Chen, Ang Li, Lichan Hong, and Ed H Chi. Snr: Sub-network routing for flexible parameter sharing in multi-task learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 216-223, 2019.
Robert Östling and Jörg Tiedemann. Continuous multilinguality with language vectors. arXiv preprint ar 1r1iv:1612.07486 2016.
Matt Post. A call for clarity in reporting BLEU scores. Proceedings of the Third Conference Machine Translation: Research Papers, pp. 186-191, Belgium, Brussels, October 2018. As-
sociation for Computational Linguistics. URL httpt:tt////.. /wwww aclweb. wtt. www.orrgaaaatootoooooooooooooo
W18 6319
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. arXiv preprint arXiv: 10683, 2019.
Maksim Riabinin and Anton Gusev. Learning @ home: Crowdsourced training of large neural net- works using decentralized mixture-of-experts. arXiv preprint arXiv:2002.04013. 2020.
Sebastian Ruder, Joachim Bingel, Isabelle Augenstein, and Anders Søgaard. Latent multi-tass chitecture learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4822-4829 2019.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. arXiv preprini arXiv: 1804.04235 2018.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outragcously large neural networks: The sparsely-gater mixture-of-expert layer. arXiv preprint arXiv: 17000065388 2017.
Aditya Siddhant, Ankur Bapna, Yuan Cao, Orhan Firat, Mia Chen, Sneha Kudugunta, Naveen Ari- vazhagan, and Yonghui Wu. Leveraging monolingual data with self-supervision for multilingual neural machine translation. arXiv preprint arXiv:2005.048) 04816, 2020.

---

Under review as a conference paper at ICLR 2021
Xu Tan, Yi Ren, Di He, Tao Qin. Zhou Zhao, and Tie Yan Liu. Multilingual neural machine trans- lation with knowledge distillation. arXiv preprint 1rXiv:1902.1000 2019.
Jörg Tiedemann. Emerging language spaces learned from massively multilingual corpora. arXiv preprint arXiv:1802.00273, 2018.
Jakob Uszkoreit, Jay Ponte, Ashok C Popat, and Moshe Dubiner. Large scale parallel docu- ment mining for machine translation. In Proceedings of the 23rd International Conference Computational Linguisties, pp. 1101-1109, Association for Computational Linguistics, 2010.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, pp. 5998-6008, 2017.
Yining Wang. Jiajun Zhang, Feifei Zhai, Jingfang Xu, and Chengqing Zong. Three strategies to im- prove one-to-many multilingual translation. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2955-2960, 2018
Shijie Wu and Mark Dredze. Beto, bentz, becas: The surprising cross-lingual effectiveness of bert. arXiv preprint arXiv:1904.0900 2019.
Brandon Yang, Gabriel Bender, Quoe V Le, and Jiquan Ngiam. Condconv: Conditionally parameter- ized convolutions for efficient inference. In Advances in Neural Information Processing Systems, pp. 1307-1318, 2019.
12


Reviews:
[{"review_id": "jkaxuPMhD0c", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper compares different routing strategies in Mixture-of-Experts for multilingual NMT, and proposes to route by tasks instead of token, which can achieve better or comparable translation accuracy measured by BLEU and also enable separation of network structures at decoding time with affordable serving costs. The paper claims that with task-level routing, the server only needs to pre-load K experts (assuming top-K routing for MoE layers) during inference, instead of loading all experts as in token/sentence level routing.\n\nOverall, I appreciate the analyses and comprehensive experiment studies in this paper, and also the large-scale in-house datasets used in experiments. The experiment findings about the different routing strategies in the encoder and decoder in en-xx and xx-en settings are also interesting. The connections between the gating distribution and the similar of languages in Figure 3 also make sense.\n\nHowever, I doubt the novelty and machine learning contribution of this paper: 1) The different routing strategies are natural and seems to have already been proposed by previous works. e.g., for task-level routing, [1] used similar kind of mixture of experts in the language level. 2) This paper simply studies different routing strategies, which is more like empirical analyses. Although the results are somewhat interesting, they are not surprising and most findings are in expectation. \n\n[1] Universal Neural Machine Translation for Extremely Low Resource Languages. https://arxiv.org/pdf/1802.05368.pdf\n\nMeanwhile, I have some questions on the experiment settings: when comparing with the single multilingual base model in Table 1 and Figure 2, the parameters of the MoE model are larger than the single multilingual model (e.g., 533M vs 142M in Table 1). Therefore, it is obvious that MoE model achieves better accuracy than smaller single multilingual model. The MoE models should compare with a single multilingual model with the same amount of parameters. \n\nAnother important thing I need to point out is that this paper seems to violate the anonymous policy of ICLR 2021. In Section 4.3.1, the paper says \"We use an in-house training corpus (Arivazhagan et al., 2019)\".  However, Arivazhagan et al., 2019 shows the authors from Google, which reveals the organization of authors in this paper.  Meta reviewer can further double confirm if this violates the policy.\n \nBesides, this paper is not carefully written and there are many typos which affect the reading. e.g., 1) Two \"wo_e\" in the line below the equation in Section 2; 2) \"to route the token to a select few experts\", there is an additional \"a\"; 3) \"a learning rate of a learning rate of 3.0\" in Section 4.1.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "VjqPw6bnR-0", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces several routing strategies for multilingual neural machine translation. The motivation is to train a single mixture model that can serve the training and prediction of multiple models. Specifically, several strategies are proposed: token-level, sentence-level and task-level. Experiments on WMT and massive multilingual NMT dataset show that the proposed approach outperforms the vanilla unified multilingual model. While this approach is simple and straightforward, I have some concerns. \n\nPros:\n- A mixture model is proposed for multilingual NMT, forming a hierarchical structure:  token-level, sentence-level, and task-level.\n- The general framework is reasonable, straightforward, and easy to implement.\n- Experiments are extensive (both the WMT data and massive NMT data).\n\n\nCons:\n- The general idea is not novel. Building a mixture model for multi-task learning has been well studied in the literature [1,2] (not cited). The relation may need to be clarified.\n- In the experiment part, the results of single models (not unified multilingual model) need to be reported.\n- Comparisons with other unified multilingual approaches are required [3,4],\n\n\n\n[1] Ma et al. Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts, KDD 2020\n\n[2] Ma et al. Snr: Sub-network routing for flexible parameter sharing in multi-task learning.\n\n[3] Tan et al. Multilingual Neural Machine Translation with Language Clustering.\n\n[4] Multilingual Neural Machine Translation with Knowledge Distillation.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "FVahvqP-LrJ", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors explore alternatives to the standard token-based routing in sparsely-gated MoE models for multilingual NMT. This exploration is motivated by the need for efficient inference in MoE models, for which token-based routing is a limitation. The alternative is task-based routing, where examples for a task are assigned to the same experts. This allows efficient device placement and request dispatch at inference time. The paper compares with the approaches as well as hybrid approaches where different parts of the network use different routing mechanisms. The results show that task level routing is comparable to token-level routing with the added benefit of inference efficiency. Performing task-based routing on the decoder side only gives better better translation quality, at the cost of inference efficiently. An analysis of routing decision in token-based routing justifies the design choices. \n\nOverall, the paper takes a focused problem and experimentally shows an solution that improves inference efficiency. The experiments are well-described and the paper is well-written. \nWhile, the result is interesting it is only a marginal contribution with little novelty in my opinion. It would also be interesting to see how this approaches compares with simpler approaches that deterministically allot parameters to different languages (Wang etal 2018, Bapna etal 2019, Zhang etal 2020)  or language groups (Fan etal 2020). \n \n- Yining Wang, Jiajun Zhang, Feifei Zhai, Jingfang Xu, and Chengqing Zong. Three strategies to improve one-to-many multilingual translation. EMNLP. 2018. \n- Ankur Bapna, Naveen Arivazhagan, and Orhan Firat. Simple, scalable adaptation for neural machine translation. EMNLP. 2019.\n- Biao Zhang, Philip Williams, Ivan Titov, and Rico Sennrich. Improving massively multilingual neural machine translation and zero-shot translation. ACL. 2020.\n- Fan, Angela and Bhosale, Shruti and Schwenk, Holger and Ma, Zhiyi and El-Kishky, Ahmed and Goyal, Siddharth and Baines, Mandeep and Celebi, Onur and Wenzek, Guillaume and Chaudhary, Vishrav and Goyal, Naman and Birch, Tom and Liptchinsky, Vitaliy and Edunov, Sergey and Grave, Edouard and Auli, Michael and Joulin, Armand. Beyond English-Centric Multilingual Machine Translation.  arXiv:2010.11125 preprint. 2020", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes routing strategies for multilingual NMT. The motivation is to train a single mixture model that can serve the training and prediction of multiple models. Several strategies are proposed: token-level, sentence-level and task-level. This is a simple and straightforward approach (which is fine). The main concerns from the reviewers regard novelty and missing comparisons. In their updated draft, the authors added comparisons to bilingual models and they added a discussion wrt related work. However, the author’s response did not address enough some of other reviewers’ concerns regarding comparison with other approaches, and the lack of novelty persists (mixture models for multi-task learning have been previously proposed in the literature), which makes me lean towards rejection. I suggest the authors address these aspects in future iterations of their work.

Decision: Reject