L ee BERT : Learned Early Exit for BERT with cross-level optimization (2021)

Abstract: Pre-trained language models like BERT are performant in a wide range of natural language tasks. However, they are resource exhaustive and computationally expensive for industrial scenarios. Thus, early exits are adopted at each layer of BERT to perform adaptive computation by predicting easier samples with the first few layers to speed up the inference. In this work, to improve efficiency without performance drop, we propose a novel training scheme called Learned Early Exit for BERT (LeeBERT). First, we ask each exit to learn from each other, rather than learning only from the last layer. Second, the weights of different loss terms are learned, thus balancing off different objectives. We formulate the optimization of LeeBERT as a bi-level optimization problem, and we propose a novel cross-level optimization (CLO) algorithm to improve the optimization results. Experiments on the GLUE benchmark show that our proposed methods improve the performance of the state-of-the-art (SOTA) early exit methods for pre-trained models.

## introduction the last couple of years have witnessed the rise of pre-trained language models (plms), such as bert @xcite , gpt @xcite , xlnet @xcite , and albert @xcite , etc. by pre-training on the unlabeled corpus and fine-tuning on labeled ones, bert-like models achieved considerable improvements in many natural language processing (nlp) tasks, such as text classification and natural language inference (nli), sequence labeling, etc. however, these plms suffer from two problems. the first problem is efficiency. the state-of-the-art (sotas) achievements of these models usually rely on very deep model architectures accompanied by high computational demands, impairs their practicalities. like general search engines or online medical consultation services, industrial settings process generally millions of requests per minute. what makes efficiency more critical is that the traffic of online services varies drastically with time. for example, during the flu season, the search requests of dingxiangyuan foot_0 are ten times more than usual. and the number of claims during the holidays is five to ten times more than that of the workdays for online shopping. many servers need to be deployed to enable bert in industrial settings, which is unbearable for many companies. second, previous literature @xcite @xcite pointed out that large plms with dozens of stacked transformer layers are over-parameterized and could suffer from the "overthinking" problem @xcite . that is, for many input samples, their shallow representations at a shallow layer are enough to make a correct classification. in contrast, the final layer's representations may be overfitting or distracted by irrelevant features that do not generalize. the overthinking problem leads to not only poor generalization but also wasted computation. to address these issues, both the industry and academia have devoted themselves to accelerating plms at inference time. standard methods include direct network pruning @xcite @xcite , knowledge distillation @xcite @xcite , weight quantization @xcite @xcite and adaptive inference @xcite @xcite . among them, adaptive inference has attracted much attention. given that real-world data is usually com-posed of easy samples and difficult samples, adaptive inference aims to deal with simple examples with only a small part of a plm, thus speeding up inference time on average. the speed-up ratio can be controlled with certain hyper-parameters to cope with drastic changes in request traffic. what's more, it can address the over-thinking problem and improve the model's generalization ability. early exiting is one of the most crucial adaptive inference methods @xcite . it implements adaptive inference by installing exits, or intermediate prediction layer, at each layer of bert and exiting "easy" samples at exits of the shallow layers to speed up inference (figure 1 ). strategies for early exiting are designed @xcite @xcite , which decides when to exit given the current obtained predictions (from previous and current layers). early exiting architectures' training procedure is essentially a multi-objective problem since each exit is trying to improve its performance. different objectives from different classifiers may conflict and interfere with one-another (phuong and lampert, 2019; @xcite . thus they incorporate distillation loss to improve the training procedure by encouraging early exits to mimic the output distributions of the last exit. the motivation is that the last exit has the maximum network capacity and should be more accurate than the earlier exits. in their work, only the last exit can act as a teacher exit. besides, the multiple objectives are uniformly weighted. in this work, we propose a novel training mechanism called learned early exiting for bert (lee-bert). our contributions are three folded. first, instead of learning from the last exit, leebert asks each exit to learn from each other. the motivation is that different layers extract features of varying granularity. thus they have different perspectives of the sentence. distilling knowledge from each other improves the expressiveness of lower exits and alleviates the overfittng of the later exits. second, to achieve the optimal trade-offs between different loss terms, their weights are treated as parameters and are learned along with model parameters. the optimization of the learnable weights and model parameters is formulated as a bi-level optimization problem, optimized with gradient descent. built upon previous literature @xcite , we propose a novel cross-level optimization (clo) algorithm to solve the bilevel optimization better. extensive experiments are conducted on the glue benchmark @xcite , and show that leebert outperforms existing sota bert early exiting methods, sometimes by a large margin. ablation study shows that: (1) knowledge distillation among all the exits can improve their performances, especially for the shallow ones; (2) our novel clo algorithm is useful in learning more suitable weights and brings performance gains. our contributions are integrated into our lee-bert framework, which can be summarized as follows: ## preliminaries in this section, we introduce the necessary background for bert early exiting. throughout this work, we consider the case of multi-class classification with samples {(x n , y n ), x @xmath0, y @xmath1, @xmath2, 2, ..., n }, e.g., sentences, and the number of classes is k. ## our leebert framework in this section, we introduce the proposed lee-bert framework. first, we present our distillation based loss design, and then we elaborate on how to optimize with learnable weights. our main contribution is a novel training mechanism for bert early exiting, which extends @xcite and phuong and lampert (2019) via mutual distillation and learned weights. ## conclusion and discussions in this work, we propose a new framework for improving plms' early exiting. our main contributions lie in two aspects. first, we argue that exits should learn and distill knowledge from each other during training. second, we propose that early exiting networks' training objectives be weighted differently, where the weights are learnable. the learnable weights are optimized with the cross-level optimization we propose. experiments on the glue benchmark datasets show that our framework can improve plms' early exiting performances, especially under high latency requirements. our framework is easy to implement and can be adapted to various early exiting strategies. we want to explore novel exiting strategies that better guarantee exiting performances in the future.