M o EBERT : from BERT to Mixture-of-Experts via Importance-Guided Adaptation (2022)

Abstract: Pre-trained language models have demonstrated superior performance in various natural 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 models drops significantly compared with the pre-trained models due to their reduced model capacity. 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 neural networks in a pre-trained model into multiple experts. As such, representation power of the pre-trained model is largely retained. During inference, only one of the experts is activated, such that speed can be improved. We also propose a layer-wise distillation method to train MoEBERT. We validate the efficiency and efficacy of MoEBERT on natural language understanding and question answering tasks. Results show that the proposed method outperforms existing task-specific distillation algorithms. For example, our method outperforms previous approaches by over 2% on the MNLI (mismatched) dataset. Our code is publicly available at https://github.com/SimiaoZuo/MoEBERT .

## introduction pre-trained language models have demonstrated superior performance in various natural language processing tasks, such as natural language understanding @xcite @xcite and natural language generation @xcite . these models can contain billions of parameters, e.g., t5 @xcite ) contains up to 11 billion parameters, and @xcite consists of up to 175 billion parameters. their extreme sizes bring challenges in serving the models to real-world applications due to latency requirements. model compression through knowledge distillation @xcite is a promising approach that reduces the computational 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 @xcite @xcite @xcite and task-specific @xcite @xcite @xcite . task-agnostic distillation pre-trains the student and then fine-tunes it on downstream tasks; while task-specific distillation directly fine-tunes the student after initializing it from a pre-trained model. note that task-agnostic approaches are often combined with task-specific distillation during fine-tuning for better performance @xcite . we focus on task-specific distillation in this work. one major drawback of existing knowledge distillation approaches is the drop in model performance caused by the reduced representation power. that is, because the student model has fewer parameters than the teacher, its model capacity is smaller. for example, the student model in distil-bert @xcite has 66 million parameters, about half the size of the teacher (bert-base, @xcite . consequently, performance of distilbert drops significantly compared with bert-base, e.g., over 2% on mnli (82.2 v.s. 84.5) and over 3% on @xcite . we resort to the mixture-of-experts (moe, @xcite structure to remedy the representation power issue. moe models can increase model capacity while keeping the inference com-putational cost constant. a layer of a moe model @xcite @xcite @xcite consists of an attention mechanism and multiple feed-forward neural networks (ffns) in parallel. each of the ffns is called an expert. during training and inference, an input adaptively activates a fixed number of experts (usually one or two). in this way, the computational cost of a moe model remains constant during inference, regardless of the total number of experts. such a property facilitates compression without reducing model capacity. however, moe models are difficult to train-fromscratch and usually require a significant amount of parameters, e.g., 7.4 billion parameters for switchbase @xcite . we propose moebert, which incorporates the moe structure into pretrained language models for fine-tuning. our model can speedup inference while retaining the representation power of the pre-trained language model. specifically, we incorporate the expert structure by adapting the ffns in a pre-trained model into multiple experts. for example, the hidden dimension of the ffn is 3072 in bert-base @xcite , and we adapt it into 4 experts, each has a hidden dimension 768. in this way, the amount of effective parameters (i.e., parameters involved in computing the representation of an input) is cut by half, and we obtain a ×2 speedup. we remark that moebert utilizes more parameters of the pretrained model than existing approaches, such that it has greater representation power. to adapt the ffns into experts, we propose an importance-based method. empirically, there are some neurons in the ffns that contribute more to the model performance than the other ones. that is, removing the important neurons causes significant performance drop. such a property can be quantified by the importance score @xcite @xcite . when initializing moebert, we share the most important neurons (i.e., the ones with the highest scores) among the experts, and the other neurons are distributed evenly. this strategy has two advantages: first, the shared neurons preserve performance of the pre-trained model; second, the nonshared neurons promote diversity among experts, which further boost model performance. after initialization, moebert is trained using a layer-wise task-specific distillation algorithm. we demonstrate efficiency and effectiveness of moebert on natural language understanding and question answering tasks. on the glue @xcite benchmark, our method significantly outperforms existing distillation algorithms. for example, moebert exceeds performance of stateof-the-art task-specific distillation approaches by over 2% on the mnli (mismatched) dataset. for question answering, moebert increases f1 by 2.6 on squad v1.1 @xcite and 7.0 on squad v2.0 @xcite compared with existing algorithms. the rest of the paper is organized as follows: we introduce background and related works in section 2; we describe moebert in section 3; experimental results are provided in section 4; and section 5 concludes the paper. ## method in this section, we first present an algorithm that adapts a pre-trained language model into a moe model. such a structure enables inference speedup by reducing the number of parameters involved in computing an input token's representation. then, we introduce a layer-wise task-specific distillation method that compensates for the performance drop caused by model compression. ## experiments in this section, we evaluate the effectiveness and efficiency of the proposed algorithm on natural language understanding and question answering tasks. we implement our algorithm using the huggingface transformers foot_3 @xcite code-base. all the experiments are conducted on nvidia v100 gpus. ## conclusion we present moebert, which uses a mixture-of-experts structure to distill pre-trained language models. our proposed method can speedup inference by adapting the feed-forward neural networks (ffns) in a pre-trained language model into multiple experts. moreover, the proposed method largely retains model capacity of the pre-trained model. this is in contrast to existing approaches, where the representation power of the compressed model is limited, resulting in unsatisfactory performance. to adapt the ffns into experts, we adopt an importance-based method, which identifies and shares the most important neurons in a ffn among the experts. we further propose a layer-wise taskspecific distillation algorithm to train moebert . we conduct systematic experiments on natural language understanding and question answering tasks. results show that the proposed method outperforms existing distillation approaches.