Exploring Routing Strategies for Multilingual Mixture-of-Experts Models

1 INTRODUCTION . Scaling up neural network models has recently received great attention , given the significant quality improvements in a variety of areas such as natural language understanding ( Raffel et al. , 2019 ; Brown et al. , 2020 ) and multilingual machine translation ( Huang et al. , 2019 ; Lepikhin et al. , 2020 ) . While training massive models on large amounts of data can almost guarantee improved quality , there are two factors affecting their practicality and applicability : ( 1 ) training efficiency and ( 2 ) inference efficiency . Large dense models are often prohibitively compute-intensive to train , with some models requiring TFlops-days of compute ( Brown et al. , 2020 ) . A recent line of work has proposed sparsely-gated Mixture-of-Experts ( MoE ) layers as an efficient alternative to dense models ( Shazeer et al. , 2017 ; Lepikhin et al. , 2020 ; Riabinin & Gusev , 2020 ) in order to address training efficiency limitations . In a vanilla sparsely-gated MoE model each token of the input sequence activates a different subset of the experts , hence the computation cost per token becomes only proportional to the size of the activated sub-network . However , they fail to meet requirements on inference efficiency . Consider a long sequence where each token of the sequence activates a disjoint subset of available experts . From a practical standpoint , the inference trace of the full sequence spans several experts independently for every token , resulting in an independent pathway for each token . Although this is a desired property adding flexibility to the model and increasing its capacity , it becomes prohibitive for inference for the following reasons : The model parameters in these large models are beyond the memory limit of a single accelerator , and require model parallelism to shard them across a cluster of devices during inference . For models with MoE Layers , the input token would be dynamically routed to different experts allocated to different devices . This further adds communication cost across devices to the overall serving cost . Moreover , due to the sequential nature of the autoregressive decoding ( Kasai et al. , 2020 ; Chen et al. , 2018 ) , the added communication cost from model parallel decoders gets multiplied by the number of decoding steps . To add to this , serving MoE models efficiently requires batching a large number of input tokens together , otherwise only a subset of the MoE network will be activated leading to device under-utilization . In this work , we study the inference efficiency of sparsely gated MoE models while taking into account the characteristics of the intended application , Multilingual Neural Machine Translation ( MNMT ) . MNMT is an inherently multi-task learning problem , aimed at building a single neural network for translating multiple language pairs simultaneously . In a MNMT model , the extent to which parameters are shared across languages determines the magnitude of positive transfer ( Baldwin & Ford , 1988 ) and conversely task interference due to the capacity bottleneck ( Arivazhagan et al. , 2019 ) . In an ideal scenario , we would want to efficiently train a single large MNMT model maximizing transfer while expanding the capacity bottleneck ; at the same time , we would like to enjoy the benefits of sparsely activated sub-networks per-task at inference time , i.e . extracting out a sub-network from the model to decode for a particular language pair to actualize inference efficiency . We propose routing algorithms for MoE models with affordable serving costs . While vanilla MoEs route each sub-word token in the input to its preferred experts , we explore alternative routing strategies that leverage global task level information to route all tokens corresponding to a particular task collectively to the same set of experts . While this strategy could be perceived to be restrictive for parameter sharing across tasks , we empirically demonstrate that routing based on task boundaries performs better when applied to MNMT . During training , we mix the inputs from different tasks in the same batch in order to learn the routing network and encourage positive transfer among the tasks . During inference , we decode different tasks separately and only load the subset of experts associated with the corresponding task . We compare our method with multilingual baselines and find that we achieve significant gains on two benchmarks : a multilingual WMT task with comparable inference cost ( +3.59 BLEU ) , described in Section 4 , and a large internal dataset ( +3.6 BLEU ) , in Section 4.3.2 ) . We see that the gains are comparable with conventional position-wise Mixture-of-Expert models while utilizing decoders with only a fraction ( 6.25 % and 1.56 % ) of their serving cost . We discuss the trade-offs of these different methods in Section 3.2 . In Section 4.3.4 , we analyze the routing decisions made in MoE models and motivate our method . 2 SCALING TRANSFORMERS WITH MIXTURE-OF-EXPERTS . The Transformer ( Vaswani et al. , 2017 ) architecture is a popular model used for neural machine translation and other natural language understanding problems . In sequence-to-sequence problems ( of which neural machine translation is one example ) , the model consists of a separate encoder and decoder , each of which contains multiple Transformer layers . For further details on Transformers , we refer the reader to the original paper ( Vaswani et al. , 2017 ) . We use the Mixture-of-Experts Transformer models used by Lepikhin et al . ( 2020 ) , where the MoE layers for the Transformers consist of E feed-forward networks ( FFN ) , such that ( FFN1 . . . FFNE ) . FFNe ( xs ) = woe · ReLU ( wie · xs ) ys = E∑ e=1 Gs , e · FFNe ( xs ) Here , xs is the input token at position s to the MoE layer and each FFNe is a two layer neural network using a ReLU activation function . wie and woe are the input and output projection weights of the e-th expert . Finally , Gs , E is vector computed by the gating network . For each expert , most values of this vector are zeros , one value being positive . We use this vector to route the token to a select few experts . The entries chosen from Gs , E determine how much the expert contributes to the final output ys . Note that , in this work we choose the top 2 weight experts for each example to be comparable with the prior work . The gating network Gs , E must be considered carefully for efficiency purposes : ( 1 ) the utilization of experts must be balanced and ( 2 ) the function must be efficient to implement at scale . For a more thorough discussion of MoE transformer , we direct the reader to Lepikhin et al . ( 2020 ) . 3 METHODS . In this section we describe our candidate routing strategies in the context of MNMT and discuss their trade-offs from the perspective of the training and inference efficiency . It is known that multi- lingual models learn different overlapping representations depending on the language - this is true for both dense ( Wu & Dredze , 2019 ; Tiedemann , 2018 ; Tan et al. , 2019 ; Östling & Tiedemann , 2016 ; Kudugunta et al. , 2019 ) and sparse models ( Section 4.3.4 ) . Therefore we propose changing the routing algorithm GATE ( xs ) of MoEs to choose different experts using more natural separations . 3.1 ROUTING STRATEGIES . Given the sequential nature of the multilingual machine translation task , the routing decisions can be made at three different granularities , from bottom up ( i ) token-level , ( ii ) sentence-level and ( iii ) task-level , as detailed below . • Token-level Routing : This is the baseline discussed in Section 2 where each token is routed independently . • Sentence-level Routing : Each sequence ( sentence ) , and all tokens that form the sequence , are routed to the same expert . We change the routing algorithm to select experts by sentence representation , calculated by taking the average token representations in a given sentence . • Task-level Routing : We select experts by task boundaries as opposed to making input- level decisions . In the context of MNMT , these task boundaries can either be defined by the target language ( French-to-English and German-to-English are the same task ) or the language pair ( French-to-English and German-to-English are different tasks ) . Gs , E = GATE ( 1 S S∑ s=1 xs ) ( Sentence-level routing ) ( 1 ) Gs , E = GATE ( task ids ) ( Task-level routing ) ( 2 ) We further illustrate the difference in Figure 1 , in token-based MoE models ( Figure 1a ) , tokens from each example are routed to different experts , whereas in task-level MoE models ( Figure 1b ) , tokens may be routed to the same expert based on task . 3.2 INFERENCE IMPLICATIONS OF ROUTING STRATEGIES . While the MoE models discussed in ( Shazeer et al. , 2017 ; Lepikhin et al. , 2020 ) train quickly relative to the number of parameters in terms of the wall-clock time , they are expensive to serve . Consider a MoE with 512 experts and 50B parameters ( Lepikhin et al. , 2020 ) . When employing token-level routing , each token can be independently routed to a different set of experts during inference . Given that the entire model is too large to load into memory on a single accelerator , the two potential solutions to utilize this model for inference are : ( i ) Loading experts dynamically from host to device depending on routing decisions , or ( ii ) Utilizing model-parallelism over multiple accelerators for serving . While the first solution incurs heavy host-device communication costs , the second introduces significantly inter-device communication overhead . Another practical approach to serve a large MoE would require model compression via quantization , pruning or distillation ( Cheng et al. , 2017 ) . While the first two strategies haven ’ t been explored in the context of conditional computation , distillation ( Hinton et al. , 2015 ; Kim & Rush , 2016 ) has been found to introduce undesirable artifacts into the student model ( Freitag et al. , 2019 ; Bogoychev & Sennrich , 2019 ) in the context of NMT . On the other hand , if we limit the number of experts available to every task in the model to a small fraction of the total available capacity , it is possible to extract task-specific models for serving , alleviating the need for complex serving strategies or compression . Since decoding time complexity for auto-regressive seq2seq models is dominated by the decoder ( Kasai et al. , 2020 ) , we can also pursue a hybrid strategy where the encoder utilizes more expensive routing strategies while the decoder of the model utilizes simpler and efficient routing . We do note , however , that MoE models that route purely by task boundaries are slower to train due to load balancing considerations . All examples in the input batch belonging to the same task would route to the same set of experts , possibly leading to some experts bearing a significant amount of the load . Balancing between these inference and training time trade-offs merits further exploration . Summarizing the effective decoding cost of the MoE models utilizing different routing strategies : • Token/Sentence level routing : The routing decisions are made dynamically . Assuming each token/sentence makes disjoint choices , the server needs to load all E experts . • Task-level routing : Tokens corresponding to each input sentence are routed to the same experts statically . The server only needs to pre-load K experts ( assuming top-K routing ) .

Summary: In this paper, the authors explore alternatives to the standard token-based routing in sparsely-gated MoE models for multilingual NMT. This exploration is motivated by the need for efficient inference in MoE models, for which token-based routing is a limitation. The alternative is task-based routing, where examples for a task are assigned to the same experts. This allows efficient device placement and request dispatch at inference time. The paper compares with the approaches as well as hybrid approaches where different parts of the network use different routing mechanisms. The results show that task level routing is comparable to token-level routing with the added benefit of inference efficiency. Performing task-based routing on the decoder side only gives better better translation quality, at the cost of inference efficiently. An analysis of routing decision in token-based routing justifies the design choices.