Accelerating BERT Inference for Sequence Labeling via Early-Exit (2021)

Abstract: Both performance and efficiency are crucial factors for sequence labeling tasks in many real-world scenarios. Although the pre-trained models (PTMs) have significantly improved the performance of various sequence labeling tasks, their computational cost is expensive. To alleviate this problem, we extend the recent successful early-exit mechanism to accelerate the inference of PTMs for sequence labeling tasks. However, existing early-exit mechanisms are specifically designed for sequence-level tasks, rather than sequence labeling. In this paper, we first propose a simple extension of sentence-level early-exit for sequence labeling tasks. To further reduce the computational cost, we also propose a token-level early-exit mechanism that allows partial tokens to exit early at different layers. Considering the local dependency inherent in sequence labeling, we employed a window-based criterion to decide for a token whether or not to exit. The token-level early-exit brings the gap between training and inference, so we introduce an extra self-sampling fine-tuning stage to alleviate it. The extensive experiments on three popular sequence labeling tasks show that our approach can save up to 66%∼75% inference cost with minimal performance degradation. Compared with competitive compressed models such as DistilBERT, our approach can achieve better performance under the same speed-up ratios of 2×, 3×, and 4×.

## introduction sequence labeling plays an important role in natural language processing (nlp). many nlp tasks can be converted to sequence labeling tasks, such as named entity recognition, part-of-speech tagging, chinese word segmentation and semantic role labeling. these tasks are usually fundamental and highly time-demanding, therefore, apart from performance, their inference efficiency is also very important. the past few years have witnessed the prevailing of pre-trained models (ptms) @xcite on various sequence labeling tasks @xcite @xcite . despite their significant improvements on sequence labeling, they are notorious for enormous computational cost and slow inference speed, which hinders their utility in real-time scenarios or mobile-device scenarios. recently, early-exit mechanism @xcite @xcite has been introduced to accelerate inference for large-scale ptms. in their methods, each layer of the ptm is coupled with a classifier to predict the label for a given instance. at inference stage, if the prediction is confident foot_0 enough at an earlier time, it is allowed to exit without passing through the entire model. figure 1(a) gives an illustration of early-exit mechanism for text classification. however, most existing early-exit methods are targeted at sequence-level prediction, such as text classification, in which the prediction and its confidence score are calculated over a sequence. therefore, these methods cannot be directly applied to sequence labeling tasks, where the prediction is token-level and the confidence score is required for each token. in this paper, we aim to extend the early-exit mechanism to sequence labeling tasks. first, we proposed the sentence-level early-exit (sen-tee), which is a simple extension of existing earlyexit methods. sentee allows a sequence of tokens to exit together once the maximum uncertainty ... it is not good it is not good it is not good layer (i+1) layer (i+2) layers (>i+2) confident? yes😎 confident? no😕 early-exit prediction confident? no😕 layer (i) negative negative positive (a) early-exit for text classification confident? (pooling) confident? (pooling) no😕 confident? (pooling) yes😎 early-exit prediction dell founded dell in 1984 dell founded dell in 1984 dell founded dell in 1984 no😕 b-org😕 b-org 😕 o 😎 b-time 😎 o 😎 b-org😎 b-per 😕 o 😎 b-time 😎 o 😎 b-org😎 b-per 😎 o 😎 b-time 😎 o 😎 ... layer (i+1) layer (i+2) layers (>i+2) layer (i) (b) sentence-level early-exit for sequence labeling dell founded dell in 1984 dell founded dell in 1984 dell founded dell in 1984 early-exit prediction (c) token-level early-exit for sequence labeling figure 1 : early-exit for text classification and sequence labeling, where represents the self-attention connection, represents simply copying, 😎 and 😕 represents the confident and uncertain prediction. the darker the hidden state block, the deeper layer it is from. due to space limit, the window-based uncertainty is not reflected. ## bert for sequence labeling recently, ptms @xcite have become the mainstream backbone model for various sequence labeling tasks. the typical framework consists of a backbone encoder and a task-specific decoder. encoder in this paper, we use bert @xcite as our backbone encoder . the architecture of bert consists of multiple stacked transformer layers @xcite . given a sequence of tokens x 1 , • • • , x n , the hidden state of l-th transformer layer is denoted by decoder usually, we can predict the label for each token according to the hidden state of the top layer. the probability of labels is predicted by ## early-exit for sequence labeling the inference speed and computational costs of ptms are crucial bottlenecks to hinder their application in many real-world scenarios. in many tasks, the representations at an earlier layer of ptms are usually adequate to make a correct prediction. therefore, early-exit mechanisms @xcite @xcite are proposed to dynamically stop inference on the backbone model and make prediction with intermediate representation. however, these existing early-exit mechanisms are built on sentence-level prediction and unsuitable for token-level prediction in sequence labeling tasks. in this section, we propose two earlyexist mechanisms to accelerate the inference for sequence labeling tasks. ## related work ptms are powerful but have high computational cost. to accelerate them, many attempts have been made. a kind of methods is to reduce its size, such as distillation @xcite , structural pruning @xcite and quantization @xcite . another kind of methods is early-exit, which dynamically adjusts the encoding layer number of different samples @xcite @xcite @xcite . while they introduced early-exit mecha-nism in simple classification tasks, our methods are proposed for the more complicated scenario: sequence labeling, where it has not only one prediction probability and it's necessary to consider the dependency of token exitings. @xcite proposed depth-adaptive transformer to accelerate machine translation. however, their earlyexit mechanism is designed for auto-regressive sequence generation, in which the exit of tokens must be in left-to-right order. therefore, it is unsuitable for language understanding tasks. different from their method, our early-exit mechanism can consider the exit of all tokens simultaneously. ## conclusion and future work in this work, we propose two early-exit mechanisms for sequence labeling: sentee and to-kee. the former is a simple extension of sequencelevel early-exit while the latter is specially designed for sequence labeling, which can conduct more finegrained computational cost allocation. we equip tokee with window-based uncertainty and selfsampling finetuning to make it more robust and faster. the detailed analysis verifies their effectiveness. sentee and tokee can achieve 2× and 3∼4× speedup with minimal performance drop. for future work, we wish to explore: (1) leveraging the exited token's label information to help the exiting of remained tokens; (2) introducing crf or other global decoding methods into early-exit for sequence labeling.