Original: SC O M O E: E FFICIENT M IXTURES OF E XPERTS WITH S TRUCTURED C OMMUNICATION
Revised: SC O M O E: E FFICIENT M IXTURES OF E XPERTS WITH S TRUCTURED C OMMUNICATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Mixture-of-Experts (MoE) models are promising architectures for massive multilingual neural machine translation and large language models due to the advantage of sublinear scaling.
Revised: Mixture-of-Experts (MoE) models are promising architectures for massive multilingual neural machine translation and large language models due to the advantage of sublinear scaling.

Original: However, the training of large MoE models is usually bottlenecked by the all-to-all communication (Lepikhin et al., 2020).
Revised: However, the training of large MoE models is usually bottlenecked by the all-to-all communication (Lepikhin et al., 2020).

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

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

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

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

Original: MoE layers.
Revised: MoE layers.

Original: The sigmoid gating in the balanced router used in the token clustering is substituted with the softmax gating with differential sorting.
Revised: The sigmoid gating in the balanced router used in the token clustering is substituted with the softmax gating with differential sorting.

Original: Experiments on massive bilingual and 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.
Revised: Experiments on massive bilingual and 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Recent years have witnessed a substantial interest in exploring sparse architectures based on Mixture of Experts for training massive multilingual machine translation (Lepikhin et al., 2020;
Revised: Recent years have witnessed a substantial interest in exploring sparse architectures based on Mixture of Experts for training massive multilingual machine translation (Lepikhin et al., 2020;

Original: Kim et al., 2021) and large language models (Fedus et al., 2021; Zhang et al., 2021; Ma et al., 2022; Du et al., 2021; Zoph et al., 2022; Rajbhandari et al., 2022; Lin et al., 2021).
Revised: Kim et al., 2021) and large language models (Fedus et al., 2021; Zhang et al., 2021; Ma et al., 2022; Du et al., 2021; Zoph et al., 2022; Rajbhandari et al., 2022; Lin et al., 2021).

Original: Experts of MoE models are distributed over multiple devices.
Revised: Experts of MoE models are distributed over multiple devices.

Original: Due to the sparse architecture where only a combination of experts are selected to process each input, the number of experts and hence the scale of MoE models can be sufficiently large while the computational cost is only sublinear to the number of parameters.
Revised: Due to the sparse architecture where only a combination of experts are selected to process each input, the number of experts and hence the scale of MoE models can be sufficiently large while the computational cost is only sublinear to the number of parameters.

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

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

Original: To mitigate this issue, we propose Structured Communication based MoE (ScoMoE), which treats the all-to-all communication in a structured way rather than equally across different devices.
Revised: To mitigate this issue, we propose Structured Communication based MoE (ScoMoE), which treats the all-to-all communication in a structured way rather than equally across different devices.

Original: 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) is faster than that across accelerators, and the bandwidth inside a node (intra-node) is faster than that across nodes (inter-node).
Revised: 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) is faster than that across accelerators, and the bandwidth inside a node (intra-node) is faster than that across nodes (inter-node).

Original: Figure 1a visualizes the hierarchical structure of communication topology with a 9 ×matrix, where different levels of communication are in different colors.
Revised: Figure 1a visualizes the hierarchical structure of communication topology with a 9 ×matrix, where different levels of communication are in different colors.

Original: We view the data flow in the all-to-all communication from the perspective of two dimensions: sequence dimension (tokens) and feature dimension (embeddings of tokens).
Revised: We view the data flow in the all-to-all communication from the perspective of two dimensions: sequence dimension (tokens) and feature dimension (embeddings of tokens).

Original: The proposed SCoMoE transforms (slicing or projecting) the data flow at either the sequence or feature dimension into three communication groups: intra-accelerator, intra-node and global (inter-node) communication, as shown in Figure 1b.
Revised: The proposed SCoMoE transforms (slicing or projecting) the data flow at either the sequence or feature dimension into three communication groups: intra-accelerator, intra-node and global (inter-node) communication, as shown in Figure 1b.

Original: For the data slicing on the sequence dimension, we select tokens for a communication group according to
Revised: For the data slicing on the sequence dimension, we select tokens for a communication group according to

Original: Accelerator 0 Accelerator
Revised: Accelerator 0 Accelerator

Original: Accelerator 2 .
Revised: Accelerator 2 .

Original: Accelerator
Revised: Accelerator

Original: A cce l er a t
Revised: A cce l er a t

Original: o r
Revised: o r

Original: A cce l er a t
Revised: A cce l er a t

Original: o r
Revised: o r

Original: A cce l er a t
Revised: A cce l er a t

Original: o r
Revised: o r

Original: A cce l er a t
Revised: A cce l er a t

Original: o r
Revised: o r

Original: Node Node
Revised: Node Node

Original: Node. .
Revised: Node. .

Original: Intra-accelerator Communication
Revised: Intra-accelerator Communication

Original: Intra-node Communication
Revised: Intra-node Communication

Original: Inter
Revised: Inter

Original: [Figure]
Revised: [Figure]

Original: node Communication (a)
Revised: node Communication (a)

Original: Token 

Revised: Token

Original: Token
Revised: Token

Original: Slicing on the Sequence Dimension
Revised: Slicing on the Sequence Dimension

Original: Feature
Revised: Feature

Original: Token 
 Token
Revised: Token 
 Token

Original: Token 2 Token
Revised: Token 2 Token

Original: Token 4 Token
Revised: Token 4 Token

Original: Sequence
Revised: Sequence

Original: T o k e n 0 ...
Revised: T o k e n 0 ...

Original: T o k e n 5 Transpose ... T o k e n
Revised: T o k e n 5 Transpose ... T o k e n

Original: Projection on the Feature Dimension
Revised: Projection on the Feature Dimension

Original: Token
Revised: Token

Original: Token 1 Token
Revised: Token 1 Token

Original: Token
Revised: Token

Original: T o k e n(
Revised: T o k e n(

Original: [Figure]
Revised: [Figure]

Original: the assignment scores of the tokens with the experts inside the group.
Revised: the assignment scores of the tokens with the experts inside the group.

Original: To organize the data transformation 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.
Revised: To organize the data transformation 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.

Original: Theoretically, structuring all-to-all communication in this way is faster than its original form, since less data are transmitted through the slow inter-node communication channel.
Revised: Theoretically, structuring all-to-all communication in this way is faster than its original form, since less data are transmitted through the slow inter-node communication channel.

Original: However, this may hurt the performance, because intra-accelerator/intra-node communication can be processed by only a part of experts.
Revised: However, this may hurt the performance, because intra-accelerator/intra-node communication can be processed by only a part of experts.

Original: To alleviate this issue, we further propose a token clustering approach to aggregating related tokens from different devices, which elevates the association of tokens inside each device (accelerator/node).
Revised: To alleviate this issue, we further propose a token clustering approach to aggregating related tokens from different devices, which elevates the association of tokens inside each device (accelerator/node).

Original: The proposed token clustering uses the balance router presented by Lewis et al.
Revised: The proposed token clustering uses the balance router presented by Lewis et al.

Original: (2021) for clustering, where each device is a cluster.
Revised: (2021) for clustering, where each device is a cluster.

Original: In the balanced router, a 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.
Revised: In the balanced router, a 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.

Original: Hence we propose to replace the sigmoid gate with the softmax gate via a straight-through trick (Bengio et al., 2013) for better gradient broadcasting.
Revised: Hence we propose to replace the sigmoid gate with the softmax gate via a straight-through trick (Bengio et al., 2013) for better gradient broadcasting.

Original: In a nutshell, our contributions are summarized as follows:
Revised: In a nutshell, our contributions are summarized as follows:

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

Original: A token clustering method is proposed to dispatch the related tokens to the same devices to alleviate the routing locality in the structured communication.
Revised: A token clustering method is proposed to dispatch the related tokens to the same devices to alleviate the routing locality in the structured communication.

Original: We propose to use softmax gate to substitue the sigmoid gate in the balanced router for better gradient broadcasting.
Revised: We propose to use softmax gate to substitue the sigmoid gate in the balanced router for better gradient broadcasting.

Original: Experiments on massive bilingual and multilingual machine translation demonstrate that SCoMoE is faster than Gshard significantly with comparable or even better translation performance.
Revised: Experiments on massive bilingual and multilingual machine translation demonstrate that SCoMoE is faster than Gshard significantly with comparable or even better translation performance.

Original: Further analysis discloses the strategies of selecting hyper-parameters for SCoMoE.
Revised: Further analysis discloses the strategies of selecting hyper-parameters for SCoMoE.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: MoE models (Jacobs et al., 1991; Jordan & Jacobs, 1994) are ensemble method to integrate multiple experts.
Revised: MoE models (Jacobs et al., 1991; Jordan & Jacobs, 1994) are ensemble method to integrate multiple experts.

Original: Shazeer et al.
Revised: Shazeer et al.

Original: (2017) propose a gating network to select a combination of experts and mix data parallelism and model parallelism to increase the batch size.
Revised: (2017) propose a gating network to select a combination of experts and mix data parallelism and model parallelism to increase the batch size.

Original: Gshard (Lepikhin et al., 2020) utilize the MoE parallelism proposed by Shazeer et al.
Revised: Gshard (Lepikhin et al., 2020) utilize the MoE parallelism proposed by Shazeer et al.

Original: (2017) to scale Transformer by replacing the MLP layer of Transformer with multiple MLP experts.
Revised: (2017) to scale Transformer by replacing the MLP layer of Transformer with multiple MLP experts.

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

Original: Deepspeed-MoE (Kim et al., 2021) combines ZerO (Rajbhandari et al., 2020) and MoE parallelism to train large MoE models.
Revised: Deepspeed-MoE (Kim et al., 2021) combines ZerO (Rajbhandari et al., 2020) and MoE parallelism to train large MoE models.

Original: Many recent large-scale pretrained language models (Fedus et al., 2021; Zhang et al., 2021; Ma et al., 2022; Du et al., 2021; Zoph et al.,
Revised: Many recent large-scale pretrained language models (Fedus et al., 2021; Zhang et al., 2021; Ma et al., 2022; Du et al., 2021; Zoph et al.,

Original: 2022) have been trained with Gshard and its variants.
Revised: 2022) have been trained with Gshard and its variants.

Original: A variety of new routers, e.g., hash router (Roller et al., 2021), random router (Zuo et al., 2021) and balanced router (Lewis et al., 2021), have been proposed in lieu of the top-2 router in GShard.
Revised: A variety of new routers, e.g., hash router (Roller et al., 2021), random router (Zuo et al., 2021) and balanced router (Lewis et al., 2021), have been proposed in lieu of the top-2 router in GShard.

Original: We employ the balanced router proposed by the BASE Layer (Lewis et al., 2021) for token clustering.
Revised: We employ the balanced router proposed by the BASE Layer (Lewis et al., 2021) for token clustering.

Original: The balanced router uses the auction algorithm (Bertsekas, 1992) for balanced routing.
Revised: The balanced router uses the auction algorithm (Bertsekas, 1992) for balanced routing.

Original: Although sparse MoEs are computationally efficient, they suffer from frequent and expensive all-toall communication across the network.
Revised: Although sparse MoEs are computationally efficient, they suffer from frequent and expensive all-toall communication across the network.

Original: Gshard (Lepikhin et al., 2020) has shown that the all-to-all communication becomes the bottleneck for training large MoEs.
Revised: Gshard (Lepikhin et al., 2020) has shown that the all-to-all communication becomes the bottleneck for training large MoEs.

Original: While the heterogeneous hierarchical network topology has been exploited to reduce the communication consumption, previous studies (Lin et al., 2020; Castiglia et al., 2021) focus on the all-reduce communication in SGD rather than the all-to-all communication in MoE. Tutel (Hwang et al., 2022) propose a hierarchical all-to-all communication approach for MoE, which performs all-to-all communication and layout transformation inside a node, and then all-to-all communication across nodes.
Revised: While the heterogeneous hierarchical network topology has been exploited to reduce the communication consumption, previous studies (Lin et al., 2020; Castiglia et al., 2021) focus on the all-reduce communication in SGD rather than the all-to-all communication in MoE. Tutel (Hwang et al., 2022) propose a hierarchical all-to-all communication approach for MoE, which performs all-to-all communication and layout transformation inside a node, and then all-to-all communication across nodes.

Original: Although the hierarchical all-to-all communication of Tutel reduces the number of communication hops, it increases the amount of communication data.
Revised: Although the hierarchical all-to-all communication of Tutel reduces the number of communication hops, it increases the amount of communication data.

Original: Hetu-MoE (Nie et al., 2022) also proposes a hierarchical all-to-all communication mechanism similar to that of Tutel.
Revised: Hetu-MoE (Nie et al., 2022) also proposes a hierarchical all-to-all communication mechanism similar to that of Tutel.

Original: The difference is that Hetu-MoE dispatches all data inside one node to one accelerator, which may cause memory overflow.
Revised: The difference is that Hetu-MoE dispatches all data inside one node to one accelerator, which may cause memory overflow.

Original: Sparse-MLP (Lou et al., 2022) shortens the communication time by eliminating 10% tokens before routing, according to the importance scores between tokens and experts.
Revised: Sparse-MLP (Lou et al., 2022) shortens the communication time by eliminating 10% tokens before routing, according to the importance scores between tokens and experts.

Original: However, removing tokens may hurt the performance.
Revised: However, removing tokens may hurt the performance.

Original: Fast-MoE (He et al., 2021) optimizes the training speed of MoE on the computation of experts, rather than communication.
Revised: Fast-MoE (He et al., 2021) optimizes the training speed of MoE on the computation of experts, rather than communication.

Revised: Task-level routing (Zhou et al., 2022) extracts a subset of experts for inference which improves the inference efficiency, but only works for multilingual/multitask MoEs and the extracted submodel can only translate in one direction.

Revised: Expert Choice Routing (Kudugunta et al., 2021) speedups the

Revised: convergence

Original:  Comparison with Tutel Although both SCoMoE and Tutel improve the all-to-all communication of MoE based on the hierarchical network topology, SCoMoE is significantly different from Tutel in the following four aspects:
Revised: speed of MoE with better load balance , but can not reduce the time cost for each step. Comparison with Tutel Although both SCoMoE and Tutel improve the all-to-all communication of MoE based on the hierarchical network topology, SCoMoE is significantly different from Tutel in the following four aspects:

Original: • 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 internode communication and increase that for intra-node communication.
Revised: • 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 internode communication and increase that for intra-node communication.

Original: • Algorithm:
Revised: • Algorithm:

Original: The communication of Tutel is hierarchical, first intra-node, and then inter-node.
Revised: The communication of Tutel is hierarchical, first intra-node, and then inter-node.

Original: However, the communication of SCoMoE is structured, where different data are communicated in different groups (intra-accelerator/intra-node/inter-node).
Revised: However, the communication of SCoMoE is structured, where different data are communicated in different groups (intra-accelerator/intra-node/inter-node).

Original: • Translation Performance:
Revised: • Translation Performance:

Original: The hierarchical all-to-all communication algorithm proposed by Tutel may improve the communication speed, but does not influence the translation performance.
Revised: The hierarchical all-to-all communication algorithm proposed by Tutel may improve the communication speed, but does not influence the translation performance.

Original: By contrast, SCoMoE changes the routing of some data, which may serve as regularization and improve the translation performance.
Revised: By contrast, SCoMoE changes the routing of some data, which may serve as regularization and improve the translation performance.

Original: • Speedup:
Revised: • Speedup:

Original: In certain situations, Tutel even slows down training while SCoMoE still improves training speed.
Revised: In certain situations, Tutel even slows down training while SCoMoE still improves training speed.

Original: The comparison of the two methods on the training speed are shown in Figure 3d.
Revised: The comparison of the two methods on the training speed are shown in Figure 3d.

Original: Despite

Original: these

Original:  significant differences, SCoMoE is orthogonal to Tutel and the two methods can be easilycombined
Revised: Despite these significant differences, SCoMoE is orthogonal to Tutel and the two methods can be easily

Original:   M ODEL 3. S PARSE M IXTURE - OF -E XPERTS AND A LL - TO -A LL C OMMINICATION
Revised: combined  M ODEL 3. S PARSE M IXTURE - OF -E XPERTS AND A LL - TO -A LL C OMMINICATION

Original: Sparse MoE models usually extend Transformer architecture with MoE layers where only a combination of experts are activated for computation each time.
Revised: Sparse MoE models usually extend Transformer architecture with MoE layers where only a combination of experts are activated for computation each time.

Original: Different experts do not share parameters, as shown in Figure 2 (in blue), while other modules (in red) are shared across all devices.
Revised: Different experts do not share parameters, as shown in Figure 2 (in blue), while other modules (in red) are shared across all devices.

Original: Given an input sequence X which contains n tokens { X 1 , ..., X i , ..., X n } with n token embeddings, sparse MoE models assign each token
Revised: Given an input sequence X which contains n tokens { X 1 , ..., X i , ..., X n } with n token embeddings, sparse MoE models assign each token

Original: X i to experts usually according to a token-to-expert assignment matrix M :
Revised: X i to experts usually according to a token-to-expert assignment matrix M :

Original: [Equation]
Revised: [Equation]

Original: W E is a projection matrix, where each column is corresponding to an expert embedding.
Revised: W E is a projection matrix, where each column is corresponding to an expert embedding.

Original: M ij is the assignment score between an expert embedding W E j ⊤ and a token embedding X i .
Revised: M ij is the assignment score between an expert embedding W E j ⊤ and a token embedding X i .

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

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

Original: Experts in devices process the received tokens, and send the processed tokens back to their original devices, which requires another all-to-all communication.
Revised: Experts in devices process the received tokens, and send the processed tokens back to their original devices, which requires another all-to-all communication.

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

Original: The intra-accelerator communication is actually copying data from the memory inside the corresponding accelerator, which is very fast.
Revised: The intra-accelerator communication is actually copying data from the memory inside the corresponding accelerator, which is very fast.

Original: The intranode communication is the communication across different accelerators inside one node, which is much faster than the inter-node communication across different nodes.
Revised: The intranode communication is the communication across different accelerators inside one node, which is much faster than the inter-node communication across different nodes.

Original: However, most of the communication in the all-to-all communication are the slowest inter-node communication, which is visualized in orange in Figure 1a.
Revised: However, most of the communication in the all-to-all communication are the slowest inter-node communication, which is visualized in orange in Figure 1a.

Original: Furthermore, the ratio of inter-node communication increases quadratically with the number of nodes.
Revised: Furthermore, the ratio of inter-node communication increases quadratically with the number of nodes.

Original: Therefore, the all-to-all communication becomes very expensive and hence the bottleneck for training large MoE models with many devices.
Revised: Therefore, the all-to-all communication becomes very expensive and hence the bottleneck for training large MoE models with many devices.

Original: 3. S TRUCTURED A LL - TO -A LL C OMMUNICATION
Revised: 3. S TRUCTURED A LL - TO -A LL C OMMUNICATION

Original: 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: intraaccelerator, intra-node and global communication.
Revised: 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: intraaccelerator, intra-node and global communication.

Original: The key idea is to increase the amount of tokens dispatched/combined in the fast intra-accelerator/node communication channels, and decrease the amount of tokens in the slow inter-node communication channel.
Revised: The key idea is to increase the amount of tokens dispatched/combined in the fast intra-accelerator/node communication channels, and decrease the amount of tokens in the slow inter-node communication channel.

Original: Intuitively, this structured all-toall 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 transforms the data into multiple groups for communication at different levels before expert computation.
Revised: Intuitively, this structured all-toall 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 transforms the data into multiple groups for communication at different levels before expert computation.

Original: Different from the vanilla MoE, we do not compute the token-to-expert assignment scores for all experts.
Revised: Different from the vanilla MoE, we do not compute the token-to-expert assignment scores for all experts.

Original: Instead, we aim to select the experts inside each accelerator for the intra-accelerator communication, and the experts inside each node for the intra-node communication.
Revised: Instead, we aim to select the experts inside each accelerator for the intra-accelerator communication, and the experts inside each node for the intra-node communication.

Original: As shown in Figure 2, the tokens in the intra-accelerator/node group can only be processed by the experts inside the accelerator/node, while the tokens in the global group are processed by all experts.
Revised: As shown in Figure 2, the tokens in the intra-accelerator/node group can only be processed by the experts inside the accelerator/node, while the tokens in the global group are processed by all experts.

Original: The processed tokens are sent back to their original devices.
Revised: The processed tokens are sent back to their original devices.

Original: 3. D ATA T RANSFORMATION FOR S TRUCTURED C OMMUNICATION
Revised: 3. D ATA T RANSFORMATION FOR S TRUCTURED C OMMUNICATION

Original: The data transformation for the structured communication could be either on the sequence or feature dimension.
Revised: The data transformation for the structured communication could be either on the sequence or feature dimension.

Original: The difference of them is visualized in Figure 1b.
Revised: The difference of them is visualized in Figure 1b.

Original: Data Slicing on the Sequence Dimension For the data transformation on the sequence dimension, we slice the tokens into each communication group according to the token-to-expert assignment scores introduced in the Section 3.1.
Revised: Data Slicing on the Sequence Dimension For the data transformation on the sequence dimension, we slice the tokens into each communication group according to the token-to-expert assignment scores introduced in the Section 3.1.

Original: The tokens to be sent to intra-accelerator/node devices have the highest assignment score with the experts inside the devices.
Revised: The tokens to be sent to intra-accelerator/node devices have the highest assignment score with the experts inside the devices.

Original: Specifically, the tokens for local devices are collected via:
Revised: Specifically, the tokens for local devices are collected via:

Original: [Equation]
Revised: [Equation]

Original: where M ij is the assignment score between expert j and token i , E local is a set of experts distributed at local devices, X local is a set of tokens that have the top- k assignment scores with E local .
Revised: where M ij is the assignment score between expert j and token i , E local is a set of experts distributed at local devices, X local is a set of tokens that have the top- k assignment scores with E local .

Original: k denotes the number of tokens to be dispatched to local devices, which is a hyper-parameter.
Revised: k denotes the number of tokens to be dispatched to local devices, which is a hyper-parameter.

Original: The tokens forboth 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 −∞ to prevent the same tokens from being selected for multiple groups of communication. Projection Token Selection
Revised: The tokens for

Original: Intra-accelerator Communication Intra-node Communication
Revised: [Figure] Intra [Figure]

Original: Global Communication
Revised: Global Communication

Original: Structured Communication on Sequence Dimension
Revised: Structured Communication on Sequence Dimension

Original: Structured Communication on FeatureDimension sequence feature sequence feature (Self/Cross) Attention Layer FFN Layer (Self/Cross) Attention Layer
Revised: Structured Communication [Figure] Feature

Original: SCoMoE Layer Token Clustering FFN Layer Recover
Revised: [Figure]

Original: Expert Expert
Revised: Expert Expert

Original: Expert ...
Revised: Expert ...

Original: Node . . .
Revised: Node . . .

Original: Node
Revised: Node

Original: Expert Expert
Revised: Expert Expert

Original: Expert ...
Revised: Expert ...

Original: Node . . .
Revised: Node . . .

Original: Node ... ... ...
Revised: Node ... ... ...

Original: Accelerator Accelerator
Revised: Accelerator Accelerator

Original: Accelerator
Revised: Accelerator

Original: Accelerator Accelerator
Revised: Accelerator Accelerator

Original: Accelerator
Revised: Accelerator

Original: Transpose 
Revised: Transpose 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 −∞ to prevent the same tokens from being selected for multiple groups of communication.

Original: Data Projection on the Feature Dimension To transform the data on the feature dimension, we project the input sequence X with the dimensionality of d model to lower-dimensional representations of d intra-accelerator , d intra-node and d global , with three different projections W intra-accelerator , W intra-node and W global :
Revised: Data Projection on the Feature Dimension To transform the data on the feature dimension, we project the input sequence X with the dimensionality of d model to lower-dimensional representations of d intra-accelerator , d intra-node and d global , with three different projections W intra-accelerator , W intra-node and W global :

Original: [Equation]
Revised: [Equation]

Original: (3) where X ∈ R n × d model , X intra-accelerator ∈ R n × d [Equation] , X [Equation]
Revised: where X ∈ R n × d model , X intra-accelerator ∈ R n × d intra-accelerator , X intra-node ∈ R n × n intra-node , X global ∈ R n × d global .

Original: [Equation]

Original: ∈ R

Original: [Equation]

Original: × n [Equation] node , X global ∈ R n × d global . X intra-accelerator , X intra-node , X global are projected sequences to be dispatched via different communication channels.
Revised: X intra-accelerator , X intra-node , X global are projected sequences to be dispatched via different communication channels.

Original: We display these projections and projected sequences in different colors in Figure 2.
Revised: We display these projections and projected sequences in different colors in Figure 2.

Original: The dimensionality of data dispatched to experts is reduced into d intra-accelerator , d intra-node or d global .
Revised: The dimensionality of data dispatched to experts is reduced into d intra-accelerator , d intra-node or d global .

Original: Hence we also need to reduce the dimensionality of weights in experts correspondingly.
Revised: Hence we also need to reduce the dimensionality of weights in experts correspondingly.

Original: As an expert is actually a fully-connected network, which contains two linear projections W 1 ∈
Revised: As an expert is actually a fully-connected network, which contains two linear projections W 1 ∈

Original: R d model × d ffn , W 2 ∈ R d ffn × d model .
Revised: R d model × d ffn , W 2 ∈ R d ffn × d model .

Original: Specifically, we set the shape of the first projection to be ( ˆ d, d ffn ) , and that of the second projection to be ( d ffn , ˆ d ) .
Revised: Specifically, we set the shape of the first projection to be ( ˆ d, d ffn ) , and that of the second projection to be ( d ffn , ˆ d ) .

Original: ˆ d ∈ { d intra-accelerator , d intra-node , d global } , corresponding to the input dimension.
Revised: ˆ d ∈ { d intra-accelerator , d intra-node , d global } , corresponding to the input dimension.

Original: Actually, the parameters of the experts are projected into three parts, each of which is corresponding to one level of communication.
Revised: Actually, the parameters of the experts are projected into three parts, each of which is corresponding to one level of communication.

Original: The outputs of experts are recast back to the original dimension ( d model ) with another three projections.
Revised: The outputs of experts are recast back to the original dimension ( d model ) with another three projections.

Original: 3. T OKEN C LUSTERING
Revised: 3. T OKEN C LUSTERING

Original: The data slicing on the sequence dimension forces more tokens to be dispatched to local devices, which may hurt the performance, as tokens dispatched to local devices are processed by only a few experts inside the devices.
Revised: The data slicing on the sequence dimension forces more tokens to be dispatched to local devices, which may hurt the performance, as tokens dispatched to local devices are processed by only a few experts inside the devices.

Original: To mitigate this issue, we propose to apply token clustering before the
Revised: To mitigate this issue, we propose to apply token clustering before the

Original: SCoMoE layers.
Revised: SCoMoE layers.

Original: The token clustering is to aggregate related tokens to the same device (accelerator/node), where each device is a cluster.
Revised: The token clustering is to aggregate related tokens to the same device (accelerator/node), where each device is a cluster.

Original: Aggregating related tokens together will make it better for the experts to process them.
Revised: Aggregating related tokens together will make it better for the experts to process them.

Original: But the router of MoE usually assigns related tokens to the experts (Zoph et al., 2022), why do we need the token clustering?
Revised: But the router of MoE usually assigns related tokens to the experts (Zoph et al., 2022), why do we need the token clustering?

Original: The reason is that the router of SCoMoE assigns more tokens to the experts inside local devices, while token clustering can mitigate this problem.
Revised: The reason is that the router of SCoMoE assigns more tokens to the experts inside local devices, while token clustering can mitigate this problem.

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

Original: We perform the token clustering before the first SCoMoE layer, and recover the original token distribution after the last SCoMoE layer.
Revised: We perform the token clustering before the first SCoMoE layer, and recover the original token distribution after the last SCoMoE layer.

Original: We average the embeddings of experts inside an accelerator/node as the embedding of the device (either the accelerator or node), and assign tokens to devices in the same way as we assign tokens to experts.
Revised: We average the embeddings of experts inside an accelerator/node as the embedding of the device (either the accelerator or node), and assign tokens to devices in the same way as we assign tokens to experts.

Original: But unfortunately, the commonly used top-2 gating (Lepikhin et al., 2020) and topgating (Fedus et al., 2021)
Revised: But unfortunately, the commonly used top-2 gating (Lepikhin et al., 2020) and topgating (Fedus et al., 2021)

Original: routers could not be used here, since these routers abandon some tokens during routing.
Revised: routers could not be used here, since these routers abandon some tokens during routing.

Original: Therefore we resort to routers that keep all tokens, e.g, greedy router (Lewis et al., 2021) and balanced router (Lewis et al., 2021).
Revised: Therefore we resort to routers that keep all tokens, e.g, greedy router (Lewis et al., 2021) and balanced router (Lewis et al., 2021).

Original: The greedy router allows tokens to be dispatched to the top-1 expert according to the token-to-expert assignment matrix.
Revised: The greedy router allows tokens to be dispatched to the top-1 expert according to the token-to-expert assignment matrix.

Original: However, different from the commonly used top-1 gating (Fedus et al., 2021), the greedy router does not set capacity for each expert, allowing any number of tokens to be sent to the same expert.
Revised: However, different from the commonly used top-1 gating (Fedus et al., 2021), the greedy router does not set capacity for each expert, allowing any number of tokens to be sent to the same expert.

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

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

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

Original: But the auction algorithm usually converges until many iterations, which is inefficient for autoregressive decoding.
Revised: But the auction algorithm usually converges until many iterations, which is inefficient for autoregressive decoding.

Original: Hence the BASE Layer (Lewis et al., 2021) uses the balanced router instead of the greedy router at inference.
Revised: Hence the BASE Layer (Lewis et al., 2021) uses the balanced router instead of the greedy router at inference.

Original: In this work, we perform token clustering with the balanced router.
Revised: In this work, we perform token clustering with the balanced router.

Original: However, we have found a weakness of Base Layer in gradient broadcasting, and proposed Differentiable Sorting to address it, which is described in Appendix A.
Revised: However, we have found a weakness of Base Layer in gradient broadcasting, and proposed Differentiable Sorting to address it, which is described in Appendix A.

Original: The token clustering aggregates the tokens of different sentences together, which cannot be input to the self-attention layer.
Revised: The token clustering aggregates the tokens of different sentences together, which cannot be input to the self-attention layer.

Original: Therefore we have to change the order of Transformer layers, putting the attention layers before the MoE layers.
Revised: Therefore we have to change the order of Transformer layers, putting the attention layers before the MoE layers.

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

Original: The design of the new architecture is due to the following reasons: (1) Press et al.
Revised: The design of the new architecture is due to the following reasons: (1) Press et al.

Original: (2020) have investigated 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 performance.
Revised: (2020) have investigated 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 performance.

Original: (2) Riquelme et al.
Revised: (2) Riquelme et al.

Original: (2021) find that the top MoE layers are more important than the bottom
Revised: (2021) find that the top MoE layers are more important than the bottom

Original: MoE layers.
Revised: MoE layers.

Original: (3) Intuitively, the training of sparse MoE layers is more unstable than the training of the dense layers (Zoph et al., 2022).
Revised: (3) Intuitively, the training of sparse MoE layers is more unstable than the training of the dense layers (Zoph et al., 2022).

Original: Hence, putting the stable dense layers at the bottom may help stabilize the training of the entire architecture.
Revised: Hence, putting the stable dense layers at the bottom may help stabilize the training of the entire architecture.

Original: We put the balanced router before the first SCoMoE layer for token clustering.
Revised: We put the balanced router before the first SCoMoE layer for token clustering.

Original: However, the token clustering also requires the all-to-all communication.
Revised: However, the token clustering also requires the all-to-all communication.

Original: We remove one MoE layer to avoid extra communication cost.
Revised: We remove one MoE layer to avoid extra communication cost.

Original: Hence the number of SCoMoE layers is N/ 2 − 1 rather than N/ 2 , as shown in Figure 2.
Revised: Hence the number of SCoMoE layers is N/ 2 − 1 rather than N/ 2 , as shown in Figure 2.

Original: To keep the same computation cost with the original model, we also insert two fully-connected layers after the balanced router.
Revised: To keep the same computation cost with the original model, we also insert two fully-connected layers after the balanced router.

Original: The balanced router together with the two fully-connected layers act as the first MoE layer.
Revised: The balanced router together with the two fully-connected layers act as the first MoE layer.

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

Original:  E XPERIMENTS 4. E XPERIMENT S ETTINGS
Revised:  E XPERIMENTS 4. E XPERIMENT S ETTINGS

Original: We evaluated the proposed SCoMoE on both bilingual and multilingual neural machine translation.
Revised: We evaluated the proposed SCoMoE on both bilingual and multilingual neural machine translation.

Original: In bilingual NMT experiments, we used the WMT17-En-Fr corpus, which consists of 35,762, training sentences, while in the multilingual NMT experiments, the OPUS-100 corpus (Zhang et al., 2020) was used, which consists of 100 languages and 107,924,846 training examples.
Revised: In bilingual NMT experiments, we used the WMT17-En-Fr corpus, which consists of 35,762, training sentences, while in the multilingual NMT experiments, the OPUS-100 corpus (Zhang et al., 2020) was used, which consists of 100 languages and 107,924,846 training examples.

Original: The details of these two datasets are shown in Table 2 of
Revised: The details of these two datasets are shown in Table 2 of

Original: Appendix B. We used the script of fairseq 4 to download and pre-process the WMT17-En-Fr dataset while the OPUS-100 dataset was directly downloaded from the OPUS web.
Revised: Appendix B. We used the script of fairseq 4 to download and pre-process the WMT17-En-Fr dataset while the OPUS-100 dataset was directly downloaded from the OPUS web.

Original: Both the WMT17-En-Fr and OPUS-100 dataset were tokenized via the BPE algorithm (Sennrich et al., 2016).
Revised: Both the WMT17-En-Fr and OPUS-100 dataset were tokenized via the BPE algorithm (Sennrich et al., 2016).

Original: For the multilingual experiments, a language token indicating the source language was prepended to each source sentence and similarly, a target language token was prepended to each target sentence.
Revised: For the multilingual experiments, a language token indicating the source language was prepended to each source sentence and similarly, a target language token was prepended to each target sentence.

Original: Models were trained in two settings: base model and large model.
Revised: Models were trained in two settings: base model and large model.

Original: The base models have 6 encoder/decoder layers, 8 experts and d model = 512 , which were trained on one node with 8 RTX
Revised: The base models have 6 encoder/decoder layers, 8 experts and d model = 512 , which were trained on one node with 8 RTX

Original: A6000 GPUs.
Revised: A6000 GPUs.

Original: The large models have 12 encoder/decoder layers, 32 experts, and d model = 1024 , which were trained on 4 nodes with 32 RTX A6000 GPUs.
Revised: The large models have 12 encoder/decoder layers, 32 experts, and d model = 1024 , which were trained on 4 nodes with 32 RTX A6000 GPUs.

Original: The 4 nodes were connected with Infiniband, the bandwidth of which is 100 Gbps.
Revised: The 4 nodes were connected with Infiniband, the bandwidth of which is 100 Gbps.

Original: The number of parameters in the base model is 0.
Revised: The number of parameters in the base model is 0.

Original: billion while 4.1 billion in the large model.
Revised: billion while 4.1 billion in the large model.

Original: The details of the base and large model configuration could be found in Table 3 of Appendix B.
Revised: The details of the base and large model configuration could be found in Table 3 of Appendix B.

Original: We compared our method with Gshard and its variant with a small capacity factor on translation performance and training speed.
Revised: We compared our method with Gshard and its variant with a small capacity factor on translation performance and training speed.

Original: We refer to our SCoMoE on the sequence dimension as SCoMoESeq, and that on the feature dimension as SCoMoE-Feat.
Revised: We refer to our SCoMoE on the sequence dimension as SCoMoESeq, and that on the feature dimension as SCoMoE-Feat.

Original: The Gshard with a small capacity factor is referred to as Small-Cap-MoE. We set the ratio of tokens dispatched through the intra-accelerator communication to be α ( 0 < α < 1 ), and the ratio through the intra-node communication to be β ( 0 < β < 1 ), hence that through the global communication is γ = 1 − α − β ( 0 < γ < 1 ).
Revised: The Gshard with a small capacity factor is referred to as Small-Cap-MoE. We set the ratio of tokens dispatched through the intra-accelerator communication to be α ( 0 < α < 1 ), and the ratio through the intra-node communication to be β ( 0 < β < 1 ), hence that through the global communication is γ = 1 − α − β ( 0 < γ < 1 ).

Original: To simplify the comparison, we set either α or β to be non-zero, and the other to be 0.
Revised: To simplify the comparison, we set either α or β to be non-zero, and the other to be 0.

Original: We refer the SCoMoE-(Seq/Feat) with α > 0 , β = 0 , γ > 0 as SCoMoE-(Seq/Feat)- α , while that with α = 0 , β > 0 , γ > 0 as SCoMoE-(Seq/Feat)- β .
Revised: We refer the SCoMoE-(Seq/Feat) with α > 0 , β = 0 , γ > 0 as SCoMoE-(Seq/Feat)- α , while that with α = 0 , β > 0 , γ > 0 as SCoMoE-(Seq/Feat)- β .

Original: Please refer to Appendix B for more details about experiment settings.
Revised: Please refer to Appendix B for more details about experiment settings.

Original: 4. M AIN R ESULTS
Revised: 4. M AIN R ESULTS

Original: We evaluated the proposed methods on the WMT17-En-Fr and OPUS-100.
Revised: We evaluated the proposed methods on the WMT17-En-Fr and OPUS-100.

Original: We report the BLEU score on the WMT17-En-Fr, and the average BLEU scores over all, English-to-any, and any-toEnglish language pairs on the OPUS-100.
Revised: We report the BLEU score on the WMT17-En-Fr, and the average BLEU scores over all, English-to-any, and any-toEnglish language pairs on the OPUS-100.

Original: Results are shown in Table 1.
Revised: Results are shown in Table 1.

Original: We also report the win-rate and the speedup of all models over Gshard.
Revised: We also report the win-rate and the speedup of all models over Gshard.

Original: We evaluated the training speed on the dummy-mt task 6 of fairseq, since the training speed evaluated on it is more stable.
Revised: We evaluated the training speed on the dummy-mt task 6 of fairseq, since the training speed evaluated on it is more stable.

Original: Each input sentence in the dummy-mt task is a series of numbers, e.g., “0, 1, 2, 3, 4, ...”, rather than the real sentence.
Revised: Each input sentence in the dummy-mt task is a series of numbers, e.g., “0, 1, 2, 3, 4, ...”, rather than the real sentence.

Original: 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 to train each model).
Revised: 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 to train each model).

Original: Base Model
Revised: Base Model

Original: The base SCoMoE was trained on a single node, hence we only need to consider the ratio for the intra-accelerator communication ( α ).
Revised: The base SCoMoE was trained on a single node, hence we only need to consider the ratio for the intra-accelerator communication ( α ).

Original: For a fair comparison on translation performance, we set α for SCoMoE-Seq to 0.3, and that for SCoMoE-Feat to 0.5, the capacity factor of Small-CapMoE to 0.8, so that these models can be trained with a comparable training speed.
Revised: For a fair comparison on translation performance, we set α for SCoMoE-Seq to 0.3, and that for SCoMoE-Feat to 0.5, the capacity factor of Small-CapMoE to 0.8, so that these models can be trained with a comparable training speed.

Original: More evaluation on the training speed of SCoMoE can be found in Appendix D
Revised: More evaluation on the training speed of SCoMoE can be found in Appendix D

Original: From Table 1, we can observe that both SCoMoE-Seq/Feat and Small-Cap-MoE are faster than
Revised: From Table 1, we can observe that both SCoMoE-Seq/Feat and Small-Cap-MoE are faster than

Original: Gshard by nearly 10%.
Revised: Gshard by nearly 10%.

Original: But SCoMoE cost less communication time than Small-Cap-MoE. Why  do they still have similar training speed?
Revised: But SCoMoE cost less communication time than Small-Cap-MoE. Why [Table] do they still have similar training speed?

Original: The reason is that SCoMoE brings some extra overhead on computation, which is analyzed in Appendix D. Although SCoMoE-Feat is slightly slower than Small-Cap-MoE, it is better in terms of BLEU on both the WMT17-En-Fr and OPUSdataset.
Revised: The reason is that SCoMoE brings some extra overhead on computation, which is analyzed in Appendix D. Although SCoMoE-Feat is slightly slower than Small-Cap-MoE, it is better in terms of BLEU on both the WMT17-En-Fr and OPUSdataset.

Original: While SCoMoE-Seq is comparable or even better than Gshard on both the WMT17-En-Fr and OPUS-100 dataset, due to the proposed token clustering and differentiable sorting.
Revised: While SCoMoE-Seq is comparable or even better than Gshard on both the WMT17-En-Fr and OPUS-100 dataset, due to the proposed token clustering and differentiable sorting.

Original: We further demonstrate the effectiveness of the two methods in Appendix E and F.[Table] [Figure]
Revised: We further demonstrate the effectiveness of the two methods in Appendix E and F.

Original: Large Model
Revised: Large Model

Original: The large SCoMoE models were trained on multiple nodes, hence both SCoMoESeq- α and - β were evaluated.
Revised: The large SCoMoE models were trained on multiple nodes, hence both SCoMoESeq- α and - β were evaluated.

Original: We set α = 0 .
Revised: We set α = 0 .

Original: 2 for SCoMoE-Seq-
Revised: 2 for SCoMoE-Seq-

Original: α and 0.5 for SCoMoE-
Revised: α and 0.5 for SCoMoE-

Original: Feat- α , which have the comparable training speed with SCoMoE-Seq/Feat- β with β = 0 .
Revised: Feat- α , which have the comparable training speed with SCoMoE-Seq/Feat- β with β = 0 .

Original: We set the capacity factor of Small-Cap-MoE to 1.0, while that of other models was set to 1.25.
Revised: We set the capacity factor of Small-Cap-MoE to 1.0, while that of other models was set to 1.25.

Original: From Table 1, it is clear that all variants of SCoMoE are faster and consume less communication time than Gshard significantly.
Revised: From Table 1, it is clear that all variants of SCoMoE are faster and consume less communication time than Gshard significantly.

Original: Among the 4 variants, SCoMoE-
Revised: Among the 4 variants, SCoMoE-

Original: Feat- α is the fastest.
Revised: Feat- α is the fastest.

Original: It is reasonable since we set α of SCoMoE-
Revised: It is reasonable since we set α of SCoMoE-

Original: Feat- α to 0.5, but that for SCoMoE-Seq-
Revised: Feat- α to 0.5, but that for SCoMoE-Seq-

Original: α to only 0.3.
Revised: α to only 0.3.

Original: We provide the training speed comparison of SCoMoE with more settings of α and β in Appendix D. It can also be seen that
Revised: We provide the training speed comparison of SCoMoE with more settings of α and β in Appendix D. It can also be seen that

Original: SCoMoE-Seq-
Revised: SCoMoE-Seq-

Original: α is faster than SCoMoE-Seq-
Revised: α is faster than SCoMoE-Seq-

Original: β
Revised: β

Original: but takes more communication time.
Revised: but takes more communication time.

Original: This is because that SCoMoE uses the top-2 routing, where each token is computed in two experts.
Revised: This is because that SCoMoE uses the top-2 routing, where each token is computed in two experts.

Original: In comparison, tokens for intra-accelerator communication are computed in a single expert inside the accelerator, which reduces computational time.
Revised: In comparison, tokens for intra-accelerator communication are computed in a single expert inside the accelerator, which reduces computational time.

Original: We have also found that all variants of SCoMoE achieve better translation performance than both
Revised: We have also found that all variants of SCoMoE achieve better translation performance than both

Original: Gshard and its variant Small-Cap-MoE on the OPUS-100 dataset.
Revised: Gshard and its variant Small-Cap-MoE on the OPUS-100 dataset.

Original: They also perform comparably to Gshard on the WMT17-En-Fr dataset.
Revised: They also perform comparably to Gshard on the WMT17-En-Fr dataset.

Original: Among the 4 variants of SCoMoE, SCoMoE-Seq- β is the best in terms of BLEU on both OPUS-100 and WMT17-En-Fr.
Revised: Among the 4 variants of SCoMoE, SCoMoE-Seq- β is the best in terms of BLEU on both OPUS-100 and WMT17-En-Fr.

Original: It is surprising that SCoMoESeq- β outperforms Gshard by 2.8 BLEU on OPUS-100, but only 0.1 BLEU on WMT17-En-Fr.
Revised: It is surprising that SCoMoESeq- β outperforms Gshard by 2.8 BLEU on OPUS-100, but only 0.1 BLEU on WMT17-En-Fr.

Original: We conjecture that MoE models are easy to overfit on low-resource languages (Costa-juss`a et al., 2022), while our method can mitigate this issue by regularizing the routing of MoE. We have shown the improvements of SCoMoE on different languages in Appendix C to validate this hypothesis.
Revised: We conjecture that MoE models are easy to overfit on low-resource languages (Costa-juss`a et al., 2022), while our method can mitigate this issue by regularizing the routing of MoE. We have shown the improvements of SCoMoE on different languages in Appendix C to validate this hypothesis.

Original: Overall, all variants of SCoMoE are faster than Gshard with comparable or better translation performance.
Revised: Overall, all variants of SCoMoE are faster than Gshard with comparable or better translation performance.

Original: The best settings of α/β for SCoMoE From Table 1, we can roughly see that SCoMoE-Seq/Featβ have better translation performance, while SCoMoE-Seq/
Revised: The best settings of α/β for SCoMoE From Table 1, we can roughly see that SCoMoE-Seq/Featβ have better translation performance, while SCoMoE-Seq/

Original: Feat- α are faster.
Revised: Feat- α are faster.

Original: We further compared the translation performance of them with different settings of α or β .
Revised: We further compared the translation performance of them with different settings of α or β .

Original: The results are shown in Figure 3a and 3b.
Revised: The results are shown in Figure 3a and 3b.

Original: According to Figure 3a,
Revised: According to Figure 3a,

Original: SCoMoE-Seq-
Revised: SCoMoE-Seq-

Original: α is better with a smaller α in terms of BLEU, which is not surprising since increasing α will reduce the amount of data for global communication.
Revised: α is better with a smaller α in terms of BLEU, which is not surprising since increasing α will reduce the amount of data for global communication.

Original: But SCoMoE-MoE- β achieves better translation performance with larger β s.
Revised: But SCoMoE-MoE- β achieves better translation performance with larger β s.

Original: Furthermore,
Revised: Furthermore,

Original: SCoMoE-MoE- β is even comparable to Gshard when β
Revised: SCoMoE-MoE- β is even comparable to Gshard when β

Original: = 0 .
Revised: = 0 .

Original: 9 , suggesting that only a small proportion (10%) of tokens need to be communicated in global communication and that intra-node communication is sufficient for the majority of tokens.
Revised: 9 , suggesting that only a small proportion (10%) of tokens need to be communicated in global communication and that intra-node communication is sufficient for the majority of tokens.

Original: SCoMoE-Seq- β with small β s underperforms
Revised: SCoMoE-Seq- β with small β s underperforms

Original: Gshard, the reason for which can be found in Table 4 where Gshard also suffers from performance drop with token clustering and reordering.
Revised: Gshard, the reason for which can be found in Table 4 where Gshard also suffers from performance drop with token clustering and reordering.

Original:  With regard to the comparison on the training speed (Figure 3b), the speedup of SCoMoE-seqα/β over Gshard increases nearly linearly with the increment of α/β .
Revised: [Figure] With regard to the comparison on the training speed (Figure 3b), the speedup of SCoMoE-seqα/β over Gshard increases nearly linearly with the increment of α/β .

Original: The rate of increment for
Revised: The rate of increment for

Original: SCoMoE-seq- α is around 1.0, while that for SCoMoE-seq- β is 0.5.
Revised: SCoMoE-seq- α is around 1.0, while that for SCoMoE-seq- β is 0.5.

Original: Therefore, to achieve the best speedup without performance drop in BLEU, we can set α of SCoMoE-seq- α to 0.3 and β of SCoMoE-seq- β to 0.9, which can lead to a speedup of 1.28x and 1.44x respectively.
Revised: Therefore, to achieve the best speedup without performance drop in BLEU, we can set α of SCoMoE-seq- α to 0.3 and β of SCoMoE-seq- β to 0.9, which can lead to a speedup of 1.28x and 1.44x respectively.

Original: We can also set β of SCoMoE-seq- β to 0.5, which achieves the highest BLEU scores and a speedup of 1.25x.
Revised: We can also set β of SCoMoE-seq- β to 0.5, which achieves the highest BLEU scores and a speedup of 1.25x.

Original: Scaling in model and batch size The speedup of SCoMoE is also dependent on the model size and batch size.
Revised: Scaling in model and batch size The speedup of SCoMoE is also dependent on the model size and batch size.

Original: Figure 3c and 3d compares the training speed of SCoMoE with different number of experts and batch sizes.
Revised: Figure 3c and 3d compares the training speed of SCoMoE with different number of experts and batch sizes.

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

Original: It is due to that the communication cost is more expensive with more devices.
Revised: It is due to that the communication cost is more expensive with more devices.

Original: Similarly, SCoMoE also has a more significant speedup with larger batch sizes, as shown in Figure 3d.
Revised: Similarly, SCoMoE also has a more significant speedup with larger batch sizes, as shown in Figure 3d.

Original: In Figure 3d, we also display the training speed of Gshard with the hierarchical all-to-all communication proposed in
Revised: In Figure 3d, we also display the training speed of Gshard with the hierarchical all-to-all communication proposed in

Original: Tutel (Hwang et al., 2022).
Revised: Tutel (Hwang et al., 2022).

Original: However, the hierarchical communication of Tutel even slows down the training.
Revised: However, the hierarchical communication of Tutel even slows down the training.

Original: The reason may be that the intra-node bandwidth in our experiments is not fast (PCIe) and that the number of devices is not very large (32).
Revised: The reason may be that the intra-node bandwidth in our experiments is not fast (PCIe) and that the number of devices is not very large (32).

Original: Other Experiments and Analysis
Revised: Other Experiments and Analysis

Original: • We analyze why SCoMoE can achieve a substantial improvement on translation performance on OPUS-100 in Appendix C. • We conduct the empirical evaluation and theoretical analysis on the training speed and communication time cost of SCoMoE in Appendix D. •
Revised: • We analyze why SCoMoE can achieve a substantial improvement on translation performance on OPUS-100 in Appendix C. • We conduct the empirical evaluation and theoretical analysis on the training speed and communication time cost of SCoMoE in Appendix D. •

Original: The effect of token clustering and differentiable sorting on the performance of SCoMoE can be found in Appendix E and F. • The analysis on the sensitiveness of MoE models to the capacity factor at inference is provided in Appendix.
Revised: The effect of token clustering and differentiable sorting on the performance of SCoMoE can be found in Appendix E and F. • The analysis on the sensitiveness of MoE models to the capacity factor at inference is provided in Appendix.

Original: G
Revised: G

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we have presented SCoMoE to reduce the all-to-all communication time of MoEs, by encouraging more data to dispatch through the fast intra-accelerator/node communication channel, and less data through the slow inter-node communication channel.
Revised: In this paper, we have presented SCoMoE to reduce the all-to-all communication time of MoEs, by encouraging more data to dispatch through the fast intra-accelerator/node communication channel, and less data through the slow inter-node communication channel.

Original: By slicing/projecting the data on the sequence/feature dimension, we transform the data into three groups corresponding to three levels of communication.
Revised: By slicing/projecting the data on the sequence/feature dimension, we transform the data into three groups corresponding to three levels of communication.

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

Original: Experiments on massive bilingual and multilingual neural machine translation demonstrate that our
Revised: Experiments on massive bilingual and multilingual neural machine translation demonstrate that our

Original: SCoMoE speeds up Gshard significantly with comparable or even better performance.
Revised: SCoMoE speeds up Gshard significantly with comparable or even better performance.

Original: The empirical evaluation and theoretical analysis on the training speed and communication time cost discloses the speedups achieved by our method under different setting and network bandwidth.
Revised: The empirical evaluation and theoretical analysis on the training speed and communication time cost discloses the speedups achieved by our method under different setting and network bandwidth.

Original: The ablation study demonstrates the benefit of the proposed token clustering and differentiable sorting.
Revised: The ablation study demonstrates the benefit of the proposed token clustering and differentiable sorting.

Original: A PPENDIX A D IFFERENTIABLE S ORTING
Revised: A PPENDIX A D IFFERENTIABLE S ORTING

Original: The auction algorithm in the balanced router is not differentiable, hence the BASE Layer (Lewis et al., 2021) learns the expert embedding via the sigmoid gate (Srivastava et al., 2015):
Revised: The auction algorithm in the balanced router is not differentiable, hence the BASE Layer (Lewis et al., 2021) learns the expert embedding via the sigmoid gate (Srivastava et al., 2015):

Original: [Equation]
Revised: [Equation]

Original: [Figure]

Original:  where the output of the MoE layer X
Revised: (4) where the output of the MoE layer X

Original: ′ i is the combination of the expert input
Revised: ′ i is the combination of the expert input

Original: X i and expert output expert j
Revised: X i and expert output expert j

Original: ( X i ) that are integrated with the normalized assignment score M ij .
Revised: ( X i ) that are integrated with the normalized assignment score M ij .

Original: The gradients are broadcast through the sigmoid gate to update the parameters of the expert embeddings.
Revised: The gradients are broadcast through the sigmoid gate to update the parameters of the expert embeddings.

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

Original: With the gradients from the sigmoid gate, we only know how to update the assignment scores of the activated experts, which is not available for other experts.
Revised: With the gradients from the sigmoid gate, we only know how to update the assignment scores of the activated experts, which is not available for other experts.

Original: A straightforward solution to this is to replace the sigmoid gate with the softmax gate.
Revised: A straightforward solution to this is to replace the sigmoid gate with the softmax gate.

Original: However, the softmax operator requires multiple inputs, but only one expert is activated for each token.
Revised: However, the softmax operator requires multiple inputs, but only one expert is activated for each token.

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

Original: Each row of the matrix could be seen as selecting a token from the sequence for an expert.
Revised: Each row of the matrix could be seen as selecting a token from the sequence for an expert.

Original: We implement the sorting via linear projection:
Revised: We implement the sorting via linear projection:

Original: [Equation]
Revised: [Equation]

Original: where the sorted sequence X sort is obtained by the projection W sort on the input sequence X .
Revised: where the sorted sequence X sort is obtained by the projection W sort on the input sequence X .

Original: If the number of sorted tokens is n , and the number of experts is m , n/m tokens will be assigned to each expert.
Revised: If the number of sorted tokens is n , and the number of experts is m , n/m tokens will be assigned to each expert.

Original: We visualize the sorting with a 4 × 4 sparse projection on a sequence of length 4 in Figure 4.
Revised: We visualize the sorting with a 4 × 4 sparse projection on a sequence of length 4 in Figure 4.

Original: The tokens in the sorted sequence are assigned to two experts.
Revised: The tokens in the sorted sequence are assigned to two experts.

Original: The assigned tokens often have the 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.
Revised: The assigned tokens often have the 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.

Original: The softmax is the approximation to the argmax, if we set the temperature to a small value.
Revised: The softmax is the approximation to the argmax, if we set the temperature to a small value.

Original: Therefore, we share the gradients of the projection matrix with the softmax matrix via the straight-through trick (Bengio et al., 2013):
Revised: Therefore, we share the gradients of the projection matrix with the softmax matrix via the straight-through trick (Bengio et al., 2013):

Original: [Equation]
Revised: [Equation]

Original: W sort is the combination of the sorting projection W auction from the auction algorithm, W softmax , which is the normalized assignment matrix, and detach( W softmax ) , which has the same value as
Revised: W sort is the combination of the sorting projection W auction from the auction algorithm, W softmax , which is the normalized assignment matrix, and detach( W softmax ) , which has the same value as

Original: W softmax but cannot broadcast the gradients.
Revised: W softmax but cannot broadcast the gradients.

Original: τ is the temperature which controls the sharpness of the softmax distribution.
Revised: τ is the temperature which controls the sharpness of the softmax distribution.

Original: The lower the temperature is, the better the softmax approximates the argmax.
Revised: The lower the temperature is, the better the softmax approximates the argmax.

Original: W sort is equal to W auction in the forward computation.
Revised: W sort is equal to W auction in the forward computation.

Original: W softmax has the same gradients as W sort in the backward propagation, although W softmax does not have influence on the training loss.
Revised: W softmax has the same gradients as W sort in the backward propagation, although W softmax does not have influence on the training loss.

Original: As the shape of the probability matrix W softmax is different from that of W sort , ( m, n ) vs. ( n, m ), we cannot add them directly.
Revised: As the shape of the probability matrix W softmax is different from that of W sort , ( m, n ) vs. ( n, m ), we cannot add them directly.

Original: Each row of W softmax is a set of probabilities assigning tokens to an expert, which are corresponding to n/m rows of W sort as we need to assign n/m tokens to each expert according to the same distribution.
Revised: Each row of W softmax is a set of probabilities assigning tokens to an expert, which are corresponding to n/m rows of W sort as we need to assign n/m tokens to each expert according to the same distribution.

Original: For the example in Figure 4, both the first and second row of sparse matrix W auction are corresponding to the first row of the probability matrix W softmax .
Revised: For the example in Figure 4, both the first and second row of sparse matrix W auction are corresponding to the first row of the probability matrix W softmax .

Original: We reshape W softmax from ( m, n ) to ( n, n ) by repeating each row of the matrix to n/m rows.
Revised: We reshape W softmax from ( m, n ) to ( n, n ) by repeating each row of the matrix to n/m rows.

Original: In Figure 4, we reshape the matrix from (2,4) to (4,4), by repeating the first and second row of the matrix
Revised: In Figure 4, we reshape the matrix from (2,4) to (4,4), by repeating the first and second row of the matrix

Original: W softmax to the first two and second two rows of the broadcast matrix.
Revised: W softmax to the first two and second two rows of the broadcast matrix.

Original: B D ETAILS ON E XPERIMENT S ETTINGS
Revised: B D ETAILS ON E XPERIMENT S ETTINGS

Original: [Table]
Revised: [Table]

Revised: Setting

Revised: Enc-Layer

Revised: Dec-Layer d model

Original: [Table]
Revised: Experts GPU Nodes Parameters (B)

Original:  We used the MoE branch of fairseq 7 to implement our MoE models.
Revised: Base      0. Large      4. We used the MoE branch of fairseq 7 to implement our MoE models.

Original: We prevented the padding tokens from being sent to experts by masking their assignment scores.
Revised: We prevented the padding tokens from being sent to experts by masking their assignment scores.

Original: The capacity factor of the base model was set to 1, while that of the large model was set to 1.25 following the suggestion of (Zoph et al., 2022).
Revised: The capacity factor of the base model was set to 1, while that of the large model was set to 1.25 following the suggestion of (Zoph et al., 2022).

Original: The tokens with small assignment scores to an expert were dropped if the expert was out-of-capacity.
Revised: The tokens with small assignment scores to an expert were dropped if the expert was out-of-capacity.

Original: All models used pre-normalization to stabilize the training.
Revised: All models used pre-normalization to stabilize the training.

Original: We have 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 to be no more than 0.5) and gradient accumulation (accumulating the gradients every 4 steps).
Revised: We have 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 to be no more than 0.5) and gradient accumulation (accumulating the gradients every 4 steps).

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

Original: We have also found that the MoE model trained in multilingual experiments could not perform well at inference, and that increasing the load-balance loss weight from 0.01 (default value) to 0.1 significantly improves the performance.
Revised: We have also found that the MoE model trained in multilingual experiments could not perform well at inference, and that increasing the load-balance loss weight from 0.01 (default value) to 0.1 significantly improves the performance.

Original: Therefore, we set the loadbalance loss weight to 0.1 in the multilingual experiments, and 0.01 in the bilingual experiments.
Revised: Therefore, we set the loadbalance loss weight to 0.1 in the multilingual experiments, and 0.01 in the bilingual experiments.

Original: The Adam optimizer (Kingma & Ba, 2015a) with learning rate = 5 e − 4 , α = 0 .
Revised: The Adam optimizer (Kingma & Ba, 2015a) with learning rate = 5 e − 4 , α = 0 .

Original: 9 , β = 0 .
Revised: 9 , β = 0 .

Original: 98 , was used to optimize our models.
Revised: 98 , was used to optimize our models.

Original: We employed the inverse-square-root learning schedule (Kingma &
Revised: We employed the inverse-square-root learning schedule (Kingma &

Original: Ba, 2015b), with the number of warm-up step being set to 4000.
Revised: Ba, 2015b), with the number of warm-up step being set to 4000.

Original: The batch size was set to betokens.
Revised: The batch size was set to betokens.

Original: The dropout rate was 0.3 for bilingual experiments and 0.1 for multilingual experiments.
Revised: The dropout rate was 0.3 for bilingual experiments and 0.1 for multilingual experiments.

Original: Both the base and large models were trained for 10 epochs.
Revised: Both the base and large models were trained for 10 epochs.

Original: The checkpoint with the best performance on the validation set was selected for testing.
Revised: The checkpoint with the best performance on the validation set was selected for testing.

Original: In the bilingual experiments, we evaluated the performance on the validation set with BLEU.
Revised: In the bilingual experiments, we evaluated the performance on the validation set with BLEU.

Original: In the multilingual experiments, the perplexity (ppl) on the validation set was used for faster validation.
Revised: In the multilingual experiments, the perplexity (ppl) on the validation set was used for faster validation.

Original: The BLEU score was computed via sacrebleu (Post, 2018).
Revised: The BLEU score was computed via sacrebleu (Post, 2018).

Original: [Figure]
Revised: [Figure]

Revised: a rich-/low-resource language.

Original: C W HY SC O M O E IS BETTER ON M ULTILINGUAL M ACHINE T RANSLATION
Revised: C W HY SC O M O E IS BETTER ON M ULTILINGUAL M ACHINE T RANSLATION

Original: In Table 1, the improvement of SCoMoE-Seq- β on multilingual OPUS-100 is much more significant than that on bilingual WMT17-En-Fr.
Revised: In Table 1, the improvement of SCoMoE-Seq- β on multilingual OPUS-100 is much more significant than that on bilingual WMT17-En-Fr.

Original: We plot the convergence curves of Gshard and SCoMoE on validation set of OPUS-100 and WMT17-En-Fr in Figure 6a and 6b respectively.
Revised: We plot the convergence curves of Gshard and SCoMoE on validation set of OPUS-100 and WMT17-En-Fr in Figure 6a and 6b respectively.

Original: It is clear that the convergence curves of SCoMoE are more smooth than Gshard on both OPUS-100 and WMT17-EnFr.
Revised: It is clear that the convergence curves of SCoMoE are more smooth than Gshard on both OPUS-100 and WMT17-EnFr.

Original: Gshard obtains its best perplexity on OPUS-100 at the 20kth step, but its perplexity increases as  the training continues.
Revised: Gshard obtains its best perplexity on OPUS-100 at the 20kth step, but its perplexity increases as [Figure] the training continues.

Original: While the convergence curve of Gshard on WMT17-En-Fr is also fluctuated, it tends towards getting better.
Revised: While the convergence curve of Gshard on WMT17-En-Fr is also fluctuated, it tends towards getting better.

Original: Therefore, Gshard only overfits on OPUS-100.
Revised: Therefore, Gshard only overfits on OPUS-100.

Original: This is reasonable, since WMT17-En-Fr contains much more data than any single language pair in OPUS-100 does.
Revised: This is reasonable, since WMT17-En-Fr contains much more data than any single language pair in OPUS-100 does.

Original: In Figure 5, we illustrate the data distribution of different language pairs in OPUS-100.
Revised: In Figure 5, we illustrate the data distribution of different language pairs in OPUS-100.

Original: It is clear that only half of language pairs have about 1 million parallel sentence pairs, while other language pairs have much less data.
Revised: It is clear that only half of language pairs have about 1 million parallel sentence pairs, while other language pairs have much less data.

Original: We also show the improvements of SCoMoE over Gshard on different language pairs of OPUS-100 in Figure 7, where we visualize different language pairs in different colors according to whether the language is low-resource or rich-resource.
Revised: We also show the improvements of SCoMoE over Gshard on different language pairs of OPUS-100 in Figure 7, where we visualize different language pairs in different colors according to whether the language is low-resource or rich-resource.

Original: We define a language pair as a low/rich-resource language pair according to whether the data size of this language pair is larger than the median size (640k) of all language pairs.
Revised: We define a language pair as a low/rich-resource language pair according to whether the data size of this language pair is larger than the median size (640k) of all language pairs.

Original: From Figure 7, we can observe that the large improvements ( > 8 BLEU) of SCoMoE are all achieved on the low-resource language pairs.
Revised: From Figure 7, we can observe that the large improvements ( > 8 BLEU) of SCoMoE are all achieved on the low-resource language pairs.

Original: Among these language pairs, en-af only has 274,821 sentence pairs, but gets the greatest improvement of nearly 16 BLEU.
Revised: Among these language pairs, en-af only has 274,821 sentence pairs, but gets the greatest improvement of nearly 16 BLEU.

Original: All of these suggest that SCoMoE alleviates the overfitting issue of Gshard especially on low-resource language pairs.
Revised: All of these suggest that SCoMoE alleviates the overfitting issue of Gshard especially on low-resource language pairs.

Original: D E VALUATION ON T RAINING S PEED AND C OMMUNICATION T IME
Revised: D E VALUATION ON T RAINING S PEED AND C OMMUNICATION T IME

Original: In this section, we compared the training speed and communication time of SCoMoE-Seq and
Revised: In this section, we compared the training speed and communication time of SCoMoE-Seq and

Original: SCoMoE-Feat with different settings of α and β .
Revised: SCoMoE-Feat with different settings of α and β .

Original: We report the training speed in the metric of words-per-second (WPS), which is the number of tokens that the model can process per second, and the time cost of each all-to-all communication in milliseconds.
Revised: We report the training speed in the metric of words-per-second (WPS), which is the number of tokens that the model can process per second, and the time cost of each all-to-all communication in milliseconds.

Original: [Figure] Base Model We tuned the ratio of the intra-accelerator communication ( α ) from 0 to 0.5, and illustrated the training speed and communication time of SCoMoE at different ratios in Figure 8a and 8b respectively.
Revised: Base Model We tuned the ratio of the intra-accelerator communication ( α ) from 0 to 0.5, and illustrated the training speed and communication time of SCoMoE at different ratios in Figure 8a and 8b respectively.

Original: In spite of training with fixed dummy inputs, the training speed of the base models is still unstable.
Revised: In spite of training with fixed dummy inputs, the training speed of the base models is still unstable.

Original: Hence we trained the same model for three times and show the mean and the standard deviation over the three runs in Figure 8a and 8b.
Revised: Hence we trained the same model for three times and show the mean and the standard deviation over the three runs in Figure 8a and 8b.

Original: From Figure 8a, we observe that the two models can be trained faster as the ratio of intra-accelerator communication ( α ) increases, except that the training of both SCoMoE-Seq and SCoMoE-Feat becomes slower when the ratio is 0.1, but with less communication time than when the ratio is 0.
Revised: From Figure 8a, we observe that the two models can be trained faster as the ratio of intra-accelerator communication ( α ) increases, except that the training of both SCoMoE-Seq and SCoMoE-Feat becomes slower when the ratio is 0.1, but with less communication time than when the ratio is 0.

Original: This suggests that the structured communication reduces the communication time but may cause additional overhead.
Revised: This suggests that the structured communication reduces the communication time but may cause additional overhead.

Original: The extra overhead is probably from the computation, since computations in each expert are also split into multiple groups corresponding to multiple communication groups, which is slower than performing all computations at the same time.
Revised: The extra overhead is probably from the computation, since computations in each expert are also split into multiple groups corresponding to multiple communication groups, which is slower than performing all computations at the same time.

Original: Even if we set the ratio of the intraaccelerator communication to 0, SCoMoE-Seq is still significantly faster than Gshard.
Revised: Even if we set the ratio of the intraaccelerator communication to 0, SCoMoE-Seq is still significantly faster than Gshard.

Original: The reason is that the communication cost of token clustering is half as that of the top-2 routing.
Revised: The reason is that the communication cost of token clustering is half as that of the top-2 routing.

Original: We also notice that the SCoMoE-Seq is always faster than SCoMoE-Feat, and requires less communication time.
Revised: We also notice that the SCoMoE-Seq is always faster than SCoMoE-Feat, and requires less communication time.

Original: This is because the token clustering is only applied to the SCoMoE-Seq.
Revised: This is because the token clustering is only applied to the SCoMoE-Seq.

Original: Large Model Both SCoMoE-Seq/
Revised: Large Model Both SCoMoE-Seq/

Original: Feat- α and - β can be used for training large models.
Revised: Feat- α and - β can be used for training large models.

Original: The differences between them is that SCoMoE-Seq/Feat- β accelerate the training speed with the intra-node communication instead of the intra-accelerator communication.
Revised: The differences between them is that SCoMoE-Seq/Feat- β accelerate the training speed with the intra-node communication instead of the intra-accelerator communication.

Original: We have compared the training speed and communication time of them with different settings of α / β in Figure 8c and 8d respectively.
Revised: We have compared the training speed and communication time of them with different settings of α / β in Figure 8c and 8d respectively.

Original: Overall, all variants of SCoMoE train faster and require less communication time as α or β increases.
Revised: Overall, all variants of SCoMoE train faster and require less communication time as α or β increases.

Original: It is also clear that SCoMoE-Seq/
Revised: It is also clear that SCoMoE-Seq/

Original: Feat- α is always faster than SCoMoE-Seq/Feat- β .
Revised: Feat- α is always faster than SCoMoE-Seq/Feat- β .

Original: This is not surprising as the intra-accelerator communication is much faster than the intra-node communication.
Revised: This is not surprising as the intra-accelerator communication is much faster than the intra-node communication.

Original: But we also notice that the training speed and communication time of SCoMoE-
Revised: But we also notice that the training speed and communication time of SCoMoE-

Original: Feat- β is not significantly better than Gshard until β reaches 0.5, while that for SCoMoE-
Revised: Feat- β is not significantly better than Gshard until β reaches 0.5, while that for SCoMoE-

Original: Seq- β linearly scales as β increases.
Revised: Seq- β linearly scales as β increases.

Original: The reason for this might be that the amount of communication time is related to the shape of communicated tensor.
Revised: The reason for this might be that the amount of communication time is related to the shape of communicated tensor.

Original: Theoretic Analysis on Communication Time We have empirically evaluated the training speed and the amount of communication time of SCoMoE by varying α and β .
Revised: Theoretic Analysis on Communication Time We have empirically evaluated the training speed and the amount of communication time of SCoMoE by varying α and β .

Original: The bandwidth of communication has a great impact on these empirical results in practice.
Revised: The bandwidth of communication has a great impact on these empirical results in practice.

Original: To complement the empirical evaluation, we further theoretically analyze the amount of time required by the all-to-all communication and estimate speedups that can be achieved by the proposed structured communication method.
Revised: To complement the empirical evaluation, we further theoretically analyze the amount of time required by the all-to-all communication and estimate speedups that can be achieved by the proposed structured communication method.

Original: We estimate the all-to-all communication time as the combination of the intra-node, and inter-node communication time as the time cost for the intra-accelerator communication (memory copy) is negligible in comparison to the other two communication channels.
Revised: We estimate the all-to-all communication time as the combination of the intra-node, and inter-node communication time as the time cost for the intra-accelerator communication (memory copy) is negligible in comparison to the other two communication channels.

Original: Suppose that the total amount of data for communication is D MB, and the number of accelerators in one node is m .
Revised: Suppose that the total amount of data for communication is D MB, and the number of accelerators in one node is m .

Original: Given that, the amount of data for the intra-/inter-node communication is Dm / D ( m − 1) m .
Revised: Given that, the amount of data for the intra-/inter-node communication is Dm / D ( m − 1) m .

Original: Suppose that the bandwidth of the intra/ -node communication is B 1 / B 2 MB/s.  amount of the all-to-all communication time (seconds) can be estimated as follows:
Revised: Suppose that the bandwidth of the intra/inter-node communication is B 1 / B 2 MB/s. The amount of the all-to-all communication time (seconds) can be estimated as follows:

Original: [Equation]
Revised: [Equation]

Original: According to the equation, increasing the bandwidth of the inter-node communication ( B 2 ) is more beneficial than increasing the bandwidth of the intra-node communication ( B 1 ) as 1 m − 1 ≤ 1 .
Revised: According to the equation, increasing the bandwidth of the inter-node communication ( B 2 ) is more beneficial than increasing the bandwidth of the intra-node communication ( B 1 ) as 1 m − 1 ≤ 1 .

Original: If we set the ratio of the intra-node communication to β ( α = 0 , 0 < β < 1 ), the amount of the all-to-all communication time can be computed as:
Revised: If we set the ratio of the intra-node communication to β ( α = 0 , 0 < β < 1 ), the amount of the all-to-all communication time can be computed as:

Original: [Equation]
Revised: [Equation]

Original: The absolute value of the derivative of β is proportional to 1 B 2 − 1 B 1 .
Revised: The absolute value of the derivative of β is proportional to 1 B 2 − 1 B 1 .

Original: Either decreasing B 2 or increasing B 1 would make the SCoMoE more efficient.
Revised: Either decreasing B 2 or increasing B 1 would make the SCoMoE more efficient.

Original: Therefore, If the bandwidth of the intra-node communication is slow, it is better to set β = 0 .
Revised: Therefore, If the bandwidth of the intra-node communication is slow, it is better to set β = 0 .

Original: Typically, there are two types of intra-node connections with different bandwidths: PCIE and NVlink/NVSwitch.
Revised: Typically, there are two types of intra-node connections with different bandwidths: PCIE and NVlink/NVSwitch.

Original: The NVlink/NVSwitch connection is much faster than the PCIE connection.
Revised: The NVlink/NVSwitch connection is much faster than the PCIE connection.

Original: Furthermore, the network interface controller (NIC) is also connected with accelerators through PCIE.
Revised: Furthermore, the network interface controller (NIC) is also connected with accelerators through PCIE.

Original: More communications sharing PCIE connections, will make the communication slower.
Revised: More communications sharing PCIE connections, will make the communication slower.

Original: Therefore, we recommend to set β > 0 , α = 0 if fast intra-node connections with NVlink/NVSwitch are available, and α > 0 , β
Revised: Therefore, we recommend to set β > 0 , α = 0 if fast intra-node connections with NVlink/NVSwitch are available, and α > 0 , β

Original: = 0 if slow intra-node connections with PCIE are used.
Revised: = 0 if slow intra-node connections with PCIE are used.

Original: We use the PCIE connections in our experiments, hence the speedup of
Revised: We use the PCIE connections in our experiments, hence the speedup of

Original: SCoMoE- β should be more significant with the fast NVlink/NVSwitch connections.
Revised: SCoMoE- β should be more significant with the fast NVlink/NVSwitch connections.

Original: E E FFECTS OF T OKEN C LUSTERING
Revised: E E FFECTS OF T OKEN C LUSTERING

Original: [Table]

Original: We compared the translation performance and training speed of Gshard and SCoMoE with vs. without token clustering.
Revised: We compared the translation performance and training speed of Gshard and SCoMoE with vs. without token clustering.

Original: Results are shown in Table 4.
Revised: Results are shown in Table 4.

Original: It could be seen that the Gshard with token clustering or layer reordering performs comparably to Gshard without using them on bilingual translation [Table] but is significantly worse than the original GShard on multilingual translation.
Revised: It could be seen that the Gshard with token clustering or layer reordering performs comparably to Gshard without using them on bilingual translation, but is significantly worse than the original GShard on multilingual translation.

Original: Results of
Revised: Results of

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

Original: The reason for that Gshard with token clustering and reordering has the same training speed (1.1x) with SCoMoE are in two folds: 1)
Revised: The reason for that Gshard with token clustering and reordering has the same training speed (1.1x) with SCoMoE are in two folds: 1)

Original: Token clustering speeds up the training of both Gshard and SCoMoE, since the communication cost of token clustering is only half as that of Top2Gate routing.
Revised: Token clustering speeds up the training of both Gshard and SCoMoE, since the communication cost of token clustering is only half as that of Top2Gate routing.

Original: 2) SCoMoE brings additional computational overhead and requires a setting of higher α to be faster than Gshard, which has been discussed in Appendix D.[Figure]
Revised: 2) SCoMoE brings additional computational overhead and requires a setting of higher α to be faster than Gshard, which has been discussed in Appendix D.

Original: F I MPORTANCE OF D IFFERENTIABLE S ORTING
Revised: F I MPORTANCE OF D IFFERENTIABLE S ORTING

Original:  We conducted experiments to investigate how the temperature of softmax affect the performance.
Revised: [Figure] together. We conducted experiments to investigate how the temperature of softmax affect the performance.

Original: We compared the translation performance of SCoMoE-Seq with different temperatures.
Revised: We compared the translation performance of SCoMoE-Seq with different temperatures.

Original: Results are plotted in Figure 9a, from which we could see that the performance of the model is sensitive to the temperature.
Revised: Results are plotted in Figure 9a, from which we could see that the performance of the model is sensitive to the temperature.

Original: We also report the performance of SCoMoE-Seq with the sigmoid gate as the baseline, which is also displayed in Figure 9a (the horizontal dashed line).
Revised: We also report the performance of SCoMoE-Seq with the sigmoid gate as the baseline, which is also displayed in Figure 9a (the horizontal dashed line).

Original: It could be seen that SCoMoE-Seq with the softmax gate is consistently better than that with the sigmoid gate at different temperatures.
Revised: It could be seen that SCoMoE-Seq with the softmax gate is consistently better than that with the sigmoid gate at different temperatures.

Original: G E FFECT OF C APACITY F ACTOR AT I NFERENCE
Revised: G E FFECT OF C APACITY F ACTOR AT I NFERENCE

Original: The capacity factor is of importance to the training of MoE models (Yang et al., 2021).
Revised: The capacity factor is of importance to the training of MoE models (Yang et al., 2021).

Original: In our multilingual NMT experiments, we have found that the models trained on the OPUS-100 dataset performs poorly on the test set, unless with a high capacity factor, substantially higher than that  used during training.
Revised: In our multilingual NMT experiments, we have found that the models trained on the OPUS-100 dataset performs poorly on the test set, unless with a high capacity factor, substantially higher than that [Table] used during training.

Original: For convenience, we refer to the capacity factor at inference as eval-capacityfactor.
Revised: For convenience [Table] we refer to the capacity factor at inference as eval-capacityfactor.

Original: The capacity during training and inference can be computed as follows:
Revised: The capacity during training and inference can be computed as follows:

Original: [Equation]
Revised: [Equation]

Revised: number

Original: [Equation]
Revised: [Equation]

Original:  We evaluated the trained Gshard model with different eval-capacity-factors and show the performance in Figure 9b.
Revised: [Equation] We evaluated the trained Gshard model with different eval-capacity-factors and show the performance in Figure 9b.

Original: The model for evaluation was trained with the capacity factor of 1.0.
Revised: The model for evaluation was trained with the capacity factor of 1.0.

Original: The performance of Gshard is even worse than the standard Transformer, when the eval-capacity-factor is smaller than 0.75.
Revised: The performance of Gshard is even worse than the standard Transformer, when the eval-capacity-factor is smaller than 0.75.

Original: Setting eval-capacity-factor=0.75 is equivalent to setting capacity-facor=3.
Revised: Setting eval-capacity-factor=0.75 is equivalent to setting capacity-facor=3.

Original: if the number of experts is 8, which is very expensive.
Revised: if the number of experts is 8, which is very expensive.

Original: The costs for communication, computation and memory are all proportional to the capacity factor, hence running Gshard on the OPUSdataset is expensive.
Revised: The costs for communication, computation and memory are all proportional to the capacity factor, hence running Gshard on the OPUSdataset is expensive.

Original: We also show the impact of eval-capacity-factor on the translation performance on the WMT17-En-Fr dataset in Figure 9b.
Revised: We also show the impact of eval-capacity-factor on the translation performance on the WMT17-En-Fr dataset in Figure 9b.

Original: It is interesting that the performance of Gshard on WMT17-En-Fr decreases as the eval-capacity-factor increases, which is opposite to the observation on OPUS-100.
Revised: It is interesting that the performance of Gshard on WMT17-En-Fr decreases as the eval-capacity-factor increases, which is opposite to the observation on OPUS-100.

Original: Increasing the eval-capacity-factor changes the routing of MoE, which may hurt the performance.
Revised: Increasing the eval-capacity-factor changes the routing of MoE, which may hurt the performance.

Original: But why do the the multilingual MoE models require a large capacity?
Revised: But why do the the multilingual MoE models require a large capacity?

Original: We conjecture that the routing of multilingual MoE is language-specific, i.e., the tokens in related languages are routed to the same experts.
Revised: We conjecture that the routing of multilingual MoE is language-specific, i.e., the tokens in related languages are routed to the same experts.

Original: In our experiments, each batch contains multiple languages during training, but only one language at inference.
Revised: In our experiments, each batch contains multiple languages during training, but only one language at inference.

Original: Therefore, sentences in the same language are routed to the same experts at inference, which causes these experts to overflow and tokens to be discarded.
Revised: Therefore, sentences in the same language are routed to the same experts at inference, which causes these experts to overflow and tokens to be discarded.

Original: To validate our hypothesis and alleviate this problem, we increase the load-balance weight in Gshard from 0.01 to 0.1, and observe significant improvements in Figure 9b.
Revised: To validate our hypothesis and alleviate this problem, we increase the load-balance weight in Gshard from 0.01 to 0.1, and observe significant improvements in Figure 9b.

Original: But we still need a high capacity factor for the inference of multilingual MoE models.
Revised: But we still need a high capacity factor for the inference of multilingual MoE models.

Original: We leave the issue to our future work.
Revised: We leave the issue to our future work.

Revised: H S ETTING α , β AND γ TO BE N ON -Z ERO S IMULTANEOUSLY

Revised: Although there are three communication channels, we have only used two of them in SCoMoE- α

Revised: ( α = 0, β > 0, γ > 0) and SCoMoE- β ( α = 0, β > 0, γ > 0) in our previous experiments, since it is difficult to tune three hyperparameters at the same time.

Revised: In this section, we investigated whether it is beneficial to set α , β , γ to be non-zero at the same time.

Revised: In Figure 3a, we have observed that α should be set to a small value for a comparable performance with Gshard.

Revised: Therefore we set α to be 0.2 and 0.25 for SCoMoE-Seq and SCoMoE-Feat respectively.

Revised: The results are shown in Table 5.

Revised: Comparing SCoMoE-Seq* with SCoMoE-Seq-

Revised: α and SCoMoE-Seq- β , we can find that keeping α and β non-zero at the same time could speed up the model training without the performance drop.

Revised: Furthermore, SCoMoE-Seq* with ( α = 0 .

Revised: 25 , β = 0 .

Revised: 5 , γ = 0 .

Revised: 25 ) is even better than that with ( α = 0 .

Revised: 25 , β = 0 .

Revised: 25 , γ = 0 . 5 )

Revised: in terms of both BLEU and training speed.

Revised: Therefore, it is better for us to set a small α and a large β for SCoMoE-Seq*.

Revised: SCoMoE-Feat also benefits from setting α , β and γ to be non-zero at the same time, especially for the improvement on BLEU.

Revised: It is reasonable because increasing communication channels for

Revised: [Figure]

Revised: SCoMoE means that each token can be computed in more experts.

Revised: This is similar to increasing heads in multi-head attention.

Revised: I C OMPARISON OF TRANSLATION SPEED

Revised: In previous experiments, all reported results on speed are those of training speed.

Revised: In this section we compared the speed of autoregressive decoding of SCoMoE and analyzed the influence of hyperparameters on translation.

Revised: For SCoMoE-Seq, the hyperparameters of each communication channel α, β, γ could be tuned at inference and be different from the values at training.

Revised: Therefore, we tuned α and β for

Revised: SCoMoE-Seq-

Revised: α and SCoMoE-Seq- β respectively and report BLEU and translation speed in Figure 10a and 10b respectively.

Revised: Overall, the speedup of SCoMoE during autoregressive decoding is weaker than that at training.

Revised: It can be seen that SCoMoE-Seq- α is faster than Gshard consistently at different settings of α .

Revised: SCoMoE-

Revised: Feat- α also shows a slight speedup over Gshard.

Revised: However, both SCoMoE-Seq- β

Revised: and

Revised: SCoMoE-

Revised: Feat- β are slightly slower than Gshard at inference.

Revised: The reason is that the communication during autoregressive inference is bottlenecked by the frequency rather than the size of communication data at inference.

Revised: But one exception is that SCoMoE-Seq- β with β = 0 is faster than Gshard due to the token clustering.

Revised: With regard to the translation performance, both SCoMoE-Seq- α and - β obtain their best performance with α/β = 0 .

Revised: It is also clear that the performance of them decreases as α/β increases.

Revised: All in all, it is better to set α/β to be 0 for SCoMoE-Seq at inference in terms of both translation performance and speedup.

Revised: We use this setting to report BLEU scores for SCoMoE-Seq in other sections.

Revised: J P ERFORMANCE -T IME C URVES

Revised: We have shown the improvements of SCoMoE over Gshard in both training speed and translation performance, but it is still unclear how much time SCoMoE could save compared with Gshard.

Revised: It is also interesting to see whether SCoMoE could benefit from more training steps.

Revised: Therefore we compare the performance-time curves of SCoMoE-Seq- α , SCoMoE-Seq- β and SCoMoE-Feat with

Revised: Gshard in Figure 11a, Figure 11b and Figure 11c respectively.

Revised: We can see that almost all SCoMoE variants are superior to Gshard after being trained for 45,000 seconds.

Revised: And they achieve 50 BLEU at around 80,000 seconds, while Gshard requires 129,000 seconds.

Revised: Among all variants, SCoMoEseq with ( α =0., β =0.5, γ =0.5) and ( α =0.3, β =0., γ =0.7) seem to be the best.

Revised: We then trained these two variants with more steps (12 epochs and 15 epochs respectively) so that they were trained with time comparable to that of Gshard (10 epochs).

Revised: The results of extra steps are illustrated with dashed

Revised: [Figure]

Revised: lines in Figure 11a and Figure 11b.

Revised: It can be seen that these two variants obtain improvements with more training steps and their final performance is much better than Gshard.

Revised: Most variants of

Revised: SCoMoE-Feat converge at around 80,000 steps, and only see little improvements with more training steps.s