BADGE : Speeding Up BERT Inference after Deployment via Block-wise Bypasses and Divergence-based Early Exiting (2023)

Abstract: Early exiting can reduce the average latency of pre-trained language models (PLMs) via its adaptive inference mechanism and work with other inference speed-up methods like model pruning, thus drawing much attention from the industry. In this work, we propose a novel framework, BADGE, which consists of two off-the-shelf methods for improving PLMs’ early exiting. We first address the issues of training a multi-exit PLM, the backbone model for early exiting. We propose the novel architecture of block-wise bypasses, which can alleviate the conflicts in jointly training multiple intermediate classifiers and thus improve the overall performances of multi-exit PLM while introducing negligible additional flops to the model. Second, we propose a novel divergence-based early exiting (DGE) mechanism, which obtains early exiting signals by comparing the predicted distributions of two adjacent layers’ exits. Extensive experiments on three proprietary datasets and three GLUE benchmark tasks demonstrate that our method can obtain a better speedup-performance trade-off than the existing baseline methods.\footnote{Code will be made publicly available to the research community upon acceptance.}

## introduction since bert @xcite , the pre-trained language models (plms) have become the default state-of-the-art (sota) models for natural language processing (nlp). recent years have witnessed the rise of many plms, such as gpt @xcite , xlnet @xcite , and albert @xcite , and so forth. these bert-style models achieved considerable improvements in many natural language processing (nlp) tasks by pre-training on the unlabeled corpus and fine-tuning on labeled tasks, such as text classification, natural language inference (nli), and named entity recognition (ner). despite their outstanding performances, their industrial usage is still limited * corresponding author: xlwang@cs.ecnu.edu.cn. by the high latency during inference @xcite . in addition, a special feature of consumer queries is that there are time intervals when the number of queries is exceptionally high. for example, food search engines will be used more often during dinner hours than usual. thus, deployed pre-trained models need to adjust their latency dynamically. a branch of literature focuses on making plms' inference more efficient via adaptive inference @xcite @xcite . the idea of adaptive inference is to process simple examples with only shallow layers of bert and more difficult queries with deeper layers, thus significantly speeding up the inference time on average while maintaining high accuracy. early exiting is one of the essential adaptive inference methods @xcite . as depicted in figure 1 , it implements adaptive inference by installing an early exit, i.e., an intermediate prediction layer, at each layer of plm (multi-exit plm) and early exiting "easy" samples to speed up inference. all the exits are jointly optimized at the training stage with bert's parameters. at the inference stage, an early exiting strategy is designed to decide whether to exit at each layer given the currently obtained predictions (from previous and current layers) @xcite @xcite . in this mode, different samples can exit at different depths. the average speedup ratio can be easily controlled with certain hyper-parameters without redeploying the model services or maintaining a group of models. in this work, we propose a novel framework, badge, to improve the early exiting performances of plms. badge consists of two modifications to the current early exiting literature. first, we propose to add a block-wise bypass to each transformer block so that two different representations can be produced, one for the current layer's exit and the other passed to the next transformer block. with this mechanism, an intermediate transformer block will not be distracted by the conflicting tasks in multi-exit bert training, thus providing high-quality representations for the subsequent layers. as a result, the overall performance of the multi-exit bert will improve. second, we propose a novel divergence-based early exiting method (dge), which is effective in mining early exiting signals by comparing the predicted probability distributions of the current and previous layers. extensive experiments and ablation studies are conducted on the glue benchmark datasets. the experimental results show that: (a) training multiexit plm with block-wise bypasses consistently performs better than the previous sota multiexit model training methods, thus providing a better backbone for early exiting. (b) we show that, with the same multi-exit plm backbone, our dge method can provide better efficiency-performance trade-offs than the previous sota early exiting methods. thus, with our framework, badge, a plm can achieve significantly better early exiting performances. the main contributions of our badge framework are two-fold: ## related work due to the length limit, readers are referred to appendix a for more related works on more inference speedup methods and dynamic early exiting mechanisms. ## conclusion in this work, we propose a novel framework, badge, to improve the early exiting of pre-trained language models. first, we propose to add blockwise bypasses to facilitate the joint training of the intermediate exits, thus improving the overall performance of multi-exit plms. second, we propose a novel divergence-based early exiting method, dge, which can effectively mine the exiting signals by comparing the predicted distributions of the current and previous intermediate layers. dge achieves better efficiency-performance trade-offs than the previous sota early exiting methods under the same multi-exit backbone. our badge is off-the-shelf and can effectively speed up the inferences of plms with less performance degradation.