TR-BERT: Dynamic Token Reduction for Accelerating BERT Inference

Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5798–5809 June 6–11, 2021. ©2021 Association for Computational Linguistics 5798 TR-BERT: Dynamic Token Reduction for Accelerating BERT Inference Deming Ye1,2, Yankai Lin4, Yufei Huang1,2, Maosong Sun1,3∗ 1Department of Computer Science and Technology, Tsinghua University, Beijing, China Institute for Artificial Intelligence, Tsinghua University, Beijing, China Beijing National Research Center for Information Science and Technology 2State Key Lab on Intelligent Technology and Systems, Tsinghua University, Beijing, China 3Beijing Academy of Artificial Intelligence 4Pattern Recognition Center, WeChat AI, Tencent Inc. ydm18@mails.tsinghua.edu.cn Abstract Existing pre-trained language models (PLMs) are often computationally expensive in in- ference, making them impractical in various resource-limited real-world applications. To address this issue, we propose a dynamic token reduction approach to accelerate PLMs’ infer- ence, named TR-BERT, which could flexibly adapt the layer number of each token in infer- ence to avoid redundant calculation. Specially, TR-BERT formulates the token reduction pro- cess 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 de- mands. 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. 1 Introduction Large-scale pre-trained language models (PLMs) such as BERT (Devlin et al., 2019), XLNet (Yang et al., 2019) and RoBERTa (Liu et al., 2019) have shown great competence in learning contextual rep- resentation of text from large-scale corpora. With appropriate fine-tuning on labeled data, PLMs have achieved promising results on various NLP appli- cations, such as natural language inference (Zhang et al., 2020b), text classification (Sun et al., 2019a) and question answering (Talmor and Berant, 2019). Along with the significant performance improve- ments, PLMs usually have substantial compu- tational cost and high inference latency, which presents challenges to their practicalities in resource-limited real-world applications, such as ∗Corresponding author: M. Sun (sms@tsinghua.edu.cn) 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 quadrat- ically with the sequence length. Therefore, re- searchers 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 (Sanh et al., 2019; Sun et al., 2019b), and adaptive inference models (Xin et al., 2020; Liu et al., 2020). Such layer-wise pruning reduces a tremendous amount of computation, but it sacrifices the models’ capa- bility in complex reasoning. Previous works (Sanh et al., 2019; Sun et al., 2019b) have found that the shallow model usually performs much worse on the relatively complicated question answering tasks than text classification tasks. It is straightfor- ward that pruning the entire layer of PLMs may not be an …

Content truncated to 500 words