SCoMoE: Efficient Mixtures of Experts with Structured Communication
Authors: zhiyuan zeng
Deyi Xiong
Venue: ICLR 2023 poster (2023)

Abstract:   Mixture-of-Experts (MoE) models are promising architectures for massively multilingual neural machine translation and large language models due to the advantage of sublinear scaling. However, the training of large MoE models is usually bottlenecked by the all-to-all communication (Lepikhin et al., 2020). To reduce the communication cost, we propose SCoMoE, an MoE architecture with structured all-to-all communication, inspired by the hierarchical architecture of the communication topology. SCoMoE encourages data to be communicated across devices through fast intra-accelerator/node communication channels, reducing communication throughput in the slow inter-node communication channel. We slice the data on the sequence dimension (SCoMoE-Seq) into three communication groups and project the data on the feature dimension (SCoMoE-Feat) into low-dimensional representations. To compensate the potential performance drop caused by the routing locality in SCoMoE, we further propose a token clustering approach to aggregating related tokens from different devices before the MoE layers. The sigmoid gating in the balanced router used in the token clustering is substituted with the softmax gating with differential sorting. Experiments on bilingual and massively multilingual machine translation demonstrate that SCoMoE achieves a speedup of 1.44x over GShard with comparable performance, and substantially outperforms Gshard (2.8 BLEU) on OPUS-100 with a speedup of 1.25x.

Published as a conference paper at ICLR 2023
EFFICIENT MIXTURES OF EXPERTS WITH STRUCTURED COMMUNICATION
SCoMoE:
Zhiyuan Zengi
Deyi Xiong' Tianjin University, GTCOM
Tianjin University zh yuay zengêt edu.cn
ABSTRACT
ofnture-of-Expertt (MoE) models are promising architectures for massively mul- tilingual neural machine translation and large language models due to the advan- tage of sublinear scaling. However, the training of large MoE models is usu- ally bottlenecked by the all-to-all communication (Lepikhin et al. (020). To re- duce the communication cost, we propose SCoMoE, an MoE architecture with structured all-to-all communication, inspired by the hierarchical architecture of the communication topology. SCoMoe encourages data to be communicated across devices through fast intra-accelerator//ode communication channels, re- ducing communication throughput in the slow inter-node communication channel. We slice the data on the sequence dimension (SCoMoE-Seq) into three communi- cation groups and project the data on the feature dimension (SCoMoF-Feat) into low-dimensional representations. To compensate the potential performance drop caused by the routing locality in SCoMoE, we further propose a token cluster- ing approach to aggregating related tokens from different devices before the MoE layers. The sigmoid gating in the balanced router used in the token clustering is substituted with the softmax gating with differential sorting. Experiments on bilingual and massively multilingual machine translation demonstrate that SCo- MoE achieves a speedup of 1.44x over GShard with comparable performance, and substantially outperforms Gshard (2.8 BLEU) OPUS-100 with speedup 1.25x. Codes are available httth::: /ttts:://iithaee com/ Zhi YuanZeng/ fa fairseq-
INTRODUCTION
Recent years have witnessed substantial interest in exploring sparse architectures based on Mixture of Experts for training massively multilingual machine translation (Lepikhin et al. 2020 Kim et al., 2021 and large language models (Fedus al. 2021 Zhang et al. [022bb Ma et al. 2022 Du al. 202T Zoph et al., 2022, Rajbhandari et al.) 2022]] Lin et al. 2021). Experts of MoE models are distributed over multiple devices. Due to the sparse architecture where only combination of experts are selected to process each input, the number experts and hence the scale of MoE models can be sufficiently large while the computational cost is only sublinear the number of parameters. Despite the advantage of efficient computation, MoE models require expensive all-to-all communication, to send the inputs and outputs of experts across the compute network. Previous study on GShard (Lepikhin et al. 2020) has shown that as MoE models scale, the all-to-all communication cost becomes the bottleneck for training.
To mitigate this issue, we propose Structured Communication based MoE (ScoMoE), which treats the all-to-all communication a structured way rather than equally across different devices. The motivation behind SCoMoE is that the network bandwidth between devices and nodes is different across the compute network: the bandwidth inside an accelerator (intra-accelerator) that across accelerators, and the bandwidth inside a node (intra-node) is faster than that across nodes (inter-node). Figure Ia visualizes the hierarchical structure of communication topology with a x matrix, where different levels of communication are in different colors. We view the data flow in
faster than
Work was done while the author was interning GTCOM. " *Corresponding author.

---

Published as a conference paper at ICLR 2023
Pestere
Node
Acvelnater?
Sepasee
Node
Note
Comnasicatk
leker aule
Acaloratec
Slidng Sequence
Projetioe
(a)
Figure 1: (a): The all-to-all communication contains three levels of communication with different bandwidth: the fast intra-accelerator communication inside each accelerator (green squares), the intra-node communication between Accelerator 0, and Accelerator 3, (blue squares), and the slow inter-node communication between Node O and Node (orange squares). (b): Slicing data on the sequence dimension (left) and Projecting data on the feature dimension (right) into three groups corresponding to three levels of communication. Each row of the data is token embedding.
the all-to-all communication from the perspective two dimensions: sequence dimension (tokens) and feature dimension (embeddings of tokens). The proposed SCoMoFE transforms (slicing or pro- jecting) the data flow at cither the sequence or feature dimension into three communication groups: antre-accel,rator, intra-node and global (inter-node) communication, shown in Figure Ib For the data slicing on the sequence dimension, we select tokens for a communication group according to the assignment scores of the tokens with the experts inside the group. To organize the data transfor- mation on the feature dimension, we linearly project features into the three communication groups with lower feature dimensionality and recast them back after the all-to-all communication.
Theoretically, structuring all-to-all communication in this way is faster than its original form, since less data transmitted through the slow inter-node communication channel. However, this may hurt the performance, because intra-acceelerator/itte--- part of experts. To alleviate this issue, further propose token clustering approach ag- gregating related tokens from different devices, which elevates the association of tokens inside each device faccelerator/node). The proposed token clustering uses the balance router presented by Lewis) et al. (2021) for clustering, where each device is cluster. In the balanced router, sigmoid gate is adopted to combine the inputs and outputs of experts, which could only broadcast the gradients to the activated experts, even though it may be more suitable to dispatch the tokens to other experts. Hence we propose to replace the sigmoid gate with the softmax gate via a trrighhtthrouug trick (Bengio et al. 2013) for better gradient broadcasting.
communication can be processed by only
In nutshell, our contributions are summarized as follows:
1. We propose SCoMoE that transforms the data flow in the all-to-all communication into three groups according to the bandwidth structure of communication topology.
2. A token clustering method is proposed to dispatch the related tokens to the same devices to alleviate the routing locality in the structured communication.
3. We propose to use softmax gate to substitue the sigmoid gate in the balanced router for better gradient broadcasting.
4. Experiments bilingual and massively multilingual machine translation demonstrate that SCoMoE is faster than Gshard (Lepikhin et al. 2020) significantly with comparable or even better translation performance. Further analysis discloses the strategies of selecting hyper-parameters for SCoMoE.
RELATED WORK
MoE models Jacobs et al. 1991 Jordan & Jacobs 1994) are ensemble methods to integrate multiple
experts. Shazeer et al. (2017) propose gating network select combination of experts and mix data parallelism and model parallelism to increase the batch size. Gshard (Lepikhin et al.) (020) utilizes the MoE parallelism proposed by Shazeer et al. (017) scale Transformer by replacing
2

---

Published as a conference paper at ICLR 2023
the MLP layer of Transformer with multiple MLP experts. Switch Transformer (Fedus et al. 2021
substitutes the top-2 routing of Gshard with the top-1 routing, and trains large sparse models with a selective precision strategy. Deepspeed-Me Moo (Kim et al., (021) combines ZerO RRajbaandari al.
2020) and MoE parallelism to train large Moł models. Many recent large language models (Fedus
2027 Zoph et al. (0022) have been trained with Gshard and its variants. A variety of new routers, c.g., hash router (Roller et al. 202
al. 2021 Zhang et al. 2021b Ma et al., [022] Du et al.]
random router (Zuo al. (021) and balanced router (Lewis et al., 202T). have been proposed lieu
of the top-2 router in GShard. We employ the balanced router proposed by the BASE Layer (Lewis
al.] 2021) for token clustering. The balanced router uses the auction algorithm (Bertsekas 1992)
for balanced routing.
Although sparse MoEs are computationally efficient, they suffer from frequent and expensive all-to- all communication across the network. Gshard (Lepikhin et al. (020) has shown that the all-to-all communication becomes the bottleneck for training large MoEs. While the heterogeneous hierar- chical network topology has been exploited to reduce the communication consumption, previous studies (Lin et al. 2020] Castiglia et al. T02T) I focus on the all-reduce communication in SGD rather than the all-to-all communication in MoE. Tutel (Hwang et al., (022) proposes hierarchical all-to-all communication approach for MoE, which performs all-to-all communication and layout transformation inside node, and then all-to-all communication across nodes. Although the hier- archical all-to-all communication of Tutel reduces the number of communication hops, it increases the amount of communication data. Hetu-MoE (Nie et al., (022) introduces a hierarchical all-to-all communication mechanism similar to that of Tutel. The difference is that Hetu-MoE dispatches all data inside one node one accelerator, which may cause memory overflow. Sparse-MLP (Lou et al. 2022) shortens the communication time by eliminating 10% tokens before routing, accord- ing to the importance scores between tokens and experts. However, removing tokens may hurt the performance. Fast-MoE (He et al. 2021) optimizes the training speed of MoE on the computation of experts, rather than communication. Task-level routing (Zhou et al., 2022) extracts a subset of experts for inference which improves the inference efficiency, but only works for multilingual/multi-
task MoEs and the extracted submodel can only translate in one direction. Expert Choice Routing Kudugunta et al. 2021) speedups the convergence speed of MoE with better load balance, but cannot reduce the time cost for each step.
Comparison with Tutel Although both SCoMoE and Tutel improve the all-to-all communica- tion of MoE based on the hierarchical network topology, SCoMoE is significantly different from Tutel in the following four aspects: 1) Motivation. Tutel attempts to decrease the communication hops and increase the message size for inter-node communication while SCoMoE is proposed to decrease the amount of data for inter-node communication and increase that for intra-node com- munication. 2) Algorithm. The communication of Tutel is hierarchical, first intra-node, and then inter-node. However, the communication of SCoMoE is structured, where different data are com- municated in different groups
3) Translation Performance. The hierarchical all-to-all communication algorithm proposed by Tutel may improve the communi- cation speed, but does not influence the translation performance. By contrast, SCoMoE changes the routing of some data, which may serve as regularization and improve the translation performance. 4) Speedup. In certain situations, Tutel even slows down training while SCoMoE still improves training speed. The comparison of the two methods on the training speed are shown in Figure [3d] Despite these significant differences, SCoMoE is orthogonal to Tutel and the two methods can be easily combined.
Massively Multilingual NMT Our work is also related to massively multilingual neural machine translation, which is briefly reviewed in Appendix due the constraint of space.
3 MODEL
3.1 SPARSE MIXTURE-OT-EXPERTY AND ALL-TO-ALL COMMINICATION
Sparse MoE models usually extend Transformer architecture with MoE layers where only combi- nation of experts are activated for computation each time. Different experts do not share parameters, as shown in Figure 2 (in blue), while other modules (in red) are shared across all devices.
3

---

Published as a conference paper at ICLR 2023
Given input sequence X which contains tokens {X1, ...,... Xi..., X,) with token embed- dings, sparse MoE models assign each token X; to experts usually according to a token-to-expert assignment matrix M:
M XWE
My WE, X;
WE is a projection matrix, where each column is corresponding an expert embedding. Mej is the assignment score between an expert embedding WE, and token embedding X,, Given the token-to-expert assignment matrix M, the top-2 gating router selects the top-2 experts for each input token, while the balanced router assigns tokens to experts with the auction algorithm.
After the assignment, tokens are sent to their target experts via the all-to-all communication across devices, where each device sends/receives (i.e., dispatches and combines) tokens to/from all other devices. Experts in devices process the received tokens, and send the processed tokens back to their original devices, which requires another all-to-all communication. The all-to-all communication contains three levels of communication: intra-accelerator,, intra-node, and inter-node, visualized in green, blue and orange respectively in Figure Ta The intra-accelerator communication is actually copying data from the memory inside the corresponding accelerator, which is very fast. The intra- node communication is the communication across different accelerators inside one node, which much faster than the inter-node communication across different nodes. However, most of the communication in the all-to-all communication are the slowest inter-node communication, which visualized in orange in Figure Ta Furthermore, the ratio of inter-node communication increases quadratically with the number of nodes. Therefore, the all-to-all communication becomes very expensive and hence the bottleneck for training large MoE models with many devices.
3.2 STRUCTURED ALL-TO-ALL COMMUNICATION
To reduce the all-to-all communication cost, we propose to transform the tokens for cross-device communication into three groups corresponding to the three levels of communication: intra- accelerator, intra-node and global communication| The key idea is increase the amount of tokens dispatched/combined in the fast intra-accelerator/nndd communication channels, and decrease the amount of tokens in the slow inter-node communication channel. Intuitively, this structured all-to- all communication acts like the multi-head attention where the inputs are projected into multiple keys, queries and values before attention computation in that our structured communication trans- forms the data into multiple groups for communication at different levels before expert computation.
Different from the vanilla MoE, we do not compute the token-to-expee expert assigmment scores for all experts. Instead, we aim to select the experts inside each accelerator for the intra-accelerator com- munication, and the experts inside each node for the intra-node communication. As shown in Figure 2 the tokens in the intra-acceleratorrnode group only be processed by the experts inside the accelerator/node, while the tokens in the global group are processed by all experts. The processed tokens are sent back to their original devices.
3.3 DATA TRANSFORMATION FOR STRUCTURED COMMUNICATION
The data transformation for the structured communication could be either on the sequence or feature dimension. The difference of them is visualized in Figure lb
Data Slicing on the Sequence Dimension For the data transformation on the sequence dimension, we slice the tokens into each communication group according the token-to-expert assignment scores introduced in the Section 3.1 The tokens to be sent to intra-acceleraatorrndde devices have the highest assignment score with the experts inside the devices. Specifically, the tokens for local devices are collected via:
X = arg TopK
Σ Mi
JEENee
where Mij is the assignment score between expert j and token i, Exucal is set of experts distributed local devices, Xxcaal of tokens that have the top-k assignment scores with Ejocal- k denotes
We do perform data transformation for inter-node communication.
4

---

Published as a conference paper at ICLR 2023
Intra-actterater Communiation
Intra-osts Communicatiin
Global Commankattoo
Strecture Commiicttion en Sequesce Dinrcesion
Stractured Communicaiio Peature Dimession
Recever
SCaMaf:1
Nade
Mareraate Export
Aeselraaee Expurt
Accipaater
Anderaane Equn
Anserrate Expen
Actliraie
Eipert
Expert
Layer
Clastrlng
N/2
Layer
Projection
Token Selection
Transpic
faaurr
N/2
(SelCCros) Attenion
feature
Figure 2: Architecture of the SCoMoE. Left: the overall structure of the model. Right: the detailed structure of a SCoMoE layer. There are 2NV layers the model, N attention layers, N shared 2
FNN layers and N MoE layers (token clustering with the 2 succeeding FFN layers can be seen as one MoE layer). 2 'The parameters of layers blue are shared across different devices, while parameters of layers in red are expert parameters. The token clustering aggregates related tokens at different devices together before structured communication layers. The aggregated tokens are sent back to the original devices (Recover) after the SCoMoE layers. The SCoMoE layers transform the data into three groups on either the sequence or feature dimension, and dispatch the transformed data to the experts via the intra-accelerator, intra-node and global communication respectively. We omit the process combining the processed tokens into their original devices in the diagram for simplicity.
the number of tokens to be dispatched to local devices, which is hyper-parameter. The tokens for both the intra-accelerator and intra-node communication are collected in this way. We primarily select token for local communication, first intra-accelerator, then intra-node and finally for global communication. We mask the token-to-expert assignment scores of the selected tokens with -oo to prevent the same tokens from being selected for multiple groups of communication.
Data Projection the Feature Dimension To transform the data the feature dimension, project the input sequence X with the dimensionality of dmoatel to lower-dimensional representations of denra accelerator; deura mode and dglabal, with three different projections W inta accelerator. Wintra mode and
W slobba::
X intra accelerator
X
W intra accelerator
Intra-ndde node
X W inntr-nndde
(3)
global
XW gglbbal
dintra accelerator d mode
+ dglobal d
where X E dunnee
X intra accelerator E T dhamerccccaae
X
Xintra nade Paumoodee X global E R
X inttr-accellratorr Xintr-node- X globbal are projected sequences to be dispatched via different communi- cation channels. We display these projections and projected sequences in different colors in Figure
2
The dimensionality of data dispatched to experts is reduced into daura-seellerater dyurr-noode dglobal- Hence we also need to reduce the dimensionality of weights in experts correspondingly. As expert actually fully-connected network, which contains two linear projections Wi Rdasn dim. W2 Iddes x Specifically, we the shape of the first projection be (d, dm),
and that of the second projection to be (dins d). d E (dimnra acselerator, diata mode, dglobal) corresponding
We ignore the bias for simplicity.
5

---

Published as a conference paper at ICLR 2023
to the input dimension. Actually, the parameters of the experts are projected into three parts, each of which is corresponding to one level of communication. The outputs of experts are recast back to the original dimension (dmotel) with another three projections.
3.4 TOKEN CLUSTERING
The data slicing on the sequence dimension forces more tokens to be dispatched to local devices, which may hurt the performance, tokens dispatched to local devices processed by only a few experts inside the devices. To mitigate this issue, we propose to apply token clustering before the SCoMoE layers. The token clustering is aggregate related tokens the same device (accelera- tor/node), where each device is cluster. Aggregating related tokens together will make it better for the experts to process them. But the router of MoE usually assigns related tokens to the experts (Zoph et al. 2022). why do we need the token clustering? The reason is that the router of SCo- MoF. assigns more tokens to the experts inside local devices, while token clustering can mitigate this problem. It gathers the tokens that are related to the experts at the same devices together so that the tokens at the same devices may have a stronger association with each other and can be better processed by the same experts. We perform the token clustering before the first SCoMoE layer, and recover the original token distribution after the last SCoMoE layer.
We average the embeddings of experts inside accelerator/node the embedding of the device (cither the accelerator or node), and assign tokens to devices in the same way as we assign tokens to experts. But unfortunately, the commonly used top-2 gating (Lepikhin et al. 2020) and top-1 gating (Fedus et af.) 202T) routers could not be used here, since these routers abandon some tokens during routing. Therefore we resort to routers that keep all tokens, e.g, greedy router (Lewis al., 2021) and balanced router (Lewis et al. 2021). The greedy router allows tokens to be dispatched to the top-1 expert according the token-to-expert assignment matrix. However, different from the commonly used top-1 gating (Fedus et al. 202I), the greedy router does not set capacity for each expert, allowing any number of tokens to be sent the same expert. When the greedy router is used to route tokens to devices, it may send too many tokens to a single device as it is unbalanced, causing out-of-memory issue. Incorporating the load balance loss (Lepikhin et al.) 2020) does not help in this case, since the out-of-memory issue may immediately happen at the beginning of training. In contrast, the balanced router does not have the problem of the greedy router, which uses the auction algorithm (Bertsekas) 1992) to assign the same number of tokens each expert. But the auction algorithm usually converges until many iterations, which is inefficient for autoregressive decoding. Hence the BASE Layer (Lewis al. (02T) uses the balanced router instead of the greedy router at inference. In this work, we perform token clustering with the balanced router. However, we have found weakness of Base Layer in gradient broadcasting, and proposed Differentiable Sorting to address it, which is described in Appendix B
The token clustering aggregates the tokens of different sentences together, which cannot be input to the self-attention layer. Therefore we have to change the order of Transformer layers, putting the attention layers before the MoE layers. The new architecture is shown in Figure 2 where the shared encoder/decoder layers are placed at the bottom, followed by the attention layers and MoE layers. The design of the new architecture due to the following reasons: (1) Press et al. (2020) have in- vestigated the impact of the order of attention layers and feed-forward layers on the performance and found that putting the attention layers before the feed-forward layers often yields good perfor- mance. (2) Riquelme et al. (2021) find that the top MoE layers more important than the bottom MoE layers. (3) Intuitively, the training of sparse MoE layers is more unstable than the training of the dense layers (Zoph et al. 2022). Hence, putting the stable dense layers the bottom may help stabilize the training of the entire architecture.
We put the balanced router before the first SCoMoE layer for token clustering. However, the token clustering also requires the all-to-all communication. We remove one MoE. layer to avoid extra com- munication cost. Hence the number of SCoMoE layers is N/2- rather than N/2, as shown in Fig- u To keep the same computation cost with the original model, we also insert two fully-connected layers after the balanced router
[ The balanced router together with the two fully-connected layers
Two fully-connected layer rather than one, since we use the top-2 gating in SCoMoE layers, where each taken processed by two experts.
6

---

Published as a conference paper at ICLR 2023
act the first MoE layer. Since the token clustering works on the sequence dimension, we only apply the token clustering to the SCoMoE layer that slices the data on the sequence dimension.
EXPERIMENTS
4.1 EXPERIMENT SETTINGS
We evaluated the proposed SCoMoE on both bilingual and massively multilingual neural machine translation. In bilingual NMT experiments, we used the WMT17-En-Frr corpus, which consists of 35,762,532 training sentences, while the massively multilingual NMT experiments, the OPUS- 100 corpus (Zhang et al., 2020)
used, which consists of 100 languages and 107, 107,924,846 training examples. The details of these two datasets are shown in Table of Appendix C We used the script of fairseq to download and pre-process the W-T17-En-
dataset while the OPUS-100 dataset was directly downloaded from the OPUS web Both the WMT17-En-Fi and OPUS-100 dataset were tokenized via the BPE algorithm (Sennrich et al. 2016). For the multilingual experiments, language token indicating the source language was prepended to each source sentence and similarly, target language token was prepended each target sentence.
Models were trained in two settings: base model and large model. The base models have 6 en- coder/decoder layers, experts and dmodel 512, which were trained on one node with 8 RTX A6000 GPUs The large models have 12 encoder/decoder layers, 32 experts, and dmudell 1024, which were trained on nodes with 32 RTX A6000 GPUs. The nodes were connected with In- finiband, the bandwidth of which is 100 Gbps. The number of parameters in the base model is 0.2 billion while 4.1 billion the large model. The details of the base and large model configuration could be found in Table 3 of Appendix C
We compared our method with Gshard and its variant with a small capacity factor on translation performance and training speed. We refer to SCoMoE on the sequence dimension as SCoMoB- Seq, and that on the feature dimension as SCoMoE Featt. The Gshard with a small capacity factor is referred as Small-Cap-MoE. We set the ratio of tokens dispatched through the intra-rccelerrtor communication to be a (0 < a < 1), and the ratio through the intra-node communication to be (0 < < 1), hence that through the global communication is 7 1 To simplify the comparison, we set either a or B to be non-zero, and the other to be 0. We refer the SCoMooB-(eq/Feat) with a > 0, 0,7
3 (0 y 1).
as SCoMoo-(Seq//eat)-o while that with = 0, 3 > 7 >> 0 as SCoMoE- (SeMooF-SSe//Fat)222 Please refer to Appendix C for more details about experiment settings.
4.2 MAIN RESULTS
We evaluated the proposed methods on the WMT17-En-Fr and OPUS-100. We report the BLEU score on the WM.17-En-Fr. and the average BLEU scores over all, English-to-any, and any-to- English language pairs on the OPUS-100. Results are shown in Table I and the speedup of all models over Gshard.
We also report the win-rate
We evaluated the training speed on the dummy-mt task x of fairseq. since the training speed evaluated on is more stable. Each input sentence in the dummy-mt task is a series of numbers, c.g., "0, 1, 2, 3, 4, ...', rather than the real sentence. We fixed the sequence length of all source and target sentences to be 30, and trained models with 100 million dummy tokens without validation (It takes 920 steps train each model).
Base Model The base SCoMoE was trained on a single node, hence we only need to consider the ratio for the intra-accelerator communication (a). For fair comparison translation performance, we set a for SCoMoE-Seq to 0.3, and that for SCoMoE-Feat to 0.5, the capacity factor Small-Cap-
prepare-wwtt/ee22rr.hh
https://githbbb.cm///yttttttt//
mt. py
7

---

Published as a conference paper at ICLR 2023
Table 1: Translation performance of the base and large models on the of OPUS-100 and WMT17-En-Fi (CoMoE-(Seq/Feat)) the SCoMoE-(Seq/MTT) with the intra-accelerator and global communication. Time refers to the average time cost (ms) of all-to-all communication.
Multilingual
Setting
Model
Bilingual
speedup
Time
All
En-Any
Any-En
win-rate (%)
Transformer
23.43
19.64
27.17
26.2
37.39
2.00x
Gshard
24.41
20.02
28.76
37.96
1.00x
1.6
Base (0.2B)
CapaM-
23.02
19.08
26.92
20.3
37.27
1.10x
1.33
SCoMoE-Seq-o
24.45
20.67
28.18
43.3
38.24
1.10x
1.14
SCoMoE-Feen
23.47
20.25
26.67
27.3
37.51
1.07x
1.13
Gshard
26.78
23.11
30.44
40.09
1.00x
27.0
Small-Cap-MoE
26.93
24.58
29.28
49.5
39.87
1.20x
21.3
Large (4.1B)
SCoMoE-Seq-a
27.19
25.68
28.69
69.5
39.31
1.20x
20.3
SCoMoE-Feat-a
27.99
25.13
30.85
73.9
39.79
1.32x
17.4
SCoMoF-Seq-BB
29.60
27.32
31.39
90.4
40.20
1.25x
16.8
SCoMoF-Feat-t/
28.19
24.90
31,47
71.2
39.93
1.10x
19.9
MoE. to 0.8, so that these models can be trained with a comparable training speed. More evaluation on the training speed of SCoMoE can be found in Appendix
From Table [ we can observe that both SCoMoE-Seq/Feat and Small-Cap-MoE are faster than Gshard by nearly 10%. But SCoMoE cost less communication time than Small-Cap-MoE, Why do they still have similar training speed? The reason is that SCoMoE brings some extra over- head on computation, which is analyzed Appendix Although SCoMoE-Feat slightly slower than Small-Capp Caa---ap.MM,,, it is better in terms of BLEU on both the WMT17-En--F and OPUS-100 dataset. While SCoMoE-Seq is comparable or even better than Gshard on both the WMT17-En-FF and 100 dataset, due to the proposed token clustering and differentiable sorting. We further demonstrate the effectiveness of the two methods in Appendix F G
Large Model The large SCoMoE models were trained on multiple nodes, hence both SCoMoE- Seq-a and -3 were evaluated. We set 0.2 for SCoMoE-Seq-oo and 0.5 for SCoMoR-Feat-o, which have the comparable training speed with SCoMoE-Seq/Featt/ with / = 0.5. We set the capacity factor of Small-Cap-MoE to 1.0, while that of other models was set to 1.25. From Table it is clear that all variants of SCoMoE are faster and consume less communication time than Gshard significantly. Among the variants, SCoMoE-Feett-a is the fastest. It reasonable since we of "CoMoE-Feat--e to 0.5, but that for SCoMoE-Seq-a to only 0.3. We provide the training speed comparison of SCoMoE with more settings of and in Appendix E It can also be seen that SCoMoF-Seq-a is faster than SCoMoF-Seq- but takes more communication time. This is because that SCoMoE uses the top-2 routing, where each token computed in two experts. comparison, tokens for intra-accelerator communication are computed in single expert inside the accelerator, which reduces computational time.
We have also found that all variants of SCoMoE achieve better translation performance than both Gshard and its variant Capall-Ca--MoE on the OPUS- 100 dataset. They also perform comparably dataset. Among the 4 variants of SCoMoE, SCoMoE-Seq-3 is the
to Gshard the WMT17-Ens best in terms of BLEU on both OPUS- 100 and WMT17 MMMTI--En-Fr is surprising that SCoMoE-Seq-a outperforms Gshard by 2.8 BLEU on OPUS-100, but only 0. BLEU WMT17-En-Frr We conjec- ture that MoF models are easy to overfit on low-resource languages (Costa-jussa et al.) 2022). while our method can mitigate this issue by regularizing the routing of MoE. We have shown the improve- ments of SCoMoE on different languages in Appendix D to validate this hypothesis. Overall, all variants of SCoMoE are faster than Gshard with comparable better translation performance.
The best settings of a/3 for SCoMoE From Table I we can roughly see that SeoMooE-Seg/esss have better translation performance, while SCoMoE-
are faster. We further compared the translation performance of them with different settings of a or 3. The results are shown in Figure Ba and 3b] According to Figure Ba SCoMoE-Seq-a is better with smaller in terms of BLEU, which is not surprising since increasing a will reduce the amount data for global communica- tion. But SCoMoE-MoF-L achieves better translation performance with larger As. Furthermore,
8

---

Published as a conference paper at ICLR 2023
14
14
14.5
Speetie
385
Scatte
2.0
Cahard
20
314
02
Number
2a
20033 5000
15220
(a)
(b)
(c)
(d)
Figure 3: BLEU and Speedup of SCoMoE different settings. a/B with different a/3. (bb; Speedup of SCoMoE-Seq-a/$ with different a/B. (c Speedup
Bap: BLEU of SCoMoE-Seq-
of SCoMoE-Seq-a/b with different number experts. (3db: Speedup of SCoMoE-Seq-a with different batch sizes.
SCoMoF-MoF--B is even comparable to Gshard when 3 0.9, suggesting that only small propor- tion (10%) of tokens need to be communicated in global communication and that intra-node com- munication is sufficient for the majority of tokens. SCoMoE-Seq-$ with small Is underperforms Gshard, the reason for which can be found in Table where Gshard also suffers from performance drop with token clustering and reordering.
With regard to the comparison on the training speed (Figure 3b). the speedup SCoMoE-seq- o/B over Gshard increases nearly linearly with the increment of a/3. The rate of increment for SCoMoF-seq-a is around 1.0, while that for SCoMoE-seq-S is 0.5. Therefore, to achieve the best speedup without performance drop in BLEU, of SCoMoE-seq-B to 0.9, which can lead to a speedup of 1.28x and 1.44x respectively. We can also set 3 of SCoMoE-seq-B 0.5, which achieves the highest BLEU scores and a speedup of 1.25x.
can set of SCoMoE-seq-a to 0.3 and
Scaling in model and batch size The speedup of SCoMoE is also dependent the model size and batch size. Figure [c and 3d compares the training speed of SCoMoE with different number of experts and batch sizes. From Figure 3c it is clear that the speedup of SCoMoE increases with more experts, but the trend is not so significant when the number of experts increases to 48. It is due that the communication cost is more expensive with more devices. Similarly, SCoMoE also has more significant speedup with larger batch sizes, as shown in Figure 3d In Figure 3d we also display the training speed of Gshard with the hierarchical all-to-all communication proposed in Tutel (Hwang et al., 2022). However, the hierarchical communication of Tutel even slows down the training. The reason may be that the intra-node bandwidth our experiments is not fast (PCle) and that the number of devices is not very large (32).
Other Experiments and Analyses In addition to the main results, we conducted further experi- ments and analyses, which are presented in Appendix, including the analysis on why SCoMoE can improve translation performance (D), extensional evaluation on training and translation speed of SCoMoE [ ablation study H tuning a, 3, simultaneously (l) and performance vs. time
evaluation K
CONCLUSION
In this paper, we have presented SCoMoE reduce the all-to-all communication time of MoEs, by encouraging more data to dispatch through the fast intra-aacelerrtorrnndde communication chan- nel, and less data through the slow inter-node communication channel. By slicing/projecting the data on the sequence/feature dimension, we transform the data into three groups corresponding to three levels of communication. To alleviate the performance drop caused by routing locality, we further propose the token clustering to aggregate related tokens from different devices, and replace the sigmoid gate of the balanced router with the softmax gate through differentiable sorting. Ex- periments bilingual and massively multilingual neural machine translation demonstrate that our SCoMoE speeds up Gshard significantly with comparable or even better performance. The empiri- cal evaluation and theoretical analysis the training speed and communication time cost discloses the speedups achieved by our method under different setting and network bandwidth. The ablation study demonstrates the benefit of the proposed token clustering and differentiable sorting.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
The present research was supported by the Key Research and Development Program of Yunnan Province (Grant No. 202203AA080009 and the joint research center between GTCOM and Tianjin University. We would like to thank the anonymous reviewers for their insightful comments.
REFERENCES
Roee Aharoni, Melvin Johnson, and Orhan Firat. Massively multilingual neural machine translation. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies. NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume I (Long and Short Papers), pp. 3874-3884, Association for Computational Linguistics. 2019. doi: 10.18653/ URL
v1/n19-1388 1388.
http:///060o.rrr111111633////111 1ttp///1.11113///11111n1 1388
org
Naveen Arivazhagan, Ankur Bapna, Orhan Firat, Dmitry Lepikhin. Melvin Johnson, Maxim Krikun, Mia Xu Chen, Yuan Cao, George F. Foster, Colin Cherry, Wolfgang Macherey, Zhifeng Chen, and Yonghui Wu. Massively multilingual neural machine translation in the wild: Findings and challenges. CoRR, abs/1907.05019 2019. URL http::////aa org. abs 1907 1907 05019
Ankur Bapna and Orhan Firat. Simple, scalable adaptation for neural machine translation. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pp. 1538-1548. Association for Computational Linguistics, 2019. doi: 10. 1L.18653/1/DD1--11 URL
1165
Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. CoRR, abs/1308.3432, 2013. URL
Dimitri P. Bertsekas. Auction algorithms for network flow problems: tutorial introduction. Com- put. Optim. Appl., 1(1):7-66, 1992. doi: 10. 10.1007/BF00247653. URL https://dol. 10. D5.100//FF002477
Timothy Castiglia, Anirban Das, and Stacy Patterson. Multi-level local SGD: distributed SGD for heterogeneous hierarchical networks. In 9th International Conference on Learning Repre- sentations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
https://openneviee net. www.:/ee/foruiitiiomm fww.mmiddccceppppae.ooo
Marta R. Costa-jussà, James Cross, Onur Çelebi, Maha Elbayad, Kenneth Heafield, Kevin Heffer- nan, Elahe Kalbassi, Janice Lam, Daniel Licht, Jean Maillard, Anna Sun, Skyler Wang. Guil- laume Wenzek, AI Youngblood, Bapi Akula, Lote Barrault, Gabriel Mejia Gonzalez, Prangthip Hansanti, John Hoffman, Semarley Jarrett, Kaushik Ram Sadagopan, Dirk Rowe, Shannon Spruit, Chau Tran, Pierre Andrews, Necip Fazil Ayan, Shruti Bhosale, Sergey Edunov, Angela Fan, Cynthia Gao, Vedanuj Goswami, Francisco Guzmán, Philipp Koehn, Alexandre Mourachko, Christophe Ropers, Safiyyah Saleem, Holger Schwenk, and Jeff Wang. No language left behind: Scaling human-centered machine translation. CoRR, abs/2207.04672. 2022. doi: 00.48550/arXiv22200 04672. UtL/httss: /ow. wrw.:...org//// org/ /0c 205.44/00.....2200000 wwtpv..iiiiv.....
Nan Du Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang. Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathy Meier-HHelsster, Toju Duke, Lucas Dixon, Kun Zhang, Quoe V. Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. Glam: Efficient scaling of language models with mixture-of-experts. CoRR, abs/21 12.06905, 2021. URL httpp: //axxivv//rvvvvi// org/ abs 2tt.//2122.00000
Angela Fan, Shruti Bhosale, Holger Schwenk, Zhiyi Ma, Ahmed El-Kishky, Siddharth Goyal, Mandeep Baines, Onur Celebi, Guillaume Wenzek, Vishrav Chaudhary, Naman Goyal, Tom Birch, Vitaliy Liptchinsky. Sergey Edumov, Michael Auli, and Armand Joulin. Beyond english- centric multilingual machine translation. J. Mach. Learn Res., 22:107:1-107:48, 2021. URL http:://im wworg//apers/v///2222 10000013-070.00 1307.html
10

---

Published as a conference paper at ICLR 2023
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion param- eter models with simple and efficient sparsity. CoRR, abs/2101.03961, 2021. URL https:
//axxiv.crg/aaaaa http://b//21000 03961
Thanh-Le Ha, Jan Niehues, and Alex Waibel. Toward multilingual neural machine translation with universal encoder and decoder. In Proceedings of the 13th International Conference on Spoken Language Trunslation, 2WS1, Seattle, WA, USA, December 8-9, 2016. International Work- shop on Spoken Language Translation, 2016. URL https:: ac antho loy.. org.2015
iws 1.
Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, and Jie Tang. Fastmoe: A fast mixture-of-expert training system. CoRR, 1122100.11262, 2021. URL https://arxiv:
Org/abs/21035 13262
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Effi cient training of giant neural networks using pipeline parallelism. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buse Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference
Neural Informa- tion Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 103- 112, 2019. URL 093f65e080a295f (ww000/666665100 25935666666235501162222222222263 html
nwwrieiripp.. ips. ccc/paper/2019/hast. Abst
Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang. Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, Joe Chau, Peng Cheng, Fan Yang, Mao Yang. and Yongqiang Xiong. Tutel: Adaptive mixture-of-experts at scale. CoRR, abs/2206.03382,, June 2022. URL https:
03382 p
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mix- tures of local experts. Neural Comput., 3(1):79-87, 1991. doi: (ttl////ecoo.991...... URL
https://doi.orgg
www.l.1622nnceec 1991.
Renren Jin and Deyi Xiong. Informative language representation learning for massively multi- lingual neural machine translation. In Proceedings of the 29th International Conference on Computational Linguistics, pp. 5158-5174, Gyeongju, Republic of Korea, October 2022. Inter- national Committee on Computational Linguistics. URL httts: org
2022.coling-1.o 458
Melvin Johnson, Mike Schuster, Quoc V. Le, Maxim Krikun, Yonghui Wu, Zhifeng Chen, Nikhil Thorat, Fernanda B. Viégas, Martin Wattenberg, Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google's ss multilingual neural machine translation system: Enabling zero-shot transla- tion. Trans. Assoc. Compul. Linguistics, 5:339-351 2017. doi: 10.1 10.1162/tacl) a) 00065. URL
https://doi /0rg/1.... wwwl1116tcccccc _a00000000___0055____
Michael I. Jordan and Robert A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Comput., 6(2)18112144 1994. doi: 19tp://ecco.199...2222 URL httpp://doi.
Young Jin Kim, Ammar Ahmad Awan, Alexandre Muzio, Andrés Felipe Cruz-Salinas, Liyang Lu, Amr Hendy, Samyam Rajbhandari, Yuxiong He, and Hany Hassan Awadalla. Scalable and ficient moe training for multitask multilingual models. CoRR, 0bs/2109.10465 2021. URL
atttps://axxe wwepxxi..r////. abs 1049.10465
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua
Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015a. htth: //arxiv.orq// abs / 1412. 6980
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015b. URL http:: 7/arxiva wwwpsivv.rgaabaaaaaaa
6980
11

---

Published as a conference paper at ICLR 2023
Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Thang Luong, and Orhan Firat. Beyond distillation: Task-level mixture-of-experts for efficient in- ference. In Marie-Franee Francine Moens, Xuanjing Huang. Lucia Specia, and Scott Wen-tau Yih (eds.), Findings of the Association for Computational Linguistics: EMNIPP 2021, Virtual Event / Punta Cana, Dominican Republic, 16-20 November, 2021, pp- 3577-3599. Association for Computational Linguistics, 2021, doi: 10018663/11/222..innniii cm..p.304. URL https:
7/doi. //doi.org/10 10. 1100://1..1155511122222 http://ww/222..... findinggggnns 2eweddingg-een..22mm 304
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. CoRR, Jbs/2006.16668, 2020. URL
org/abs/ 16668
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. BASE layers: Simplifying training of large, sparse models. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML, 202 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 6265-6274, PMLR, 2021. URL http::/procceeeeeeeee
m.r. wtt.://mm..ppp.s/seeee.. perres//113//ewiie
html
Junyang Lin, An Yang. Jinze Bai, Chang Zhou, Le Jiang, Xianyan Jia, Ang Wang, Jie Zhang, Yong Li, Wei Lin, Jingren Zhou, and Hongxia Yang. M6-10T: A sharing-delinking paradigm for efficient multi-trillion parameter pretraining. CoRR, abs/21 10.03888, 2021. URL https:
7/ 7/arxiv. abs 2210. 03888
Tao Lin, Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use local SGD. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL UtL/htttps://ppenrrr///
net/forum?id-
Yuxuan Lou, Fuzhao Xue, Zangwei Zheng, and Yang You. Cross-token modeling with conditional computation, 2022.
Zixuan Ma, Jiaao He, Jiezhong Qiu, Huanqi Cao, Yuanwei Wang, Zhenbo Sun, Liyan Zheng, Haojie Wang, Shizhi Tang, Tianyu Zheng. Junyang Lin, Guanyu Feng, Zeqiang Huang, Jie Gao, Aohan Zeng. Jianwei Zhang, Runxin Zhong, Tianhui Shi, Sha Liu, Weimin Zheng, Jie Tang, Hongxia Yang. Xin Liu, Jidong Zhai, and Wenguang Chen. Bagualu: targeting brain scale pretrained mod- els with over 37 million cores. In Jaejin Lee, Kunal Agrawal, and Michael F. Spear (eds.), PPoPP 22: 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, Seoul, Republic of Korea, April 2 - 6, 2022, pp. 192-204. ACM, 2022. doi: 10.1 10.1145/3503221.355541 3508417
URL https: // org 10. 1115/11033322
Xiaonan Nie, Pinxue Zhao, Xupeng Miao, and Bin Cui. Hetumoe: An efficient trillion-scale mixture-of-exper distributed training system. ArXiv, abs/2203.14685 2022.
Matt Post. A call for clarity in reporting BLEU scores. In Proceedings of the Third Con- ference on Machine Translation: Research Papers, pp. 186-191, Brussels, Belgium, Octo- ber 2018. Association for Computational Linguistics. doi: 10.18655/w//111-631 URL
ltttpss//cllanthlooyccoo org. htt.:oorgg/w111111 /W11-6119
Ofir Press, Noah A. Smith, and Omer Levy. Improving transformer models by reordering their sublayers. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault (eds.), Pro- ceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pp. 2996-3005, Association for Computational Linguisties, 2020. doi: 10. 10.186553w1/2020.acc ac0.200..cmmma..2707 URL httts://ool..cr//..B855///////
acl acl-main. 270
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations
toward training trillion parameter models. In Christine Cuicchi, Irene Qualters, and William T.
Kramer (eds.), Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2020, Virtual Event / Atlanta, Georgia, USA, November 9-19, 2020, pp. 20. IEEE/ACM, 2020. doi: 111.9////444455222000000 7/doi. org/10.. 110.///41110///114000000 00000////4440 2020. 00024
URL hetps:
12

---

Published as a conference paper at ICLR 2023
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Am- mar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-man moo: Advancing mixture-of- inference and training to power next next-generation AI scale. CoRR, abs/2201.0559% 2022. URL
ht.tps://arxxt org abss///////// 2201 0
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Su- sano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Anmual Conference
Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. URL
8583 8595, 2021.
neurips......................... wttp:/ppppppppeeeee pttpppppppppcc11/////p/ htm
http:/221//// 2021/hash.///////p /hash/
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances Neural Information Processing Systems 34: Annual Conference
Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 17555- 2021. URL httpttpppppppcceeiiggg..... neurips cC/ .html
paper/ 202rr2221/ htthh/hhh////// 92bf f 6 407377 6ea Saw.aaaae3666666
Rico Sennrich, Barry Haddow. and Alexandra Birch. Neural machine translation of rare words with subword units. Proceedings of the 54th Annual Meeting of the Association for Com- putational Linguistics, ACL 2016, August 7-12, 2016, Berlin, Germany, Volume I: Long Pa- doi: 10.18653/11/16--1122 URL
pers. The Association for Computer Linguistics, 2016.
https:
wtrg/10.11655)
Noam M. Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. ArXiv, abs/1701.06533, 2017.
Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. CoRR, Abs/1505.0038/ 2015. |ttpp:::://axxxvvooo
aww./105. 00387
Haoran Sun and Deyi Xiong. Language branch gated multilingual neural machine translation. In Proceedings of the 29th International Conference on Computational Linguistics, pp. 5046-5053, Gyeongju, Republic of Korea, October 2022. International Committee on Computational Linguis- tics. URL http::/ss///laaaahhllaoooooo wag.annhhhoooyggo...o... wgw../g/yy.......o....// 2022. cow.nn-1..111100
Yining Wang, Jiajun Zhang, Feifei Zhai, Jingfang Xu, and Chengqing Zong. Three strategies to improve one-to-many multilingual translation. In Ellen Riloff, David Chiang, Julia Hock- enmaier, and ichi Tsuji (eds.), Proceedings of the 2018 Conference on Empirical Meth- ods in Natural Language Processing, Brussels, Belgium. October 31 November 4, 2018, pp. 2955 2960. Association for Computational Linguistics, 2018. doi: 10.18653/v1/d18.. 1326. URL https: orgg 10. 186336 v/33vvvv/ll --31136)
Yining Wang, Long Zhou, Jiajun Zhang, Feifei Zhai, Jingfang Xu, and Chengqing Zong. A compact and language-sensitive multilingual translation method. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1213-1223, Florence, Italy, July 2019. Association for Computational Linguisties. doi: 10.18653/11/11--1117 URL https: htt////aannntol....rrroo/oo/ 1117
Hongfei Xu, Qiuhui Liu, Josef van Genabith, and Deyi Xiong. Modeling task-aware MIMO cardi- nality for efficient multilingual neural machine translation. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing, ACI/IJCNIPP 2021, 1, (Volume 2: Short Papers), Virtual Event, August 1-6, 2021, pp- 361-367. Association for Computational Linguistics. 2021. doi: 1/186653/1122221acllleeee URL oogg:::://oo..ogg/1..1. wtt.://o0.....1///////1/13
acl ath.o//t short.hs
An Yang, Junyang Lin, Rui Men, Chang Zhou, Le Jiang. Xianyan Jia, Ang Wang. Jie Zhang, Jiamang Wang, Yong Li, Di Zhang, Wei Lin, Lin Qu, Jingren Zhou, and Hongxia Yang. Ex- ploring sparse expert models and beyond. CoRR, abs/2105.15082. 2021. URL https: warxiv.org/aba.. 15082
13

---

Published as a conference paper at ICLR 2023
Biao Zhang, Philip Williams, Ivan Titov, and Rico Sennrich. Improving massively multilingual neural machine translation and zero-shot translation. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Associ- ation for Computational Linguistics, ACI. 2020, Online, July 5-10, 2020, pp. (628-1639. As- sociation for Computational Linguistics, 2020. doi: main.
148. URL
https://doi.org: http://...1653/1//22000222
Biao Zhang, Ankur Bapna, Rico Semnrich, and Orhan Firat. Share or not? learning to sched- ule language-specific capacity for multilingual translation. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, net wtt.:/oooonnd-w10000000000
2021a. URL
Zhengyan Zhang, Yuxian Gu, Xu Han, Shengqi Chen, Chaojun Xiao, Zhenbo Sun, Yuan Yao, Fanchao Qi, Jian Guan, Pei Ke, Yanzheng Cai, Guoyang Zeng. Zhixing Tan, Zhiyuan Liu, Minlie Huang, Wentao Han, Yang Liu, Xiaoyan Zhu, and Maosong Sun. CPM-2: large-scale cost-effective pre-trained language models. CoRR, abs/2106.10715, 2021b. URL https:
//arxiv.orq/ abs/ 2106. 10715
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Y. Zhao, Andrew M. Dai, Zhifeng Chen, Quoc Le, and James Laudon. Mixture-of-experts with expert choice routing. CoRR, abs/2202.09368, 2022. httthhtps//aaat /ttppii///axiv..../ wttp:///gg//aaaaa/////.osaa http:///22.099.... 09368
Yaoming Zhu, Jiangtao Feng, Chengqi Zhao, Mingxuan Wang, and Lei Li. cnnterr-nterfference adapter for multilingual machine translation. In Aranc-Franc Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), Findings of the Association for Computa- tional Linguistics: EMNLP 2021, Virtual Event Punta Cana, Dominican Republic, 16- 20 November, 2021, pp. 2812-2823, Association for Computational Linguisties, 2021. doi: URL https: http::://ool.org/0...855//1//// 111.//0..11633/1122222222 www./0...ooo..ooooo..
1 www.diinggeennl..24444 emn emm.p.g2-emmpp2240
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. Designing effective sparse expert models. CoRR, abs/2202.08906, 2022. URL
https://arxe www.or////ooo 2202 08906
Simiao Zuo Xiaodong Liu, Jian Jiao, Young Jin Kim, Hany Hassan, Ruofei Zhang, Tuo Zhao, and Jianfeng Gao. Taming sparsely activated transformer with stochastic experts. CoRR, abs/2110.0426e 2021. httptttsss///iiiii.p//pa http:///ww.....////////100 abs/2/10. 04260
APPENDIX
A RELATED WORK ON MASSIVELY MULTILINGUAL MACHINE TRANSLATION
Recent years have witnessed an upsurge of interest in multilingual neural machine translation (MNMT) that trains a single NMT model to translate multiple languages (Johnson al. 2017: Ha
et al. 2016, Aharoni et al. 2019 Arivazhagan et al.) 2019). Efforts have been devoted to pushing the limits of MNMT in terms of the number of languages supported by adding language-specifff pa-
rameters (Wang et al. 201 8 2019
Zhang et al. 2020 Jin & Xiong 2022 Xu et al. (022). adapters
2021 Zhang al. 2021aa Sun Xiong 2022). increasing the depth (Huang et al. 2019) and width Fan et al. (021) of the model. In addition to these, MoE
(Bapna & Firat| 2019 Zhu al.
has become a popular architecture
scale MNMT to massively MNMT (Lepikhin al., 2020
Kim et al.] [02T]] Costa-jussà et al., 2022). In this work, we also conduct experiments on massively MNMT show the effectiveness SCoMoF.
B DIFFERENTIABLE SORTING
The auction algorithm in the balanced router is not differentiable, hence the BASE Layer (Lewis) ct al. 202I learns the expert embedding via the sigmoid gate (Srivastava et al. 2015):
X = (Mi)) expert, (X,) (1 a (M.,)) X;
14

---

Published as a conference paper at ICLR 2023
Anciion Algeritm
Vaciium
Expert
Expert
W
Xast
n/m
W Weottmax
detach (W sottmox)
Sol max
M
Breodcast
12/m
Figure 4: The visualization of differentiable sorting. The sorting projection W son from the auction algorithm, and the broadcast probability matrix W Solmmux are combined together via the straight- through trick to sort the sequence. The sorted four tokens are divided into two groups, per corre- sponding to one expert.
where the output of the MoE layer X is the combination of the expert input X; and expert out- put expert, (X;) that are integrated with the normalized assignment score Mij. The gradients are broadcast through the sigmoid gate to update the parameters of the expert embeddings. But only the embeddings of the activated experts can receive the gradients, even though we have computed the assignment scores between any tokens and experts. With the gradients from the sigmoid gate, only know how to update the assignment scores of the activated experts, which is not available for other experts. A straightforward solution to this is to replace the sigmoid gate with the softmax gate. However, the softmax operator requires multiple inputs, but only one expert is activated for each token. We notice that the auction algorithm assigns tokens to experts via sorting, which could be implemented as a linear projection with a sparse matrix, where each row and column only contains one non-zero value. Each row of the matrix could be seen as selecting a token from the sequence for an expert. We implement the sorting via linear projection:
(5) where the sorted sequence Xon obtained by the projection W.n on the input sequence X. If the number of sorted tokens is n, and the number of experts is n/m tokens will be assigned to each expert. We visualize the sorting with a x 4 sparse projection on a sequence of length 4 in Figure 4 assigned two experts. The assigned tokens often have the
X
Wson
The tokens in the sorted sequence max or top-k scores with the expert, therefore each row could be seen as the argmax or argtopk on the token-to-expert assignment matrix. The softmax is the approximation the temperature to a small value. Therefore, we share the gradients of the projection matrix with the softmax matrix via the straight-through trick (Bengio et al. 2013):
the argmax, if we set
W sort W Waeetion
W soffmmx detach W softmax
M T
(6)
W sofimax mof
W sor is the combination of the sorting projection W Wanction from the auction algorithm, W sootmaxx which is the normalized assignment matrix, and detach( softmax), which has the same value as W Wostmmxx but cannot broadcast the gradients. is the temperature which controls the sharpness of the softmax distribution. The lower the temperature is, the better the softmax approximates the argmax. Wsen seet equal to W Wasction in the forward computation. W seftmax has the same gradients W son in the backward propagation, although W sofimux does not have influence on the training loss.
As the shape of the probability matrix W Wanttmax different from that W snnss (m, n) vs. (n, m)) we cannot add them directly. Each row of W Wantumax is set of probabilities assigning tokens to an expert, which are corresponding to n/m rows of Wson as we need to assign n/m m tokens to each expert according to the same distribution. For the example in Figure both the first and second row of sparse matrix Wasetion are corresponding to the first row the probability matrix Wottmax- We reshape W sottmax from (m, n) (n, by repeating each row of the matrix to n m rows. In Figure we reshape the matrix from (2,4) to (4.4), by repeating the first and second row of the matrix Wsofimax to the first and second two of the broadcast matrix.
15

---

Published as a conference paper at ICLR 2023
C
DETAILS ON EXPERIMENT SETTINGS
Table The statistics of the WMT17-En-Fr and OPUS- 100 dataset.
Datasets
Languages
Train
Validation
Test
WMT17-En-Fr En-Fr
35,762,532
26,854
3003
OPUS-100
100
107,924,846
188 1000
188
1000
Table 3:
The configuration of base and large models. Fnc/Dec-LLayer denotes the number of en- coder/decoder layers.
Setting
Enc-Layer
Dec-Layer
dmnntel
Experts
GPU
Nodes
Parameters (B)
Base 6
6
512
8
8
1
0.2
Large
12
12
1024
32
32
4
4.1
Figure 5: Data distribution of different languages in OPUS-100.
We used the MoE branch of fairseq to implement MoE models. We prevented the padding tokens from being sent to experts by masking their assignment scores. The capacity factor of the base model was of (Zoph et al. (022). The tokens with small assignment scores to an expert were dropped the expert was out-of-capacity. All models used pre-normalization found that the MoE models are unstable to train in the large setting, hence we reduced the gradient norm by gradient clipping (i.e., gradient norm was clipped be more than 0.5) and gradient accumulation (accumulating the gradients every 4 steps). The gradients of experts were divided by the square root of the number of accelerators, which also reduced the gradient norm and may improve the stabilization of training. We have also found that the MoE model trained in multilingual experiments could not perform well inference, and that increasing the nad-balance loss weight from 0.01 (default value) to 0.1 significantly improves the performance. Therefore, we set the load- balance loss weight to 0.1 in the multilingual experiments, and 0.01 the bilingual experiments.
1, while that of the large model was set to 1.25 following the suggestion
stabilize the training. We have
The Adam optimizer (Kngmaa & Ba 2015a with learning rate 5 0.9, 0.98, was used to optimize our models. We employed the searse-square-rooo learning schedule (Kingma & Ba 20156). with the number of warm-up step being to 4000. The batch size was set to be 4096 tokens. The dropout rate was 0.3 for bilingual experiments and 0.1 for multilingual experiments. Both the base and large models were trained for 10 epochs. The checkpoint with the best perfor- mance on the validation set was selected for testing. In the bilingual experiments, we evaluated the performance on the validation set with BLEU. the multilingual experiments, the perplexity (ppl) on the validation set was used for faster validation. The BLEU score was computed via sacrebleu (Post 2018
"The signature of the wttpb//ww..pooorrrr//aaaaa/eeeee
https: ptrps://ee //i thub. http:tcc.//ttttocc//etattaeeeeae
-k-uurrk htw.eeeuurrffssssmooooooooo #ssmooth.exp-tok 3a+versioon. 1.5.1.
denotes the language direction, e.g. en-fr
16

---

Published as a conference paper at ICLR 2023
Gshard
50
SCoMoE-Feat-p
SCoMoE-Seq-8
Perplexity
44
Gshard
SCoMoE-Feat-f
SCaMcE Seq
40
10000
20000
30000
40000
5000
10000 15000 Step
20000
25000
Step
(a)
(b)
Figure 6: Convergence curves of Gshard and SCoMoB on the validation set of OPUS-100 WMT17-En-FF (6b).
6a and
22 20000000000 2009
Figure 7: Improvements of SCoMoF-Seq- over Gshard in terms of BLEU on different languages
of OPUS-100. Blue/yellow indicate that the language is a rich-/lowrrsooree language.
D WHY SCOMoE IS BETTER ON MULTILINGUAL MACHINE TRANSLATION
In Table the improvement of SCoMoE-Sepe -Seu-B on multilingual OPUS-100 is much more significant than that on bilingual WMT17-En-Fi We plot the convergence curves of Gshard and SCoMoE on validation set of OPUS- 100 and WMT17-FEn-FF in Figure 6a and 6b respectively It is clear that the convergence curves of SCoMoF are more smooth than Gshard on both OPUS- 100 and WMT17-En- Fr. Gshard obtains its best perplexity OPUS-I00 at the 20kth step. but its perplexity increases as the training continues. While the convergence curve Gshard on WMT17-En-Fr is also fluctuated, it tends towards getting better. Therefore, Gshard only overfits on OPUS-100. This is reasonable, since WMT17-En-FF contains much more data than any single language pair in OPUS- 100 does. In Figure 5 we illustrate the data distribution of different language pairs in OPUS- 100. It is clear that only half of language pairs have about million parallel sentence pairs, while other language pairs have much less data. We also show the improvements of SCoMoE over Gshard different language pairs of DPUS-100 in Figure 7 where we visualize different language pairs in different colors according whether the language is low-resource or rich-resource. We define a language pair as a low/rich-resource language pair according to whether the data size of this language pair larger than the median size (640k) of all language pairs. From Figure 7 we can observe that the large improvements (> 8 BLEU) of SCoMoE are all achieved on the low-resource language pairs. Among these language pairs, en-af only has 274,80 721 sentence pairs, but gets the greatest improvement of nearly 16 BLEU. All of these suggest that SCoMoE alleviates the overfitting issue of Gshard especially on low-resource language pairs.
E EVALUATION ON TRAINING SPEED AND COMMUNICATION TIME
In this section, we compared the training speed and communication time of SCoMoE-San and SCoMoE- Feat with different settings of and 3. We report the training speed in the metric of words-per-secome (WPS), which is the number of tokens that the model can process per second, and the time cost of each all-to-al communication in milliseconds.
17

---

Published as a conference paper at ICLR 2023
ScoMeC
Ochars
401
(b)
(d)
Figure 8: (8a): The speed of base SCoMoE with different (8b): The communication time of base SCoMoE with different a. (8c): The speed of large SCoMoE with different a/B. (8d) The com- munication time of large SCoMoE with different o/33 SCoMoF-Seq/Feat- denotes the SCoMoE- Seq/Feat with the intra-accelerator and global communication, for which the x-axis is the ratio of the intra-accelerator communication (ax). SCoMoE-Seq/Feen the intra-node and global communication, for which the x-axis is the ratio of the intra-node commu- nication (8).
denotes the SCoMoB-Seq/Fent with
Base Model We tuned the ratio of the intra-accelerator communication (a) from 0 to 0.5, and illustrated the training speed and communication time of SCoMoR. at different ratios in Figure [a and 8b] respectively. In spite of training with fixed dummy inputs, the training speed of the base models is still unstable. Hence we trained the same model for three times and show the mean and the standard deviation over the three runs in Figure Sa and 8b
From Figure Sa we observe that the two models can trained faster as the ratio of intra-accelerator communication (a) increases, except that the training of both SCoMoB-Seq and SCoMoE-Feall be- comes slower when the ratio is 0.1, but with less communication time than when the ratio is 0. This suggests that the structured communication reduces the communication time but may cause addi- tional overhead. The extra overhead is probably from the computation, since computations in each expert also split into multiple groups corresponding to multiple communication groups, which is slower than performing all computations at the same time. Even if we set the ratio of the intra- accelerator communication to O, SCoMoB-Seq is still significantly faster than Gshard. The reason is that the communication cost of token clustering is half as that of the top-2 routing. We also notice that the SCoMoE- Seq is always faster than SCoMoE-Fees and requires less communication time. This is because the token clustering is only applied to the SCoMoE-Seq,
Large Model Both SCoMoE -eelee-tee//eataa and -3 can be used for training large models. The dif- ferences between them is that SCoMoE-Seq/Peeee Seq/Feat-8 accelerate the training speed with the intra-node communication instead of the intra-accelerator communication. We have compared the training speed and communication time of them with different settings of a/3 in Figure &c and Rd respec- tively. Overall, all variants of SCoMoE train faster and require less communication time as a or [8 increases. It is also clear that SCoMoE-Seq/Feate is always faster than SroMooo-Seg/Fatt.. is not surprising as the intra-accelerator communication is much faster than the intra-node commu- nication. But we also notice that the training speed and communication time of SCoMoE-Feat is not significantly better than Gshard until 3 reaches 0.5, while that for SCoMoB-Seq- linearly scales as 3 increases. The reason for this might be that the amount communication time is related to the shape of communicated tensor. It seems that the communication will be faster if the feature dimension is the power 2.
This
Theoretic Analysis on Communication Time We have empirically evaluated the training speed and the amount of communication time of SCoMoE by varying and 3. The bandwidth of com- munication has a great impact on these empirical results in practice. To complement the empirical evaluation, we further theoretically analyze the amount of time required by the all-to-all communica- tion and estimate speedups that can be achieved by the proposed structured communication method. We estimate the all-to-all communication time the combination of the intra-node, and inter-node communication time as the time cost for the intra-accelerator communication (memory copy) is neg- ligible in comparison to the other two communication channels. Suppose that the total amount of data for communication is MB, and the number of accelerators in one node is m. Given that, the amount of data for the intra-/inter-node communication is D(m-1), Suppose that the bandwidth of the intra/inter-node communication B MB/s. The amount of the all-to-all communication
mh m
18

---

Published as a conference paper at ICLR 2023
time (seconds) can be estimated as follows:
D D(m 1)
T Tall-to-all
mB1
mBB
D
1)
B2 T
+
(m - 1)
According to the equation, increasing the bandwidth of the inter-node communication (B₂) is more - beneficial than increasing the bandwidth of the intra-node communication (Bi) as m-I S 1. If we set the ratio of the intra-node communication to 3 (a = 0, < B < 1), the amount of the all-to-all communication time can be computed as:
D
To (1 aol-tolll +
131
D
D(m - 1)
(1 3) mB1
+
(8)
mB²
B1
D(m
1
D
= (1 B)
+
m
B2
B
B
The absolute value of the derivative of 3 is proportional to B3 B. Either decreasing B2 in- creasing B1 would make the SCoMoE more efficient. Therefore, If the bandwidth of the intra-node communication is slow, is better to set B O. Typically, there are two types of intra- connec- tions with different bandwidths: PCIE and NVlink/NVSwicch The NVlink/NV much faster than the PCIE connection. Furthermore, the network interface controller (NIC) is also connected with accelerators through PCIE. More communications sharing PCIE connections, will make the communication slower. Therefore, we recommend to set B > 0, a = 0 if fast intra-node connections with NVlink/NVSwitch are available, and 0, 0 if slow intra-node connec- tions with PCIE are used. We use the PCIE connections in our experiments, hence the speedup of SCoMoE-3 should be more significant with the fast NVlink/NVSwithh connections.
connection
F EFFECTS OF TOKEN CLUSTERING
Table 4: The translation performance and training speed of Gshard and SCoMoE-Seq with/without token clustering and layer reordering. vI x denotes with/without. Clustering denotes the token clustering. Reordering denotes changing the order of layers according to the architecture in Figure
Model
Clustering Reordering
Multilingual
Bilingual
Speed
Gshard
24.41
37.96
1x
Gshard
V
22.73
38.29
1x
Gshard
V
22.86
37.93
1.1x
SCoMoE-Seq
X 22.81
38.0
1.05x
SCoMoE-Seq
V
24.45
38.24
1.1x
We compared the translation performance and training speed of Gshard and SCoMoE with vs. with- token clustering. Results are shown in Table 4 could be seen that the Gshard with token clustering or layer reordering performs comparably to Gshard without using them on bilingual trans- lation, but is significantly worse than the original GShard on multilingual translation. Results of SCoMoE demonstrate that token clustering can significantly benefit our model on both bilingual and multilingual machine translation in terms of both BLEU and training speed. The reason for that Gshard with token clustering and reordering has the same training speed (1.1x) with SCoMoE are two folds: 1) Token clustering speeds up the training of both Gshard and SCoMoE, since the communication cost of token clustering is only half that of Top2Gate routing. 2) SCoMoE brings additional computational overhead and requires a setting of higher a to be faster than Gshard, which has been discussed in Appendix
19

---

Published as a conference paper at ICLR 2023
375
BLEU (WMT17-En-FF
BLEU (OPUS- (PUS-100) 100)
BLEU (WMT17-En-FF)
BLEU (OPUS
OPUS-100 weght-0.01
27.3
OPUS-100 100 wegh
03
0.5
077
1.0
0.5
0.6
07
0.8
0.5
10
Temperature
Capacity Factor at Inference
(b)
Figure 9: (a) The translation performance SCoMoE-Seq with different temperatures on the test of OPUS-100 and WMT17-En-FFr. The horizontal dashed line denotes the BLEU scores of the SCoMoR-Seq with the sigmoid gate on OPUS-100 and WMT17-En-Frr (b) The translation performance of Gshard with different eval-capacity-facte the BLEU scores of the standard Transformer on OPUS-100 (blue) and WMT17-En-Fr (red). The two dashed lines in Figure (a) overlapped together.
The horizontal dashed lines denote
G IMPORTANCE OF DIFFERENTIABLE SORTING
We conducted experiments to investigate how the temperature of softmax affect the performance. We compared the translation performance of SCoMoE-Soe with different temperatures. Results are plotted Figure 9a from which we could see that the performance of the model is sensitive to the temperature. We also report the performance of SCoMoE-Seq with the sigmoid gate the baseline, which is also displayed in Figure 9a (the horizontal dashed line). It could be seen that SCoMoE-Sed with the softmax gate is consistently better than that with the sigmoid gate different temperatures.
H EFFECT OF CAPACITY FACTOR AT INFERENCE
The capacity factor is of importance to the training of MoE models (Yang al. 2021). In multilingual NMT experiments, we have found that the models trained on the OPUS-100 dataset performs poorly the test set, unless with high capacity factor, substantially higher than that used during training. For convenience, we refer to the capacity factor at inference as eval-capacity- factor. The capacity during training and inference can be computed as follows:
x number of tokens
capacity train number of experts
x capacity-factor
(9)
capacity number of tokens eval-capacity-
(10)
We evaluated the trained Gshard model with different eval-capacity-facee factors mance in Figure 9bb The model for evaluation was trained with the capacity factor of 1.0. The performance of Gshard is even worse than the standard Transformer, when the eval-capacity-face is smaller than 0.75. Setting wallccppccty-ffccor--la if the number of experts is 8, which is very expensive. The costs for communication, computation and memory are all proportional to the capacity factor, hence running Gshard on the OPUS- dataset expensive. We also show the impact of eval-capacity-ficn mance on the WMT17-En-Frr dataset in Figure [9b] It is interesting that the performance of Gshard on WMT17-En-FF decreases as the eval-capacity-face vation on OPUS-100. Increasing the eval-c eval-capacity-facte hurt the performance. But why do the the multilingual MoE models require large capacity? We conjecture that the routing of multilingual MoE is language-specific, i.e., the tokens in related lan- guages are routed to the same experts. In our experiments, each batch contains multiple languages during training, but only routed to the same experts at inference, which causes these experts to overflow and tokens to be dis- carded. To validate hypothesis and alleviate this problem, increase the load-halance weight
-actorss and show the perfor-
is equivalent to setting capacity-facor=s3
on the translation perfor-
increases, which is opposite to the obser- changes the routing of MoE, which may
language inference. Therefore, sentences in the same language
20

---

Published as a conference paper at ICLR 2023
Table The results of SCoMoE with different settings of a, and y on the testset WMT17-En- Fr. SCoMoE-Seq* and SCoMoE-Feat* indicate the models where 8 and y are all positive.
Model
B
Speedup
BLEU
Gshard
1x
40.09
SCoMoE-Seq-a
0.2
0
0.8
1.2x
39.3
SCoMoE-Seq-B
0
0.5
0.5
1.25x
40.20
SCoMoE-Seq"
0.2
0.2
0.6
1.23x
39.6
SCoMoE-Seq**
0.2
0.5
0.3
1,41x
40.07
SCoMoE-Feaate
0.5
0
0.5
1.32x
39.79
SCoMoE-FFeat
0
0.5
0.5
1.10x
39.93
SCoMoF-Feat*
0.25
0.25
0.5
1.16x
40.39
SCoMoE-Feen 0.25
0.5
0.25
1.26x
40.50
in Gshard from 0.01 to 0.I, and observe significant improvements in Figure 9b But still need high capacity factor for the inference of multilingual MoE models. We leave the issue to our future work.
I SETTING a, B AND TO BE NON-ZERO SIMULTANEOUSLY
Although there are three communication channels, we have only used two of them in SCoMoE-a (a =0, >0, y >0) and SCoMoE-3 =0, 3 >0, 7 >0) our previous experiments, since is difficult to tune three hyperparameters at the same time. In this section, we investigated whether it beneficial to set a, 3, y be non-zero at the same time. In Figure Bal we have observed that a should be set to a small value for a comparable performance with Gshard. Therefore we set a to be 0.2 and 0.25 for SCoMoE-Sec, and SCoMoE-Feat respectively. The results are shown in Table Comparing SCoMoE-Seq* with SCoMoB-Seqq-a and SCoMoE-Seq-B3 we can find that keeping and B non-zero at the same time could speed up the model training without the performance drop. Furthermore, SCoMoE-Seyy with (a 0.25, 3 0.5, 0.25) is even better than that with (a = 0.25, 3 = 0.25, = 0.5) in terms of both BLEU and training speed. Therefore, it is better for set small a and large for SCoMoF-Seq*..
SCoMoE-Feat also benefits from setting a, ,3 and y to be non-zero at the same time, especially for the improvement BLEU, is reasonable because increasing communication channels for SCoMoE means that each token can be computed in more experts. This is similar to increasing heads in multi-head attention.
J COMPARISON OF TRANSLATION SPEED
In previous experiments, all reported results on speed are those of training speed. In this section we compared the speed of autoregressive decoding of SCoMoE and analyzed the influence of hyper- parameters on translation. For SCoMoE-Sec, the hyperparameters of each communication channel a, 3, y could tuned at inference and be different from the values training. Therefore, we tuned and / for SCoMoE-Seq-a and SCoMoE-Seq-B respectively and report BLEU and translation speed in [Oa rerrepecciilssrss
Overall, the speedup of SCoMoE during autoregressive decoding is weaker than that at training. can be seen that SCoMoF-Seq-a is faster than Gshard consistently at different settings of a. SCoMoE-Feat-a also shows slight speedup over Gshard. However, both SCoMoE-Seq-33 and SCoMoE-Feates are slightly slower than Gshard at inference. The reason is that the communication during autoregressive inference is bottlenecked by the frequency rather than the size of communica- tion data at inference. But one exception is that SCoMoE-Seq-- with / = faster than Gshard due to the token clustering.
With regard to the translation performance, both SCoMoE-Seq-a and - obtain their best perfor- mance with a/3 3 = 0. It is also clear that the performance of them decreases a/ increases.
21

---

Published as a conference paper at ICLR 2023
40.30
130
100
99.50
29.00
30.75
900MeE Foudd
Pozt
61
022
015 04 0.5
0.0 01 02
04
0.
ag
(a)
(b)
Comparison SCoMoE during autogressive decoding in terms BLEU 10a) and the decoding speed (10b) on the testset of WMT17-En-Frr The x-axis (a/ B) represent the hyper- parameters tuning at inference instead of training, which different from that Figure [3 and 3b SCoMoF-Seg/Feat-o and SCoMoF-Seq/Feat-) are trained with (a 0.5, 3 0,, 0.5) and (a = 0., 3 = 0.5, = 0.5) respectively.
Figure 10:
BLEU
BIEU
Gehand
Gehant
SOOMoE Sequl
0.7)
900MIE Soai ti 0. 6.3. 0.7)
022
SidaMet-Sate Sager G.S. 5- 0.. 0.5)
Sooteit-Saee Saaid to 0. G. 2.5)
SCoMeE Sequd 6.7, 6. y...
S0OMAE See (x- 0.7.y----.-... 0.3)
442
20000
40000
90000
50000 100000
2200000
20000
A2000
00000
50000
300000
120000
Time (Seconds)
Thme (pnconds)
(a)
42
-4
Gahard
SColle Fest m.25
Figure 11: Performance- Time curves of SCoMoE-Seq-a 1la.. SCoMoE-Seq.PP SCoMoE- Feat (Tlc) on validation set of WMT17-En-Fir
11b
and
All in all. it is better set a/3 to be 0 for SCoMoE-Seg at inference in terms of both translation performance and speedup. We use this setting to report BLEU scores for SCoMoE-Seg in other sections.
22

---

Published as a conference paper at ICLR 2023
K PERFORMANCE-TIME CURVES
We have shown the improvements SCoMoE over Gshard in both training speed and translation performance, but is still unclear how much time SCoMoE could save compared with Gshard. It is also interesting to see whether SCoMoE could benefit from more training steps. Therefore we compare the performance-iimg curves of (CoMoE-Seq-o, SCoMoE-Seq-2 and SCoMoEE Gshard in Figure Tla Figure [1b and Figure IIc respectively. We can see that almost all SCoMoE variants are superior to Gshard after being trained for 45,000 seconds. And they achieve 50 BLEU at around 80,000 seconds, while Gshard requires 129,000 seconds. Among all variants, SCoMoEE- seem
with
the best. We then trained these two variants with more steps (12 epochs and 15 epochs respectively) time comparable to that of Gshard (10 epochs). The results of extra steps are illustrated with dashed lines in Figure 11a and Figure Tib It be seen that these two variants obtain improvements with more training steps and their final performance much better than Gshard. Most variants of SCoMoE-Feat converge around 80,000 steps, and only see little improvements with more training steps.
seq with (0=00.,30..5,=0055) and (0=0.3,,=0.,,=001
that they were trained with
23


Reviews:
[{"review_id": "wRb90xhesOR", "reviewer": "Reviewer_2udu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "This paper has the main motivation to improve the “sharded” distributed MoE by improving its all-to-all communication by not treating them agnostically and equally, but instead characterizing them hierarchically based on their locality relationships, namely accelerator, node and global levels. The motivated application is bi-/multi-lingual machine translation. Experiments show better performance than the baseline Gshard. ", "summary_of_the_review": "The paper has good motivation and direction, as well as their implementation. However, I think it needs improvement to be well convincing. ", "strength_and_weaknesses": "Strength: Having a good motivation and direction in ML since many of us work on higher abstraction levels, and will benefit from it in saving time and resources. \n\n\nSome Weaknesses\n1/ The main motivation is not enough without the help from token clustering and the softmax gating + differential sorting in place of sigmoid routing. As a result, it would be helpful to study how each of those 3 contributes to the main results. \n\n2/ Experiments show improvements, but directly compared to Gshard, it seems not so much given how far the solution goes with a combination of methods. \n\nThe all-to-all time is very helpful. However, compared to Gshard in terms of practicality of being simplistic, what is the overhead of this solution given the hierarchy and added interfaces? Likewise, it would be also helpful to see the training time, inference time which a typical user would also like to see. \n\n\n3/ The paper mentions the balancing to some good extent, yet there is no space of analysis to this. Probably it needs a section for regularization of balancing, for clustering might not work perfectly in terms of balancing, and so many tokens must be dropped as some expert reaches its max capacity (assuming the balancing regime is similar to Gshard – please correct if I am wrong). \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "pTMdTZYY-X", "reviewer": "Reviewer_UrrU", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "The paper proposes a new MoE architecture (and its implementation) that is faster than traditional MoE models. MoE's usually involve sending tokens from one device (where the token is routed) to others (where the selected experts are hosted). This is usually very expensive and one of the main reasons why MoEs are not more popular yet (otherwise the gap in performance vs time wrt dense models would be even larger). This paper proposes a way to partially overcome this issue.\n\nThe main idea is to send more tokens to nearby devices. Of course, one needs to be able to choose the right tokens as otherwise the actual extra capacity of having many experts will fade away. There are three types of device communication involved: intra-accelerator, intra-node, and inter-node (from cheaper to more expensive). Assume that every accelerator device contains only one expert. The proposed routing algorithm splits tokens into three groups, those that will be assigned to experts requiring intra-accelerator, intra-node, and inter-node communication, respectively. Note that sending tokens only within a node implies that not all experts will be available to that token (only those in their node!) -- and this will vary depending on where the tokens happen to be. So, if we enforce sending more tokens locally, as in principle those tokens come from a random subset of images, the quality of routing and expert specialization may suffer. The paper acknowledges and confirms this intuition.\n\nThe paper proposes two ways to select which tokens are provided with each type of \"communication-based\" routing. The first one (\"sequence-based\") aggregates the standard routing weights across experts groups. Depending on the local device for the token, there's experts that require intra-accelerator, intra-node, and inter-node communication, respectively. The router first looks at the sum of weights per token across all the intra-accelerator experts. It ranks the tokens based on this sum, and takes the top-K (where K is a hparam here, a % of the local batch) and applies routing and dispatching to these tokens only, across the intra-accelerator experts. Once those tokens and experts are removed, the same idea is applied to the intra-node experts, and so on. To account for \"removed\" experts, the missing logits are replaced with -inf before normalizing.\n\nThe second routing type, feature-based, projects each token into three representations of different dimensions (corresponding to intra-accelerator, intra-node, and inter-node) summing up to the original dimension. Every expert is also \"split\" into three smaller subexperts, one that can handle tokens from each of the three dimensions. It's not fully clear to me how the \"partitioned tokens\" are then routed. Are there three routers (one per communication dimension), and the same logic as for \"sequence-based\" is applied? If so, different \"parts\" of a token may end up --not only on different experts-- but on different communication levels?\n\nIn order to maximize the chances of having good local routing, the paper proposes to pre-cluster tokens, dispatch similar tokens to the same devices, and then apply the ideas described above. Accordingly, the paper must change the attention-mlp structure to do token clustering only once (so that tokens from the same image can end up in different devices). Figure 2 (left) shows the proposed way to do this. I think this is a cool contribution.\n\nThe paper contains a number of experiments for language models where performance and speed are compared between SCoMoE and gShard (a well-known sparse model for language). My takeaway was that SCoMoE can match gShard performance while saving 20-30% of the *overall* time (or is it just all-to-all communication time?), which is quite a nice achievement! They also study what fraction of tokens should be assigned to each of the three groups depending on the routing method.", "summary_of_the_review": "The paper proposes a new MoE architecture that tries to send most tokens to nearby devices/experts, and applies a couple of tricks to improve the required routing. First, tokens are clustered so that each device gets similar tokens. To make this possible, attention and MLPs are re-arranged so that attention layers aren't applied once MoEs start being used. Two routing techniques are presented, both of which try to identify the tokens that will be better served locally.\n\nIt's a nice paper that I really enjoyed. There's a couple of questions regarding the experimental setup and the results whose answer could make the case and my score for the paper stronger.\n\n----------------------------------------------------------------------------------------------\nAfter reading the authors' reply, I've decided to increase my score from 6 to 8. ", "strength_and_weaknesses": "I have a couple of questions or comments.\n\nFirst, I think the work would benefit from indeed having cases where both \\alpha, \\beta and \\gamma are *simultaneously* non-zero. Most of the paper describes this case, but I think it's not covered in any of the experiments.\n\nSecond, is the speed-up multiplier computed with respect to the total step time? or just the all-to-all communication? Obviously, we care the most about the former. MoE models only have a few sparse layers in general. Thus, saving a 10% wrt to a total of 5% ends up being a very small saving to justify severe engineering extra complexity.\n\nThird (and related to second), if a model saves X% time with respect to another, a nice way to get a sense of how they compare is to run both for the same amount of total time, and look at the final performances. In other words, figure out if the cheap model can re-use the time to improve performance too (or if, unfortunately, the cheap model also has a lower ceiling). A truly informative plot in these cases is performance vs total training time. I'd love to see SCoMoE vs gShard in that type of plot to see how meaningful the gains are.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TBrGQAEfC8L", "reviewer": "Reviewer_NyCC", "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": "The paper proposed a novel routing algorithm for mixture of experts that took the network bandwidths among the accelerators into consideration, which allows efficient communication for token shuffling especially on the GPUs.\n", "summary_of_the_review": "Novel routing algorithm for mixture of experts whose effectiveness were demonstrated by solid experimental results.\n", "strength_and_weaknesses": "Strengths:\n1) The authors took practical device constraints into consideration when designing the neural network architecture, and reported real benchmarks instead of just bio-O analyses. \n\n2) As the emergence of ultra-expensive large language models, how to enable affordable large models using more efficient algorithms becomes an important research area. \n\n3) The authors demonstrated the effectiveness of SCoMoE with well designed experiments and evaluations. Real world metrics like BLEU scores and step time are reported. The paper also reported the analysis of capacity factor during inference, which is a critical hyper-parameter for MoE performance.\n\n4) I want to point out that the reported speedup against GShard might seem small (less than 1.5x), however, it's due to the relative smaller number of experts used in the experiments (I guess bounded by compute budget). The speedup achieved by SCoMOE would become more significant when the number of experts is even larger. In the original GShard paper, thousands of experts were used where the communication cost would dominate the total compute cost. \n\n\nThis paper can be even better if the author can address the following concerns:\n\n1) What's the efficiency improvement of this algorithm during autoregressive decoding, where the model decode one token after another sequentially? \n\n2) It's better to estimate how large a dense model needs to be in order to achieve similar translation quality as SCoMoE. This would demonstrate the cost saving from SCoMoE comparing to a dense model with similar quality.\n\n3) Need to discuss the pros and cons of this proposed algorithm comparing to related routing algorithms also aiming to reduce communication costs, such as:\nBeyond Distillation: Task-level Mixture-of-Experts for Efficient Inference (https://arxiv.org/pdf/2110.03742.pdf)\nMixture-of-Experts with Expert Choice Routing (https://arxiv.org/pdf/2202.09368.pdf)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7IHgjV29P00", "reviewer": "Reviewer_5Gem", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "This paper proposes SCoMoE in order to reduce the all-to-all communication in mixture-of-experts model training. The core technique design is based on projecting the feature dimension (SCoMoE-Feat) into low-dimensional representations and using token clustering to compensate for the potential performance drop. ", "summary_of_the_review": "I think the idea in SCoMoE is interesting and suggests good statistical efficiency; however, the presentation of the paper should be improved and some comparison with some more advanced state-of-the-arts should be included to show the hardware efficiency. \n", "strength_and_weaknesses": "Strength:\n\n+ The idea of using feature project and token clustering is a reasonable relaxation of the original all-to-all communication in the MoE parallel training paradigms. \n\n+ The experiments suggest the good statistical efficiency of the proposed method (in terms of no drop in model quality). \n\n\nWeakness\n\n+ Some technique presentation is unclear and hard to track.\n\n+ Lack of comparison with some more advanced state-of-the-arts. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster