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

By Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh-Thang Luong, Orhan Firat. 2021

136 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Quoc V. Le
2. Barret Zoph
3. Irwan Bello
4. Mohit Bansal
5. Luke S. Zettlemoyer
6. William Fedus
7. Joseph Dean
8. Jianfeng Gao

In Computer Vision and Pattern Recognition

1. Quoc V. Le
2. Luc Van Gool
3. Dacheng Tao
4. Alan L. Yuille
5. Yi Yang
6. Trevor Darrell
7. Xiaogang Wang
8. Ming-Hsuan Yang

In Model Architectures

1. Irwan Bello
2. Quoc V. Le
3. Barret Zoph
4. Joseph Dean
5. William Fedus
6. Noam Shazeer
7. Yanping Huang
8. Nan Du

In Efficiency & Compression

1. Irwan Bello
2. Joseph Dean
3. Barret Zoph
4. William Fedus
5. Noam Shazeer
6. Yanping Huang
7. Nan Du
8. Sameer Kumar

Citations in the corpus

• Efficiently Scaling Transformer Inference (2022, neutral): “Sparsity techniques, such as task-based mixture of expert architectures (Fedus et al., 2022; Kudugunta et al., 2021; Lepikhin et al., 2020; Shazeer et al., 2017), and adaptive computation techniques that allocate different amounts of compute per input and generation timestep (Jaszczur et al.”
• A Review of Sparse Expert Models in Deep Learning (2022, background): “Kudugunta et al. (2021) instead had an MoE architecture route at the task or sentence level, which allows for more efficient inference and serving.”
• Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models (2022, unknown): “Similarly, Pfeiffer et al. (2022) develop a multilingual expert model with language-specific routing, and Kudugunta et al. (2021) develop a multi-task expert model with task-specific routing.”
• MoE-LLaVA: Mixture of Experts for Large Vision-Language Models (2024, methodology): “This mode is also widely applied in the task-specific MoE (Li et al., 2023e; Zhu et al., 2022; Ma et al., 2023; Kudugunta et al., 2021).”
• Designing Effective Sparse Expert Models (2022, background): “…was shown in Fedus et al. (2021) to be moderately effective, but recent approaches modified the routing to instead route full sentences or tasks (Kudugunta et al., 2021b; Zuo et al., 2021) which then permits extraction of sub-networks at time of serving (e.g. deploy only the network associated…”
• ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, background): “Alternatively, future work could study simply removing the experts in the decoder layer, which also confers benefits during autoregressive decoding (Kudugunta et al., 2021a).”
• A Survey on Mixture of Experts in Large Language Models (2024, unknown): “Kudugunta et al. [75] diverge from the prevalent token-level gating strategies by introducing a sentence-level gating mechanism.”
• Mixture-of-Experts in Remote Sensing: A Survey (2026, unknown): “Sparsely activated multi-task MoEs [41] and modular MoE [14] designs such as Mod-Squad [68] extend this idea by allowing tasks to share only subsets of experts and by routing different tasks or labels to different expert combinations, which can improve resistance to negative transfer and…”