TR - BERT : Dynamic Token Reduction for Accelerating BERT Inference (2021)

Abstract: Existing pre-trained language models (PLMs) are often computationally expensive in inference, making them impractical in various resource-limited real-world applications. To address this issue, we propose a dynamic token reduction approach to accelerate PLMs’ inference, named TR-BERT, which could flexibly adapt the layer number of each token in inference to avoid redundant calculation. Specially, TR-BERT formulates the token reduction process as a multi-step token selection problem and automatically learns the selection strategy via reinforcement learning. The experimental results on several downstream NLP tasks show that TR-BERT is able to speed up BERT by 2-5 times to satisfy various performance demands. Moreover, TR-BERT can also achieve better performance with less computation in a suite of long-text tasks since its token-level layer number adaption greatly accelerates the self-attention operation in PLMs. The source code and experiment details of this paper can be obtained from https://github.com/thunlp/TR-BERT .

## introduction large-scale pre-trained language models (plms) such as bert @xcite , xlnet @xcite and roberta @xcite have shown great competence in learning contextual representation of text from large-scale corpora. with appropriate fine-tuning on labeled data, plms have achieved promising results on various nlp applications, such as natural language inference @xcite ), text classification @xcite and question answering @xcite . along with the significant performance improvements, plms usually have substantial computational cost and high inference latency, which presents challenges to their practicalities in resource-limited real-world applications, such as real-time applications and hardware-constrained mobile applications. even worse, these drawbacks become more severe in long-text scenarios because self-attention operation in plms scales quadratically with the sequence length. therefore, researchers have made intensive efforts in plm's inference acceleration recently. the mainstream approach is to reduce the layer number of plms such as knowledge distillation models @xcite , and adaptive inference models @xcite . such layer-wise pruning reduces a tremendous amount of computation, but it sacrifices the models' capability in complex reasoning. previous works @xcite have found that the shallow model usually performs much worse on the relatively complicated question answering tasks than text classification tasks. it is straightforward that pruning the entire layer of plms may not be an optimal solution in all scenarios. in this paper, we introduce a dynamic token reduction method tr-bert to find out the wellencoded tokens in the layer-by-layer inference process, and save their computation in subsequent layers. the idea is inspired by recent findings that plms capture different information of words in different layers (e.g., bert focuses on the word order information @xcite in the bottom layers, obtains the syntactic information @xcite in the middle layers, and computes the task-specific information in the top layers @xcite ). hence, we could adapt different tokens to different layers according to their specific roles in the context. as shown in figure 1 , tr-bert formulates the token reduction process as a multi-step selection problem. specially, for each selection phase, tr-bert finds out the words that require high-level semantic representations, and then selects them to higher layers. the main challenge in tr-bert is how to determine each token's importance for text understanding in the token selection. it is highly task-dependent and requires to consider the correlation and redundancy among various tokens. tr-bert employs the reinforcement learning (rl) method to learn the dynamic token selection strategy automatically. after the token reduction, the rl reward involves the confidence of the classifier's prediction based on the pruned network to reflect the quality of token selection. moreover, we also add a penalty term about the number of selected tokens to the reward, by adjusting which, tr-bert can utilize the different pruning intensities in response to various performance requirements. in tr-bert, by selecting a few important tokens to go through the entire pipeline, the inference speed turns much faster and no longer grows quadratically with the sequence length. we conduct experiments on eleven nlp benchmarks. experimental results show that tr-bert can accelerate bert inference by 2-5 times to meet various performance demands, and significantly outperform previous baseline methods on question answering tasks. it verifies the effectiveness of the dynamic token reduction strategy. moreover, benefiting from the long-distance token interaction, tr-bert with 1,024 input length reaches higher performance with less inference time compared to the vanilla bert in a suite of long-text tasks. ## background and pilot analysis to investigate the potential impact of the token reduction in plms, we first introduce the trans-former architecture. after that, we conduct pilot experiments as well as empirical analyses for the lower and upper bound of the token reduction in this section. the transformer architecture @xcite has been widely adopted by the pre-trained language models (plms) for inheriting its high capacity. basically, each transformer layer wraps a self-attention module (self-att) and a feed-forward-network module (ffn) by the residual connection and layer normalization. formally, given a sequence of n words, the hidden state of the i-th layer, h @xmath0, h 2 , . . . , h n ), is computed from the previous layer state: for each transformer layer, the complexity of the self-attention module scales quadratically with the sequence length. therefore, the speed of transformer architecture will decline heavily when the sequences become longer. previous findings @xcite ) reveal that some words, such as function words, do not require high-layer modeling, since they store little information and have been well handled by plms in bottom layers. hence, selecting only the important words for high-layer computation may be a possible way to accelerate the plms' inference. to verify this assumption, we conduct a theoretical token elimination experiment in question answering (on squad 2.0 @xcite ) and text classification (on imdb @xcite ). we use the full-layer representations for the selected tokens and the early-layer representation of the deleted tokens for the prediction. to be specific, we eliminate tokens immediately after the @xmath1: random strategy (lower bound) selects tokens randomly, assuming that all tokens are equivalent for understanding. residual strategy (upper bound) directly utilizes the model prediction of the original model to guide the token selection. specially, we define a token's importance according to the influence on the model prediction when it's not selected. when substituting the r-th layer representation h r with the l-th layer representation h l (@xmath2, we define the approximate variation to model loss as the token importance: @xmath3. here, we set @xmath4. note that we could not obtain the model loss in the prediction stage. hence, the residual strategy could be viewed as an upper bound of token selection to some extent when we ignore the correlation and redundancy among the selected tokens. attention strategy is adopted by power-bert @xcite and l-adaptive @xcite . it accumulates the attention values from other tokens to a given token. it selects the tokens receiving the greatest attentions, considering them responsible for retaining and disseminating the primary information of the context. as shown in figure 2 , both attention strategy and residual strategy achieve considerable results, which demonstrates that to select important tokens is feasible for accelerating the inference of plms. besides, the residual strategy outperforms the attention strategies by a margin, especially at the low token remaining proportion (+31.8% f1 on squad 2.0 and +9.5% accuracy on imdb when selecting 10% tokens). it suggests that the accumulated attention values still cannot well reflect tokens' importance in text understanding, which requires further explorations. ## methodology in this section, we present tr-bert, which adopts a cascade token reduction to prune the bert model at token-level granularity dynamically. in a onestep token reduction process, tr-bert estimates the importance of each token, reserves the important ones, and delivers them to the higher layer. to better select important tokens for text understanding while satisfying various acceleration requirements, we employ the reinforcement learning (rl) method to automatically learn a dynamic token selection strategy. ## experiment in this section, we first introduce the baseline models and the evaluation datasets. after that, we verify the effectiveness of tr-bert on eleven nlp benchmarks. finally, we conduct a detailed analysis and case study on tr-bert to investigate the selected tokens' characteristics. ## related work researchers have made various attempts to accelerate the inference of plms, such as quantization @xcite , attention head pruning @xcite , dimension reduction @xcite , and layer reduction @xcite @xcite . in current studies, one of the mainstream methods is to dynamically select the layer number of transformer layers to make a on-demand lighter model @xcite @xcite . however, these methods operate at the whole text and they cannot perform pruning operations in a smaller granularity, such as the token-level granularity. to consider the deficiencies of layer-level pruning methods, researchers decide to seek solutions from a more meticulous perspective by developing methods to extend or accelerate the selfattention mechanism of the transformer. for example, sparse trasformer @xcite , long-former @xcite and big bird @xcite employ the sparse attention to allow model to handle long sequences. however, these methods only reduce the cuda memory but cannot be not faster than the full attention. besides, researchers also explore the feasibility of reducing the number of involved tokens. for example, funnel-transformer @xcite reduces the sequence length with pooling for less computation, and finally up-samples it to the full-length representation. universal transformer @xcite builds a self-attentive recurrent sequence model, where each token uses a dynamic halting layer. and dynsan @xcite ap-plies a gate mechanism to measure the importance of tokens for selection. spurred by these attempts and positive results, we introduce tr-bert in this study, which can creatively prune the network at the token level. to be specific, our work aims to accelerate the transformer by deleting tokens gradually as the layer gets deeper. compared with these models, tr-bert is easy to adapt to the current plms models without a significant amount of pretraining and is flexible to adjust the model speed according to different performance requirements. the main idea of tr-bert is to select essential elements and infuse more computation on them, which is widely adopted in various nlp tasks. id-lstm @xcite selects important and task-relevant words to build sentence representation for text classification. sr-mrs @xcite retrieves the question-related sentences to reduce the size of reading materials for question answering. tr-bert can be viewed as a unified framework on the transformer for the important element selection, which can be easy to be applied in wide-range tasks. ## conclusion and future work in this paper, we propose a novel method for accelerating bert inference, called tr-bert, which prunes bert at token-level granularity. specifically, tr-bert utilizes reinforcement learning to learn a token selection policy, which is able to select general meaningful tokens in the bottom layers and select task-relevant tokens in the top layers. experiments on eleven nlp tasks demonstrate the effectiveness of tr-bert as it accelerates bert inference by 2-5 times for various performance demand. besides, tr-bert achieves a better quality and speed trade-off on long-text tasks, which shows its potential to process large amounts of information in the real-world applications. in the future, we would like to attempting to apply tr-bert in the pre-training process of plms. through the automatically learned token reduction module, it is possible to reveal how bert stores syntactic and semantic information in various tokens and different layers. and it's also worth speeding up the time-consuming pre-training process.