SlowGuess/NAACL_2021_OCR huggingface dataset:
(1)  $\mathrm{TR - BERT}_{12}$  achieves higher performance while using less computation on all span-extraction QA datasets compared to all the baselines. For example,  $\mathrm{TR - BERT}_{12}$  outperforms DynaBERT by 1.8 F1 with faster speed.  $\mathrm{TR - BERT}_{12}$  even achieves better performance than BERT at low speedup rate, which demonstrates that discarding some redundant information in the top layer helps to find the correct answer. For multiple-choice RACE,  $\mathrm{TR - BERT}_{12}$  achieves better performance than DeFormer while doesn't need to pre-compute the passage representation.
(2)  $\mathrm{TR - BERT}_6$  performs better than PoWER-BERT by a margin in text classification tasks. It shows that the fixed pruning configuration and the attention-based selection strategy adopted by PoWER-BERT may not be flexible to accelerate inference for various input sequences. In contrast,
Figure 3: Quality and efficiency trade-offs for TR-BERT $_{12}$  and TR-BERT $_{6}$ .
our dynamic token selection can automatically determine the proper pruning length and tokens for each example according to the actual situation, which leads to a more effective model acceleration.
Overall, TR-BERT retains most of BERT's performance though it omits lots of token interactions in the top layers. It shows that TR-BERT learns a satisfactory token selection strategy through reinforcement learning, and could effectively reduce the redundant computation of tokens that have been extracted enough information in the bottom layers.
4.5 Fuse Layer-wise and Token-wise Pruning
Since layer-wise pruning and token-wise pruning are compatible, we also explore the incorporation of these two pruning strategies. We apply our dynamic token reduction on the 6-layer DistilBERT to obtain TR-BERT $_6$ . The trade-off comparison of
Model	HotpotQA	TriviaQA
F1	FLOPs	F1	FLOPs
BERT	57.33	1.00x	68.75	1.00x
BERTL	65.45	0.91x	69.69	0.92x
TR-BERTL	65.57	1.56x	70.41	1.20x
Model	WikiHop	Hyperparisan
Acc.	FLOPs	Acc.	FLOPs
BERT	67.67	1.00x	74.39	1.00x
BERTL	68.49	0.92x	76.83	0.92x
TR-BERTL	67.78	4.65x	74.39	1.96x
Table 3: Comparison performance and FLOPs (speedup) between TR-BERT $_L$  and BERTs with different maximum input sequence.
TR-BERT $_{12}$  and TR-BERT $_{6}$  is shown in Figure 3, from which we have the following findings:
(1) In general, as the speedup ratio increases, the performance of all models decreases, which indicates that retaining more token information usually results in a more potent model.
(2) TR-BERT $_6$  consistently outperforms TR-BERT $_{12}$  on all tasks at a high speedup ratio. In this situation, the budget doesn't allow enough tokens to go through the top layers. TR-BERT $_6$  makes a more elaborate pruning than TR-BERT $_{12}$  at bottom layers to obtain a better effectiveness.
(3) At low speedup ratio,  $\mathrm{TR - BERT_{12}}$  performs better than  $\mathrm{TR - BERT_6}$  on the question answering tasks, but worse on the text classification tasks. In general, a deep Transformer architecture can offer multi-turn feature extraction and information propagation, which can meet the complex reasoning requirements for question answering. In contrast, the result of text classification usually depends on the keywords in the context, for which a shallow model is an affordable solution. To obtain a better trade-off, we can flexibly employ a deep and narrow model for question answering and a shallow and wide model for text classification.
4.6 Results on Long-text Tasks
With token pruning, TR-BERT is able to process a longer sequence. We apply our dynamic token pruning strategy on  $\mathrm{BERT}_{\mathrm{L}}$ , which can process sequence with up to 1,024 tokens, to obtain TR-BERT $_{\mathrm{L}}$ , and conduct experiments on four datasets with longer documents, including HotpotQA, TriviaQA, WikiHop and Hyperparisan. Results on long-text tasks are shown in Table 3, from which we have the following observations:
(1)  $\mathrm{BERT_L}$  achieves better performance than BERT, especially on HotpotQA and WikiHop,
which require the long-range multi-hop reasoning;
(2) Compared to the vanilla BERT, TR-BERT $_{\mathrm{L}}$  achieves  $8.2\%$  F1 improvement with 1.56x speedup on HotpotQA, obtains  $1.7\%$  F1 improvement with 1.24x speedup on TriviaQA, gains 4.65x speedup on WikiHop and 1.96x speedup on Hyperparisan without performance drops. Compared to BERT which can only deal with up to 512 tokens at a time,  $\mathrm{BERT}_{\mathrm{L}}$  considers a longer-range token interaction and obtains a more complete reasoning chain. However, the running time of  $\mathrm{BERT}_{\mathrm{L}}$  also increase as the input sequence's length extends, which poses a challenge to the utilization of longer text. TR-BERT $_{\mathrm{L}}$  inherits the broader view from  $\mathrm{BERT}_{\mathrm{L}}$  to get a better performance with a faster inference. Moreover, the inference acceleration effect of TR-BERT $_{\mathrm{L}}$  is relatively better than TR-BERT within 512 tokens, which is coincident to the above complexity analysis section. With a longer sequence, TR-BERT can achieve extra speedup, because it significantly saves the time of the Self-Attention module, which demonstrates that TR-BERT can be further applied to process much longer tokens with limited computation.
To investigate the characteristics of the selected tokens, we conduct a detailed case study on various datasets. As shown in Table 4, TR-BERT chooses to abandon the function word, such as the, and, with, in the first token reduction module as the first module is placed at the bottom layer of BERT. The second token reduction module is placed at the middle layer of BERT, and we could observe that it is used to retaining task-specific tokens. In the first example about question answering, the second token reduction module maintains the whole question and the question-related tokens from the context for further propagating messages. In the second and third examples about movie review sentimental classification, the second token reduction module chooses to select sentimental words, such as great, excited, disappointed to determine whether the given sequence is positive or negative.
Although we train the token reduction module without direct human annotations, TR-BERT can remain the meaningful tokens in the bottom layer and select the higher layer's task-relevant tokens. It demonstrates that the pruned network's ground-truth probability is an effective signal to facilitate the reinforcement learning for token selection.
Dataset	Example
SQuAD 2.0	Question: How long did Western Europe control Cyprus?
Paragraph: The conquest of Cyprus by the Anglo-Norman forces of the Third Crusade opened a new chapter in the history of the island, which would be under Western European domination for the following 380 years. Although not part of a planned operation, the conquest had much more permanent results than initially expected.
Answer: 380 years
IMDB Positive	The Buddy Holly Story is a great biography with a super performance from Gary Buse . Busey did his own singing for this film and he does a great job .
IMDB Negative	I was really excited when I read “The Canterville Ghost” would be shown on TV. However, I was deeply disappointed. I loved the original story written by Oscar Wilde and sadly nothing of that was transferred by the movie .
Table 4: Case study of the token reduction. Skipped tokens in first token reduction module are colored with light blue. Skipped tokens in the second token reduction module are colored with blue. The final selected token are colored with dark blue.
Researchers have made various attempts to accelerate the inference of PLMs, such as quantization (Shen et al., 2020; Zhang et al., 2020a), attention head pruning (Michel et al., 2019; Hou et al., 2020), dimension reduction (Sun et al., 2020; Chen et al., 2020), and layer reduction (Sanh et al., 2019; Sun et al., 2019b; Jiao et al., 2019). 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 (Fan et al., 2020; Xin et al., 2020; Liu et al., 2020). 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 self-attention mechanism of the Transformer. For example, Sparse Trasformer (Child et al., 2019), LongFormer (Beltagy et al., 2020) and Big Bird (Zaheer et al., 2020) 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 (Dai et al., 2020) reduces the sequence length with pooling for less computation, and finally up-samples it to the full-length representation. Universal Transformer (Dehghani et al., 2019) builds a self-attentive recurrent sequence model, where each token uses a dynamic halting layer. And DynSAN (Zhuang and Wang, 2019) 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. IDLSTM (Zhang et al., 2018) selects important and task-relevant words to build sentence representation for text classification. SR-MRS (Nie et al., 2019) 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.
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.
This research is mainly supported by Science & Tech Innovation 2030 Major Project "New Generation AI" (Grant no. 2020AAA0106500) as well as supported in part by a grant from the Institute for Guo Qiang, Tsinghua University.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. CoRR, abs/2004.05150.
Qingqing Cao, Harsh Trivedi, Aruna Balasubramanian, and Niranjan Balasubramanian. 2020. Deformer: Decomposing pre-trained transformers for faster question answering. In Proceedings of ACL, pages 4487-4497.
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. 2020. The lottery ticket hypothesis for pretrained BERT networks. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long sequences with sparse transformers. CoRR, abs/1904.10509.
Christopher Clark and Matt Gardner. 2018. Simple and effective multi-paragraph reading comprehension. In Proceedings of ACL, pages 845-855.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoc V. Le. 2020. Funnel-transformer: Filtering out sequential redundancy for efficient language processing. CoRR, abs/2006.03236.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. 2019. Universal transformers. In Proceedings of ICLR.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, pages 4171-4186.
Angela Fan, Edouard Grave, and Armand Joulin. 2020. Reducing transformer depth on demand with structured dropout. In Proceedings of ICLR. OpenReview.net.
Adam Fisch, Alon Talmor, Robin Jia, Minjoon Seo, Eunsol Choi, and Danqi Chen. 2019. MRQA 2019 shared task: Evaluating generalization in reading comprehension. pages 1-13.
Saurabh Goyal, Anamitra Roy Choudhury, Venkatesan T. Chakaravarthy, Saurabh ManishRaje, Yogish Sabharwal, and Ashish Verma. 2020. Powerbert: Accelerating BERT inference for classification tasks.
Dan Hendrycks and Kevin Gimpel. 2017. A baseline for detecting misclassified and out-of-distribution examples in neural networks. In Proceedings of ICLR.
John Hewitt and Christopher D. Manning. 2019. A structural probe for finding syntax in word representations. In Proceedings of NAACL-HLT, pages 4129-4138.
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015. Distilling the knowledge in a neural network. CoRR, abs/1503.02531.
Lu Hou, Lifeng Shang, Xin Jiang, and Qun Liu. 2020. Dynabert: Dynamic BERT with adaptive width and depth. CoRR, abs/2004.04037.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. 2017. Imitation learning: A survey of learning methods. ACM Computing Survey, 50(2):21:1-21:35.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2019. Tinybert: Distilling BERT for natural language understanding. CoRR, abs/1909.10351.
Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. 2017. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of ACL, pages 1601-1611.
Johannes Kiesel, Maria Mestre, Rishabh Shukla, Emmanuel Vincent, Payam Adineh, David Corney, Benno Stein, and Martin Potthast. 2019. Semeval-2019 task 4: Hyperpartisan news detection. In Proceedings of Workshop on NAACL-HLT, pages 829-839.
Gyuwan Kim and Kyunghyun Cho. 2020. Length-adaptive transformer: Train once with length drop, use anytime with search. CoRR, abs/2010.07003.
Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In Proceedings of ICLR.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur P. Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. 2019. Natural questions: a benchmark for question answering research. TACL, 7:452-466.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard H. Hovy. 2017. RACE: large-scale reading comprehension dataset from examinations. In Proceedings of EMNLP, pages 785-794.
Ken Lang. 1995. Newsweeder: Learning to filter netnews. In Machine Learning, Proceedings of the Twelfth International Conference on Machine Learning, Tahoe City, California, USA, July 9-12, 1995, pages 331-339. Morgan Kaufmann.
Yongjie Lin, Yi Chern Tan, and Robert Frank. 2019. Open sesame: Getting inside bert's linguistic knowledge. CoRR, abs/1906.01698.
Weijie Liu, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. 2020. Fastbert: a self-distilling BERT with adaptive inference time. In Proceedings of ACL, pages 6035-6044.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. RoBERTa: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11692.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. 2011. Learning word vectors for sentiment analysis. In Proceedings of ACL, pages 142-150.
Paul Michel, Omer Levy, and Graham Neubig. 2019. Are sixteen heads really better than one? In Proceedings of NeurIPS, pages 14014-14024.
Yixin Nie, Songhe Wang, and Mohit Bansal. 2019. Revealing the importance of semantic retrieval for machine reading at scale. In Proceedings of EMNLP, pages 2553-2566.
Pranav Rajpurkar, Robin Jia, and Percy Liang. 2018. Know what you don't know: Unanswerable questions for squad. In Proceedings of ACL, pages 784-789.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. SQuAD: 100, 000+ questions for machine comprehension of text. In Proceedings of EMNLP, pages 2383-2392.
Anna Rogers, Olga Kovaleva, and Anna Rumshisky. 2020. A primer in bertology: What we know about how BERT works. Trans. Assoc. Comput. Linguistics, 8:842-866.
Victor Sanh, Lysandre Debut, Julien Chaumont, and Thomas Wolf. 2019. Distilbert, a distilled version of BERT: smaller, faster, cheaper and lighter. CoRR, abs/1910.01108.
Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W. Mahoney, and Kurt Keutzer. 2020. Q-BERT: hessian based ultra low precision quantization of BERT. In Proceedings of AAAI, pages 8815-8821.
Chi Sun, Luyao Huang, and Xipeng Qiu. 2019a. Utilizing BERT for aspect-based sentiment analysis via constructing auxiliary sentence. In Proceedings of NAACL-HLT, pages 380-385.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019b. Patient knowledge distillation for BERT model compression. In Proceedings of ACL, pages 4322-4331.
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. 2020. Mobilebert: a compact task-agnostic BERT for resource-limited devices. In Proceedings of ACL, pages 2158-2170.
Richard S. Sutton, David A. McAllester, Satinder P. Singh, and Yishay Mansour. 1999. Policy gradient methods for reinforcement learning with function approximation. In Proceedings of NeurIPS, pages 1057-1063.
Alon Talmor and Jonathan Berant. 2019. MultiQA: An empirical investigation of generalization and transfer in reading comprehension. In Proceedings of ACL, pages 4911-4921.
Adam Trischler, Tong Wang, Xingdi Yuan, Justin Harris, Alessandro Sordoni, Philip Bachman, and Kaheer Suleman. 2017. Newsqa: A machine comprehension dataset. In Proceedings of the 2nd Workshop on Representation Learning for NLP, Rep4NLP@ACL 2017, Vancouver, Canada, August 3, 2017, pages 191-200.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proceeding of NeurIPS, pages 5998-6008.
Johannes Welbl, Pontus Stenetorp, and Sebastian Riedel. 2018. Constructing datasets for multi-hop reading comprehension across documents. Trans. Assoc. Comput. Linguistics, 6:287-302.
Ronald J. Williams. 1992. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach. Learn., 8:229-256.
Ji Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. 2020. DeeBERT: Dynamic early exiting for accelerating BERT inference. In Proceedings of ACL, pages 2246-2251.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. 2019. Xlnet: Generalized autoregressive pretraining for language understanding. In Proceedings of NeurIPS, pages 5754-5764.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. 2018. Hotpotq: A dataset for diverse, explainable multi-hop question answering. In Proceedings of EMNLP, pages 2369-2380.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. 2020. Big bird: Transformers for longer sequences. CoRR, abs/2007.14062.
Tianyang Zhang, Minlie Huang, and Li Zhao. 2018. Learning structured representation for text classification via reinforcement learning. In Proceedings of AAAI, pages 6053-6060. AAAI Press.
Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. 2020a. Ternarybert: Distillation-aware ultra-low bit BERT. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 509-521. Association for Computational Linguistics.
Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. 2015. Character-level convolutional networks for text classification. In Proceedings of NeurIPS, pages 649-657.
Zhuosheng Zhang, Yuwei Wu, Hai Zhao, Zuchao Li, Shuaijiang Zhang, Xi Zhou, and Xiang Zhou. 2020b. Semantics-aware BERT for language understanding. In Proceedings of AAAI, pages 9628-9635.
Chen Zhao, Chenyan Xiong, Corby Rosset, Xia Song, Paul N. Bennett, and Saurabh Tiwary. 2020. Transformer-xh: Multi-evidence reasoning with extra hop attention. In Proceedings of ICLR.
Yimeng Zhuang and Huadong Wang. 2019. Token-level dynamic self-attention network for multipassage reading comprehension. In Proceedings of ACL, pages 2252-2262.
We evaluate models on seven question-answering datasets, including SQuAD 2.0 (Rajpurkar et al., 2018), NewsQA (Trischler et al., 2017), NaturalQA (Kwiatkowski et al., 2019), RACE (Lai et al., 2017), HotpotQA (Yang et al., 2018), TriviaQA (Joshi et al., 2017) and WikiHop (Welbl et al., 2018). Besides, we evaluate models on four long-text classification datasets, including YELP.F (Zhang et al., 2015), 20NewsGroups (Lang, 1995), IMDB (Maas et al., 2011),
and Hyperpartisan (Kiesel et al., 2019). We use the MRQA (Fisch et al., 2019) version of NewsQA and NaturalQA. Details of all evaluation datasets are shown below:
SQuAD 2.0 (Rajpurkar et al., 2018) is a large-scale reading comprehension dataset. Compared to its former SQuAD 1.0 (Rajpurkar et al., 2016), SQuAD 2.0 involves 54k unanswerable questions, which empirically makes the task much harder. For SQuAD 2.0, TR-BERT not only needs to find the question-relevant tokens, but also requires to check sufficient evidence to make a waiver decision when no answer is available. To predict the answer, we attach a span predictor on the top of BERTs and set the answer of the unanswerable question as a span of  $[0,0]$ .
MRQA (Fisch et al., 2019) integrates several existing datasets to a unified format, which provides a single context within 800 tokens for each question, ensuring at least one answer could be accurately found in the context. We adopt the NewsQA and NaturalQA dataset from the MRQA benchmark.
RACE (Lai et al., 2017) is composed of 98k multiple-choice questions collected from English examinations. The model needs to figure out the correct answer from four options for a given question and passage. Passages in RACE cover a variety of topics, which can examine the generalization of our token selection.