MoEBERT: from BERT to Mixture-of-Experts via Importance-Guided Adaptation

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1610 - 1623 July 10-15, 2022 ©2022 Association for Computational Linguistics MoEBERT: from BERT to Mixture-of-Experts via Importance-Guided Adaptation Simiao Zuo‡∗, Qingru Zhang‡, Chen Liang‡, Pengcheng He⋄, Tuo Zhao‡ and Weizhu Chen⋄ ‡Georgia Institute of Technology ⋄Microsoft {simiaozuo,qzhang441,cliang73,tourzhao}@gatech.edu {Pengcheng.H,wzchen}@microsoft.com Abstract Pre-trained language models have demon- strated superior performance in various natu- ral language processing tasks. However, these models usually contain hundreds of millions of parameters, which limits their practicality because of latency requirements in real-world applications. Existing methods train small compressed models via knowledge distillation. However, performance of these small mod- els drops significantly compared with the pre- trained models due to their reduced model ca- pacity. We propose MoEBERT, which uses a Mixture-of-Experts structure to increase model capacity and inference speed. We initialize MoEBERT by adapting the feed-forward neu- ral networks in a pre-trained model into multi- ple experts. As such, representation power of the pre-trained model is largely retained. Dur- ing inference, only one of the experts is acti- vated, such that speed can be improved. We also propose a layer-wise distillation method to train MoEBERT. We validate the efficiency and effectiveness of MoEBERT on natural lan- guage understanding and question answering tasks. Results show that the proposed method outperforms existing task-specific distillation algorithms. For example, our method outper- forms previous approaches by over 2% on the MNLI (mismatched) dataset. Our code is pub- licly available at https://github.com/ SimiaoZuo/MoEBERT. 1 Introduction Pre-trained language models have demonstrated superior performance in various natural language processing tasks, such as natural language under- standing (Devlin et al., 2019; Liu et al., 2019; He et al., 2021b) and natural language generation (Rad- ford et al., 2019; Brown et al., 2020). These models can contain billions of parameters, e.g., T5 (Raffel et al., 2019) contains up to 11 billion parameters, and GPT-3 (Brown et al., 2020) consists of up to ∗Corresponding author. 175 billion parameters. Their extreme sizes bring challenges in serving the models to real-world ap- plications due to latency requirements. Model compression through knowledge distil- lation (Romero et al., 2015; Hinton et al., 2015) is a promising approach that reduces the compu- tational overhead of pre-trained language models while maintaining their superior performance. In knowledge distillation, a large pre-trained language model serves as a teacher, and a smaller student model is trained to mimic the teacher’s behavior. Distillation approaches can be categorized into two groups: task-agnostic (Sanh et al., 2019; Jiao et al., 2020; Wang et al., 2020, 2021; Sun et al., 2020a) and task-specific (Turc et al., 2019; Sun et al., 2019; Li et al., 2020; Hou et al., 2020; Sun et al., 2020b; Xu et al., 2020). Task-agnostic dis- tillation pre-trains the student and then fine-tunes it on downstream tasks; while task-specific dis- tillation directly fine-tunes the student after ini- tializing it from a pre-trained model. Note that task-agnostic approaches are often combined with task-specific distillation during fine-tuning for bet- …

Content truncated to 500 words