Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference (2021)

Abstract: Sparse Mixture-of-Experts (MoE) has been a successful approach for scaling multilingual translation models to billions of parameters without a proportional increase in training computation. However, MoE models are prohibitively large and practitioners often resort to methods such as distillation for serving. In this work, we investigate routing strategies at different granularity (token, sentence, task) in MoE models to bypass distillation. Experiments on WMT and a web-scale dataset suggest that task-level routing (task-MoE) enables us to extract smaller, ready-to-deploy sub-networks from large sparse models. On WMT, our task-MoE with 32 experts (533M parameters) outperforms the best performing token-level MoE model (token-MoE) by +1.0 BLEU on average across 30 language pairs. The peak inference throughput is also improved by a factor of 1.9x when we route by tasks instead of tokens. While distilling a token-MoE to a smaller dense model preserves only 32% of the BLEU gains, our sub-network task-MoE, by design, preserves all the gains with the same inference cost as the distilled student model. Finally, when scaling up to 200 language pairs, our 128-expert task-MoE (13B parameters) performs competitively with a token-level counterpart, while improving the peak inference throughput by a factor of 2.6x.

## introduction scaling up neural network models has recently received great attention, given the significant quality improvements on a variety of tasks including natural language understanding @xcite and multilingual machine translation @xcite . 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 @xcite . a recent line of work has proposed sparsely-gated mixture-of-experts (moe) layers as an efficient alternative to dense models @xcite @xcite 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 device, 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 @xcite , 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 severe device under-utilization. in this work, we study the inference efficiency of sparsely gated moe models while taking into ac-count 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 @xcite and conversely task interference due to the capacity bottleneck @xcite . in an ideal scenario, we would want to efficiently train a single large mnmt model maximizing transfer while expanding the capacity bottleneck; meanwhile, we would like to enjoy the benefits of sparsely activated sub-networks per-task at inference time, i.e. extracting out a sub-network to decode for a particular language pair to actualize inference efficiency. an alternative way to enjoy high inference efficiency from a large model is knowledge distillation @xcite . however, @xcite found that only a small fraction of quality gains from a large sparse model can be preserved in the student models. instead; ## scaling transformers with mixture-of-experts the transformer @xcite architecture is a popular model used for neural machine translation and other natural language understanding/generation problems. in sequence-to-sequence problems, the model consists of an encoder and decoder, each of which contains multiple transformer layers. for further details, we refer the reader to the original paper @xcite . we use the mixture-of-experts transformer models proposed by @xcite , where the moe layers for the transformers consist of e feed-forward networks (ffn), such that here, x s is the input token at position s to the moe layer and each ffn e is a two layer neural network using a relu activation function. wi e and wo e are the input and output projection weights of the e-th expert. finally, g s,e is vector computed by the gating network (also referred as router). 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 g s,e determine how much the expert contributes to the final output y s . 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 g s,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 transformers, we direct the reader to @xcite . ## 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. multilingual models learn joint representations across languages to the extent of the parameters being shared (wu and dredze, 2019; tiedemann, 2018; tan et al., 2019; zhang et al., 2020; östling and tiedemann, 2016; kudugunta et al., 2019). while being beneficial for transfer, extreme sharing of the parameters exacerbates interference. allowing dedicated (unshared) parameters are known to be effective at mitigating interference (zhang et al., 2021; kong et al., 2021) and moe variants are inherently learn such partitioning across languages/tasks. therefore we study the routing algorithm gate(x s ) of moes to mitigate interference, while enabling transfer and effective at inference. ## experiments on 30 language pairs we compare routing strategies at multiple levels in both, the encoder and the decoder, by conducting extensive experiments on two benchmarks: the public wmt dataset with 30 language pairs (section 4.1) and an in-house web-scale dataset with 200 language pairs (section 5). we start with wmt setup. ## scaling up to 200 language pairs we now scale our results up to a larger internal dataset with over 200 language pairs, while also scaling the number of parameters to beyond 10 billion weights. in addition, we look more closely at the gating decisions made by these sparse models and discuss their implications. ## conclusions in this work we discussed more inference friendly algorithms for routing examples in multilingual sparse mixture-of-experts models by making use of task boundaries. we empirically demonstrated that this new algorithm performs as well as, or better than, conventional token-based routing algorithms on two different datasets: a multilingual wmt setup covering 30 language pairs and a large internal dataset covering 200 language pairs, in terms of machine translation quality evaluated with bleu. by carefully comparing inference throughput across different routing approaches and distilled models, we demonstrated the superiority of task-based routing algorithms over either serving a token-based moe model as-is (in terms of peak throughput) and over distilling a large moe model into a smaller dense model (in terms of bleu). we conclude by highlighting that algorithms that are more inference friendly while retaining the quality gains of moe models are a promising direction for future exploration, motivating research on inference efficiency for large models. although we studied some hybrid routing strategies where encoder and decoder networks utilize different routing schemes, we believe that future research on more granular routing hybrids or hierarchical variants will deliver more gains and advance our understanding of large scale, sparsely gated, massively multi-task networks. shijie wu and mark dredze. 2019. beto, bentz, becas: the surprising cross-lingual effectiveness of bert. arxiv preprint arxiv:1904.09077. brandon yang, gabriel bender, quoc v le, and jiquan ngiam. 2019. condconv: conditionally parameterized convolutions for efficient inference. in advances in neural information processing systems, pages 1307-1318. biao zhang, ankur bapna, rico sennrich, and orhan firat. 2021. share or not? learning to schedule language-specific capacity for multilingual translation. in international conference on learning representations. biao zhang, philip williams, ivan titov, and rico sennrich. 2020. improving massively multilingual neural machine translation and zero-shot translation. arxiv preprint arxiv:2004.11867.