LECO : Improving Early Exiting via Learned Exits and Comparison-based Exiting Mechanism (2023)

Abstract: Recently, dynamic early exiting has attracted much attention since it can accelerate the inference speed of pre-trained models (PTMs). However, previous work on early exiting has neglected the intermediate exits’ architectural designs. In this work, we propose a novel framework, Learned Exits and COmparison-based early exiting (LECO) to improve PTMs’ early exiting performances. First, to fully uncover the potentials of multi-exit BERT, we design a novel search space for intermediate exits and employ the idea of differentiable neural architecture search (DNAS) to design proper exit architectures for different intermediate layers automatically. Second, we propose a simple-yet-effective comparison-based early exiting mechanism (COBEE), which can help PTMs achieve better performance and speedup tradeoffs. Extensive experiments show that our LECO achieves the SOTA performances for multi-exit BERT training and dynamic early exiting.

## introduction despite achieving state-of-the-art (sota) performances on almost all the natural language processing (nlp) tasks @xcite , large pre-trained language models (plms) still have difficulty being applied to many industrial scenarios with low latency requirements. many research works are devoted to speeding up the inference of bert or other plms, such as network pruning @xcite @xcite , student network distillation @xcite @xcite , and early exiting @xcite @xcite . due to its potential in applications, early exiting has attracted much attention in the research field @xcite . early exiting requires a multi-exit bert, a bert backbone with an intermediate classifier (or exit) installed on each layer. and then, a dynamic early exiting mechanism is applied during the forward pass to ensure efficient inference. early exiting is in parallel with and can work together with static model compression methods @xcite . however, the literature focuses less on the training of multi-exit bert @xcite @xcite and there is no literature systematically discussing the architectural design of the intermediate exits. in this work, we propose a novel framework, learned exits and comparison-based early exiting (leco), designated to discover the full potentials of multi-exit bert in early exiting. first, we design a suitable and comprehensive search space for architectural learning of the intermediate exits (see figure 1 ). our search space contains candidate activation functions, encoding operations, and pooling operations. we follow the differentiable neural architecture search (dnas) framework like @xcite ; @xcite ; @xcite to learn a set of intermediate exits with different architectures automatically. second, reflecting on the limitations of the patience-based early exiting method pabee @xcite , we propose a comparison-based early exiting (cobee) mechanism. cobee makes early exiting decisions by comparing the predicted distributions of adjacent intermediate layers. we conduct extensive experiments and ablation studies on the glue benchmark @xcite . we show that learned intermediate exits of leco outperform the previous sota multi-exiting bert training methods while adding fewer trainable parameters. furthermore, our novel dynamic early exiting mechanism cobee outperforms the previous sota early exiting mechanisms. further analysis shows that: (a) our leco framework can help to boost the performance of multi-exiting bert under different training strategies. (b) our novel dynamic early exiting strategy outperforms the baseline early exiting methods. our contributions are as follows: ## preliminaries in this section, we introduce the necessary background for bert early exiting. we consider the case of multi-class classification with k classes, @xmath0, 2, ..., k}. the dataset consists of n samples {(x i , y i ), @xmath1, 2, ..., n }} , where x i is an input sentence consisting of l words, and y @xmath2. ## comparison-based early exiting the patience-based mechanism @xcite validates the early exiting decisions among the previous layers, providing a promising direction for designing early exiting mechanisms. the early exiting condition in pabee is coarse: it directly compares the predicted labels. however, tt is common for bert to change its predictions after a few intermediate layers. thus, pabee's early exiting performances with low patience parameters may not be reliable. to summarize, we need a more fine-grained criterion to generate more reliable early exiting signals. we now introduce our comparison-based early exiting method, cobee. the inference procedure is illustrated in figure 1 . assume the forward pass has reached layer @xmath3. we now compare the predicted distributions of layer m and layer m ′ (@xmath4. denote the label that receives the highest probability mass at layer m as k * m , and the probability distribution of exit m is denoted as pr m , then the disagreement between layer m and layer m ′ is calculated as: for simplicity, we denote di m,m ′ = di(pr m , pr m ′ ) ∈ r. the smaller the value of di m,m ′ , the predicted distributions pr m and pr m ′ are more consistent with each other. we use a counter cnt to store the number of times the disagreement scores between adjacent layers are less than the pre-defined exiting threshold τ . at layer m, cnt m is calculated as: if di m,m-1 is less than the pre-defined threshold, then the patience counter is increased by 1. otherwise, the patience counter is reset to 0. if cnt m reaches the pre-defined patience value t, the model stops inference and exits early. otherwise, the model goes to the next layer. however, if the model does not exit early at intermediate layers, the model uses the final classifier f m for prediction. ## experiments 6.1 datasets we evaluate our proposed approach to the classification tasks on glue benchmark @xcite . we only exclude the sts-b task since it is a regression task, and we exclude the wnli task following previous work @xcite @xcite . since the original test sets are not publicly available, we follow @xcite and @xcite to construct the train/dev/test splits as follows: (a) for datasets with fewer than 10k samples (rte, mrpc, cola), we divide the original validation set in half, using one half for validation and the other for testing. (b) for larger datasets, we split 1k samples from the training set as the development set, and use the original development set as the test set. the detailed dataset statistics are presented in for mnli, we report acc, which is the average of the accuracy scores on the matched and mismatched test set. for mrpc and qqp, we report acc-f1, which is the average of accuracy and f1 scores. for cola, we report mcc, which is the matthews correlation. for all other tasks, we report accuracy (acc). ## conclusion in this work, we propose a novel framework, leco. our contributions are three-fold. first, leco designs a unified search space for architectural designs of intermediate exits. second, we apply the differentiable nas framework of darts to learn the optimal exit architectures automatically. third, we propose a novel comparison based early exiting mechanism, cobee. experiments on the glue benchmark and ablation studies demonstrate that our leco framework can achieve sota on multi-exit bert training and outperforms the previously sota dynamic early exiting methods.