BER xi T : Early Exiting for BERT with Better Fine-Tuning and Extension to Regression (2021)

Abstract: The slow speed of BERT has motivated much research on accelerating its inference, and the early exiting idea has been proposed to make trade-offs between model quality and efficiency. This paper aims to address two weaknesses of previous work: (1) existing fine-tuning strategies for early exiting models fail to take full advantage of BERT; (2) methods to make exiting decisions are limited to classification tasks. We propose a more advanced fine-tuning strategy and a learning-to-exit module that extends early exiting to tasks other than classification. Experiments demonstrate improved early exiting for BERT, with better trade-offs obtained by the proposed fine-tuning strategy, successful application to regression tasks, and the possibility to combine it with other acceleration methods. Source code can be found at https://github.com/castorini/berxit .

## introduction large-scale pre-trained language models such as bert @xcite have brought the natural language processing (nlp) community large performance gain but at the cost of heavy computational burden. while pre-trained models are available online and fine-tuning is typically done without a strict time budget, inference poses a much lower latency tolerance, and the slow inference speed of these models can impede easy deployment. it becomes even more difficult when inference has to be done on edge devices due to limited network capabilities or privacy concerns. early exiting @xcite @xcite has been proposed to accelerate the inference of bert and models with similar architecture, i.e., those comprising multiple transformer layers @xcite with a classifier at the top. instead of using only one ⋮ ⋮ transformer 1 classifier 1 confidence / lte exit continue transformer 2 classifier 2 confidence / lte exit continue transformer 𝑛 classifier 𝑛 exit input figure 1: multi-output structure of early exiting bert. while existing early exiting papers provide promising quality-efficiency trade-offs, improvements are necessary for two important components: fine-tuning strategies and exiting decision making. in these papers, fine-tuning strategies are relatively simple and fail to take full advantage of the pretrained model's effectiveness; we propose a novel fine-tuning strategy, alternating, for this multioutput model. moreover, previous work makes exiting decisions based on the confidence of output probability distributions, and is hence only applicable to classification tasks; we extend it to other tasks by proposing the learning-to-exit idea. with carefully designed fine-tuning strategies and methods for making exiting decisions, the model can achieve better quality-efficiency trade-offs and can be extended to regression tasks. we refer to our proposed ideas collectively as berxit (bert+exit), and apply it to muppets 1 including bert, roberta @xcite , and albert @xcite ; we also apply it on top of another bert acceleration method, distil-bert @xcite . we conduct experiments 1 bert and his friends. on datasets including classification and regression tasks, and show that our method can save up to 70% of inference time with minimal quality degradation. our contributions include the following: (1) an effective fine-tuning method alternating; (2) the learning-to-exit idea that extends early exiting to tasks other than classification; (3) extensive experiments that show the effectiveness of our ideas and the successful combination of early exiting with other bert acceleration methods; (4) additional experiments that provide insight into the inner mechanism of pre-trained models. in the general deep learning context, there are a number of well-explored methods to accelerate model inference. pruning @xcite @xcite removes unimportant parts of the neural model, from individual weights to layers and blocks. quantization @xcite reduces the number of bits needed to operate a neural model and to store its weights. distillation @xcite transfers knowledge from large teacher models to small student models. these methods typically require pre-training muppets from scratch foot_0 and produce only one small model with a predetermined target size. early exiting requires only fine-tuning and also produces a series of small models, from which the user can choose flexibly. it extends the idea of adaptive computation @xcite for recurrent neural networks, and is also closely related to @xcite , multi-scaled densenet @xcite , and slimmable network @xcite . early exiting for muppets has been explored by rtj foot_1 @xcite , deebert @xcite , and fastbert @xcite . despite their promising results, there is still room for im-provement regarding the fine-tuning strategies of rtj and deebert. fastbert, on the other hand, uses self-distillation @xcite for fine-tuning, which works well for small muppets such as bert base . however, our preliminary experiments foot_2 show that selfdistillation is unstable on larger muppets such as bert large , suggesting that future work is necessary for fully understanding and robustly applying self-distillation on muppets. all these three methods make early exiting decisions based on confidence (or its variants) of the predicted probability distribution, and are therefore limited to classification tasks. runtime neural pruning @xcite , skipnet @xcite , and blockdrop @xcite use reinforcement learning (rl) to decide whether to execute a network module. universal transformer @xcite and depth-adaptive transformer (dat, @xcite use learned decisions for early exiting in sequence-to-sequence tasks. concurrently, pabee @xcite proposes patience-based early exiting which is applicable to regression, but it relies on inter-layer prediction consistency and is therefore not very efficient for exiting at early layers. inspired by them, we propose a method to extend early exiting for muppets to regression tasks, using only layerspecific information as in classification. moreover, our method requires neither rl nor complicated distribution fitting as in dat, but uses a straightforward layer-wise certainty estimation, and achieves performance comparable with confidence-based early exiting on classification tasks. ## model structure and fine-tuning we start from a pre-trained muppet model (the backbone model), attach additional classifiers to it, fine-tune the model, and use it for accelerated inference by early exiting. backbone model the backbone model is an nlayer pre-trained muppet model. we denote the i th layer hidden state corresponding to the [cls] token as h i : classifiers in the original bert paper @xcite , the way to fine-tune is to attach a classifier to the final transformer layer, and then to jointly update both the backbone model and the classifier. the classifier is a one-layer fullyconnected network. it takes as input the final layer hidden state h n and outputs a prediction. its output is a probability distribution over all classes for classification tasks and a scalar for regression 5 tasks. to enable early exiting, we instead attach a classifier to every transformer layer, i.e., there are n classifiers in total. each classifier can make its own prediction, and therefore the model can accelerate inference by exiting earlier. fine-tuning strategies we discuss how to finetune this multi-output network. the loss function for the i th layer classifier is the most straightforward fine-tuning strategy is perhaps minimizing the sum of all classifiers' loss functions and jointly updating all parameters in the process. we refer to this strategy as joint, and it is also used in rtj @xcite : if we hope to preserve the best model quality for the final layer, the desired fine-tuning strategy is two-stage, which is also used in deebert @xcite . the first stage is identical to vanilla bert fine-tuning: updating the backbone model and only the final classifier. in the second stage, we freeze all parameters updated in the first stage, and fine-tune the remaining classifiers. objectives in the two stages are as follows. stage 1: min these two fine-tuning strategies are not ideal. intuitively, in this multi-output network, loss functions of different classifiers interfere with each other in a negative way. transformer layers have to provide hidden states for two competing purposes: immediate inference at the adjacent classifier and gradual feature extraction for future classifiers. therefore, achieving a balance between the classifiers is critical. two-stage produces final classifiers with optimal quality at the price of earlier layers, since most parameters are solely optimized for the final classifier. joint treats all classifiers equally, and therefore its final classifier is less effective than that of two-stage. to combine the advantages, we propose a novel fine-tuning strategy, alternating. it alternates between two objectives (taken from equation 3 and 4) for odd-numbered and even-numbered iterations. ## exiting decision making after the entire model (including the backbone and all classifiers) is fine-tuned, it can perform early exiting for an inference sample. in this section we discuss two methods to make exiting decisions. ## conclusion and future work to improve early exiting for muppets, we present berxit, including the alternating fine-tuning strategy which outperforms methods from previous papers and the lte idea which extends early exiting to a broader range of tasks. experiments show the effectiveness of alternating in providing better quality-efficiency trade-offs and the successful application of lte to regression tasks. they also show that early exiting is cumulative with other acceleration methods such as distilbert and has the potential for model interpretation. future work the fundamental question of early exiting for muppets is how many transformer layers are sufficient for making good predictions. we draw inspiration from the limit performance of muppets: the score of limit at the i th layer is obtained by taking the first i transformer layers from the pre-trained muppet model, attaching a classifier to the i th layer, and fine-tuning this single-output model. limit estimates the upper bound for any fine-tuning methods by removing inter-classifier interference. we compare the limit performance of bert base and bert large in figure 9 , and notice that with the same number of layers (and identical fine-tuning strategy), bert base almost always outperforms bert large by a large margin. this suggests that most transformer layers' potential to provide information for early exiting is limited by the single-output nature of pre-training. if we want to further improve early exiting muppets for better trade-offs, adding more exiting paths in pretraining would be a promising direction.