Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Simiao Zuo, Qingru Zhang, Chen Liang, Pengcheng He, Tuo Zhao, Weizhu Chen

TL;DR

MoEBERT presents a practical approach to compress and accelerate large pre-trained language models by converting the feed-forward networks in Transformer layers into a Mixture-of-Experts structure. It uses importance-guided adaptation to share the most critical neurons across experts and trains the model with layer-wise distillation to preserve performance. Empirical results on GLUE and SQuAD show MoEBERT outperforming both task-agnostic and task-specific distillation baselines, with notable improvements on MNLI and SQuAD v2.0, while achieving about a 2× inference speedup. The work suggests a viable path for deploying high-capacity language models in latency-sensitive applications without sacrificing accuracy.

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 effectiveness 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.

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

TL;DR

MoEBERT presents a practical approach to compress and accelerate large pre-trained language models by converting the feed-forward networks in Transformer layers into a Mixture-of-Experts structure. It uses importance-guided adaptation to share the most critical neurons across experts and trains the model with layer-wise distillation to preserve performance. Empirical results on GLUE and SQuAD show MoEBERT outperforming both task-agnostic and task-specific distillation baselines, with notable improvements on MNLI and SQuAD v2.0, while achieving about a 2× inference speedup. The work suggests a viable path for deploying high-capacity language models in latency-sensitive applications without sacrificing accuracy.

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 effectiveness 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.
Paper Structure (20 sections, 9 equations, 4 figures, 7 tables)

This paper contains 20 sections, 9 equations, 4 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Backbone: Transformer
  4. Mixture-of-Experts Models
  5. Pre-trained Language Models
  6. Knowledge Distillation
  7. Method
  8. Importance-Guided Adaptation of Pre-trained Language Models
  9. Layer-wise Distillation
  10. Model Training
  11. Experiments
  12. Datasets
  13. Baselines
  14. Implementation Details
  15. Main Results
  16. ...and 5 more sections

Figures (4)

  • Figure 1: Adapting a two-layer FFN into two experts. The blue neuron is the most important one, and is shared between the two experts. The red and green neurons are the second and third important ones, and are assigned to expert one and two, respectively.
  • Figure 2: Ablation study on MNLI. We report the average accuracy of MNLI-m and MNLI-mm. As default settings, we have expert dimension $768$, number of experts $4$, and shared dimension $512$.
  • Figure 3: Experimental results of model variants on MNLI (average of m and mm). Our methods are denoted Import, All and Hash-r in the subfigures, respectively.
  • Figure 4: Inference speed (examples/second, CPU) on the SST-2 dataset.