E fficient BERT : Progressively Searching Multilayer Perceptron via Warm-up Knowledge Distillation (2021)

Abstract: Pre-trained language models have shown remarkable results on various NLP tasks. Nevertheless, due to their bulky size and slow inference speed, it is hard to deploy them on edge devices. In this paper, we have a critical insight that improving the feed-forward network (FFN) in BERT has a higher gain than improving the multi-head attention (MHA) since the computational cost of FFN is 2~3 times larger than MHA. Hence, to compact BERT, we are devoted to designing efficient FFN as opposed to previous works that pay attention to MHA. Since FFN comprises a multilayer perceptron (MLP) that is essential in BERT optimization, we further design a thorough search space towards an advanced MLP and perform a coarse-to-fine mechanism to search for an efficient BERT architecture. Moreover, to accelerate searching and enhance model transferability, we employ a novel warm-up knowledge distillation strategy at each search stage. Extensive experiments show our searched EfficientBERT is 6.9 \times smaller and 4.4 \times faster than BERT \rm_{BASE} , and has competitive performances on GLUE and SQuAD Benchmarks. Concretely, EfficientBERT attains a 77.7 average score on GLUE test , 0.7 higher than MobileBERT \rm_{TINY} , and achieves an 85.3/74.5 F1 score on SQuAD v1.1/v2.0 dev , 3.2/2.7 higher than TinyBERT _4 even without data augmentation. The code is released at https://github.com/cheneydon/efficient-bert .

## introduction diverse pre-trained language models (plms) (e.g., bert @xcite ) have been intensively investigated by designing new pretext tasks, architectures, or attention mechanisms @xcite @xcite . the performances of these plms far exceed the traditional methods on a variety of natural language * corresponding author. processing (nlp) tasks. nevertheless, their shortcomings are still evident (including a considerable model size and low inference efficiency), limiting real-world application scenarios. to alleviate the aforementioned limitations, many model compression methods have been proposed, including quantization, weight pruning, and knowledge distillation (kd) @xcite @xcite . among them, kd @xcite that transfers the knowledge from larger teacher models to smaller student models with minimal performance sacrifice is most widely used due to its plug-and-play feasibility and its scalability in the rapid delivery of new models. specifically, kd allows us to train our own bert architecture significantly faster than training from scratch. hence, we adopt kd in this paper. besides, inspired by the impressive results by neural architecture search (nas) in vision tasks @xcite @xcite , adopting nas to further boost the performance of plms or reduce the computational cost has attracted increasing attentions @xcite @xcite . although considerable progress has been made in the field of kd for plms, the compression of the feed-forward network (ffn) has been rarely studied. this contradicts the fact that the computational cost of ffn is 2∼3 times larger than that of the multi-head attention (mha). in addition, @xcite have proved that the multilayer perceptron (mlp) in ffn can prevent the undesirable rank collapse caused by self-attention and thus can improve bert optimization. these motivate us to investigate the nonlinearity of ffn in bert. in this paper, we make the first attempt to compress and improve the barely-explored multilayer perceptron (mlp) in ffn and propose a novel coarse-to-fine nas approach with warm-up kd to find the optimal mlp architectures, aiming to search for a universal small bert model with competitive performance and strong transferability. specifically, we design a rich and flexible search space to discover an excellent ffn with maximal nonlinearity and a minimal computational cost. our search space contains various mathematical operations, stack numbers, and expansion ratios of intermediate hidden size. to efficiently search from our vast search space, we progressively shrink the search space in three stages. even with this elegant coarse-to-fine search strategy, pre-training each candidate model still needs a lot of time to converge during searching. to solve this problem, different from the conventional kd strategy @xcite , we propose a warm-up kd strategy to fast transfer the knowledge, where a pre-trained supernet is additionally introduced to perform a joint warm-up for all candidate models. note that the warm-up strategy in the third stage is slightly different from that of the first two stages. during the first two stages, each candidate model initially inherits its weights from a frozen warmedup supernet to accelerate searching. but in the third stage, since there is no need to search mathematical operations, an unfrozen warmed-up supernet sharing weights across different candidate models is allowed, i.e., each model can inherit weights from this activated warmed-up supernet for a quick launch and is then trained with weight sharing in a multi-task manner to enhance transferability. extensive experimental results show that our searched architecture, named efficientbert, is 6.9× smaller and 4.4× faster than bert base , and has competitive performance. on the test set of glue benchmark, efficientbert attains an average score of 77.7, which is 0.7 higher than mobilebert tiny , and achieves an f1 score of 85.3/74.5 on the squad v1.1/v2.0 dev dataset, which is 3.2/2.7 higher than tinybert 4 even without data augmentation. nevertheless, the compression of the feedforward network (ffn) has not been well studied, although its computational cost is 2∼3 larger than the multi-head attention (mha) as pointed out by teacher model student model gelu add & norm … … embedding classifier mha ffn add & norm linear operations linear operations x gelu max tanh add & norm linear operations mathematical operations linear operations l  x mathematical operations mathematical operations add & norm mathematical operations x mathematical operations mathematical operations (a) search stage 1 (b) search stage 2 (c) search stage 3 activated operations unused operations activated channels unused channels l  activated stack number l  frozen stack number activated path frozen path … … embedding classifier mha ffn x 1  l  l  warm-up kd frozen supernet candidate subnet single-task dataset mnli search space (search stage 1) ffn mathematical operation ffn stack number ffn intermediate expansion ratio search space (search stage 2) ffn mathematical operation search space (search stage 3) ffn stack number ffn intermediate expansion ratio activated supernet multi-task dataset mnli rte … candidate subnet warm-up kd (search stage 1, 2) warm-up kd (search stage 3) search space both in the search and retraining stages, we use a novel warm-up knowledge distillation to transfer the teacher model's knowledge (middle). specifically, each candidate or retrained subnet first inherits the weights from a frozen or activated warmed-up supernet, then conducts pre-training and fine-tuning with single-task or multi-task dataset. @xcite . in contrast, compressing ffn is our main focus in this work. neural architecture search. motivated by the success of neural architecture search (nas) in computer vision @xcite @xcite , increasing attention has been paid to applying nas to nlp tasks @xcite @xcite , aiming to automatically search for optimal architectures from a vast search space. evolved transformer @xcite employs nas to search for a better transformer architecture with an evolutionary algorithm. hat @xcite applies nas to search for efficient hardware-aware transformer models based on a transformer supernet. adabert @xcite searches for task-adaptive small models with kd and differentiable nas method. nas-bert @xcite proposes a task-agnostic nas method for adaptive-size model compression, where several acceleration techniques (including block-wise search, search space pruning, and performance approximation) are introduced to speed up the searching process. differently, in this paper, we design a comprehensive search space towards the nonlinearity of multilayer perceptron (mlp) in ffn, and propose a novel coarse-to-fine nas approach with warm-up kd to find the optimal mlp architectures. unlike adabert, we apply nas to search ## our efficientbert we aim at discovering a lightweight mlp architecture with better nonlinearity in each ffn layer, ensuring the searched model can achieve compelling performance. we first present the search space towards better nonlinearity of mlp in ffn, as described in section 3.1. then we propose a novel coarse-to-fine nas method with warm-up kd as discussed in section 3.2. ## experiment this section demonstrates the superior performance and transferability of our efficientbert on a wide range of downstream tasks. ## conclusion in this paper, we focus on the compression and improvement of ffn and design a profound search space over the nonlinearity of mlp in ffn, aiming at searching for better mlp architectures to improve the model performance. due to the enormous search space, we conduct nas in a progressive manner and employ a novel warm-up kd strategy at each search stage to accelerate searching and enhance model transferability. extensive experiments show that our searched architecture efficientbert is 6.9× smaller and 4.4× faster than bert base , and has competitive performance and strong generalization ability. in the future, we will leverage nas to discover more dynamic plms w.r.t different hardwares and downstream tasks.