SlowGuess/NAACL_2021_OCR huggingface dataset:
Deming Ye $^{1,2}$ , Yankai Lin $^{4}$ , Yufei Huang $^{1,2}$ , Maosong Sun $^{1,3*}$
$^{1}$ Department of Computer Science and Technology, Tsinghua University, Beijing, China  
 $^{2}$ Institute for Artificial Intelligence, Tsinghua University, Beijing, China  
 $^{3}$ Beijing National Research Center for Information Science and Technology
$^{2}$ State Key Lab on Intelligent Technology and Systems, Tsinghua University, Beijing, China  
 $^{3}$ Beijing Academy of Artificial Intelligence  
 $^{4}$ Pattern Recognition Center, WeChat AI, Tencent Inc.
ydm18@mails.tsinghua.edu.cn
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.
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 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 (Zhang et al., 2020b), text classification (Sun et al., 2019a) and question answering (Talmor and Berant, 2019).
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 (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' capability 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 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 well-encoded 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 (Lin et al., 2019) in the bottom layers, obtains the syntactic information (Hewitt and Manning, 2019) in the middle layers, and computes the task-specific information in the top layers (Rogers et al., 2020)). 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
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
Figure 1: An illustration of the TR-BERT. TR-BERT gradually compresses the text sequence as the layer gets deeper. In RL training, we sample actions of Select or Skip for each token via the policy network.
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.
2 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 (Vaswani et al., 2017) 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_{i} = (h_{1}, h_{2}, \dots, h_{n})$ , is computed from the previous layer state:
$$
\boldsymbol {M} _ {i - 1} = \operatorname {L N} \left(\boldsymbol {H} _ {i - 1} + \text {S e l f - A T T} \left(\boldsymbol {H} _ {i - 1}\right)\right),
$$
$$
\boldsymbol {H} _ {i} = \operatorname {L N} \left(\boldsymbol {M} _ {i - 1} + \operatorname {F F N} \left(\boldsymbol {M} _ {i - 1}\right)\right), \tag {1}
$$
where  $i \in [1, L]$ ,  $L$  is the number of stacked Transformer layers, LN denotes the LayerNorm layer. 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 (Rogers et al., 2020) 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.
Figure 2: Performance under different selected proportion of tokens and different heuristic strategies.
To verify this assumption, we conduct a theoretical token elimination experiment in question answering (on SQuAD 2.0 (Rajpurkar et al., 2018)) and text classification (on IMDB (Maas et al., 2011)). 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  $l = 4th$  layer and adopt the following three strategies to select the retained tokens:
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}$  ( $r > l$ ), we define the approximate variation to model loss as the token importance:  $I = \frac{\partial loss}{\partial H_{r}} (H_{r} - H_{l})$ . Here, we set  $r = 9$  since other values get a little worse results. 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 PoWERBERT (Goyal et al., 2020) and L-Adaptive (Kim and Cho, 2020). 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.
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 one-step 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.
Figure 1 shows the model architecture of TR-BERT. To inherit the high capacity from the PLMs, TR-BERT keeps the same architecture as BERT. Differently, as the layer gets deeper, TR-BERT gradually shortens the sequence length via token reduction modules, aiming to reduce the computational redundancy of unimportant tokens.
The token reduction modules are required to measure the importance of tokens and offer an integral selection scheme. Due to the lack of direct supervision, we employ the policy network for training the module, which adopts a stochastic policy and uses a delayed reward to guide the policy learning. In one-step reduction, we perform action sampling for the current sequence. The selected tokens are conveyed to the next Transformer layer for further computation. In contrast, the unselected tokens are terminated with their representation remaining unchanged. After all the actions are decided, we fetch each token's representation from the layer where it terminated, and compute the golden label's likelihood as a reward. To be specific, we introduce state, action, reward, and objective function as follows:
State State  $s_t$  consists of the token representations inherited from the previous layer before the  $t$ -th token reduction layer.
Action We adopt two alternative actions for each token, {Select, Skip}, where the token can be selected for further computation or be skipped to the
final layer. We implement the policy network as a two-layer feed-forward network with GeLU activation (Hendrycks and Gimpel, 2017):
$$
\pi \left(a _ {t} \mid s _ {t}; \boldsymbol {\theta}\right) = \sigma \left(\boldsymbol {W} _ {2} \left(\operatorname {G e L U} \left(\boldsymbol {W} _ {1} \boldsymbol {H} _ {s _ {t}} + b _ {1}\right)\right) + b _ {2}\right), \tag {2}
$$
where  $a_{t}$  denotes the action at state  $s_t$  for sequence representation  $\pmb{H}_{s_t} = \{h_1,h_2,\dots ,h_n\}$  at  $t$ -th reduction,  $\pmb {\theta} = \{\pmb {W}_1,\pmb {W}_2,b_1,b_2\}$  are trainable parameters, and  $\sigma (.)$  is sigmoid activation function. For the selected token set  $\{t_1,t_2,\ldots ,t_{n^*}\}$ , where  $n^{*}\leq n$ , we conduct a Transformer layer operation on their corresponding representations:
$$
\mathbf {H} ^ {\prime} = \operatorname {T r a n s f o r m e r} \left(\left[ h _ {t _ {1}}, h _ {t _ {2}}, \dots , h _ {t _ {n ^ {*}}} \right]\right). \tag {3}
$$
For the selected tokens, their representation  $\mathbf{H}'$  is conveyed to the next layer for further feature extraction and information aggregation. For the other skipped tokens, their representations in the current layer are regarded as their final representations.
Reward Aiming to select significant tokens for making a precise decision in the prediction layer, we adopt the likelihood of predicting the golden label as a reward. For example, when classifying the input sequence  $\mathbf{X}$ , we use the models' predicting probability of the ground-truth label  $\mathbf{Y}$  to reflect the quality of the token selection. In addition, to encourage the model to delete more redundant tokens for accelerating, we include an additional punitive term by counting the number of selected tokens. Hence, the overall reward  $R$  is defined as:
$$
R = \log \Pr (y = \mathbf {Y} | \mathbf {X}) - \lambda \cdot \sum_ {t} \left| \left\{a _ {t} = \text {S e l e c t} \right\} \right|, \tag {4}
$$
where  $\sum_{t}|\{a_{t} = Select\}|$  denotes the total number of the selected tokens in all token reduction modules, and  $\lambda$  is a harmonic coefficient to balance two reward terms.
Objective Function We optimize the policy network to maximize the expected reward. Formally, our objective function is defined as:
$$
\begin{array}{l} {J (\pmb {\theta})} = {\mathbb {E} _ {(s _ {t}, a _ {t}) \sim \pi (a _ {t} | s _ {t}; \theta)} r [ (s _ {1}, a _ {1}) \dots (s _ {T}, a _ {T}) ]} \\ = \sum_ {(s _ {1}, a _ {1}).. (s _ {T}, a _ {T})} \prod_ {t} \pi_ {\theta} (a _ {t} | s _ {t}) \cdot R, \tag {5} \\ \end{array}
$$
where  $T$  is the number of states. According to the REINFORCE algorithm (Williams, 1992) and policy gradient method (Sutton et al., 1999), we update network with the policy gradient as below:
$$
\nabla_ {\theta} J (\theta) = \sum_ {t = 1} ^ {T} R \cdot \nabla_ {\theta} \log \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right). \tag {6}
$$
Our policy network is integrated into the original Transformer network, and we train both of them simultaneously. The entire training process involves three steps:
(1) Fine-tune the PLM model for downstream tasks with the task-specific objective;
(2) Freeze all the parameters except that of the policy network, conduct reinforcement learning (RL), and update the policy network to learn token reduction strategy;
(3) Unfreeze all parameters and train the entire network with the task-specific objective and RL objective simultaneously.
Due to the large searching space, RL learning is difficult to converge. We adopt imitation learning (Hussein et al., 2017) for warming up the training of the policy network. To be specific, in the RL training, we sample several action sequences via the policy network to compute rewards. And we guide the optimization direction by providing heuristic action sequences sampled by the Residual Strategy during the early training period, which could roughly select the most important tokens. The heuristic action sequence is defined as selecting the top  $K$  important tokens and skipping the others, where  $K$  is defined as the expected selected number of the current policy network. In our preliminary experiment, both the heuristic action sequence and expected selected number mechanism are beneficial to the stable training.
To further improve the performance of our pruned model, we also adopt Knowledge Distillation (KD) (Hinton et al., 2015) to transfer knowledge from the intact original fine-tuned model.
For a Transformer layer with a hidden size of  $d$  and an input sequence of  $n$  tokens, the Self-Attention module consumes  $O(n^{2}d)$  time and memory complexity while the Feed-Forward Network takes  $O(nd^{2})$ . That is, our token reduction gains near-linear speedup when  $n$  is relatively smaller than  $d$ . Therefore, when the input sequence gets longer, such as up to 1,024 tokens, our method can enjoy a more effective speedup.
In the RL training, we compute loss on the pruned model, so the acceleration is still valid for this stage. Since we focus on accelerating BERT inference, we consider the extra training consumption on the pruned model is acceptable.
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.
We adopt two pre-trained models and three pruned networks as our baselines for comparison:
BERT (Devlin et al., 2019) is a Transformer-based pre-trained model. We use the  $\mathrm{BERT}_{\mathrm{BASE}}$  model, which consists of 12 Transformer layers and supports a maximum sequence length of 512.
$\mathbf{BERT}_{\mathbf{L}}$  is our implemented BERT, which can support input sequences with up to 1,024 tokens. We initialize the parameters of  $\mathrm{BERT}_{\mathrm{L}}$  with that of BERT, where the additional position embedding is initialized with the first 512 ones. After that, we continue to train it on Wikipedia<sup>2</sup> for 22k steps.
DistilBERT (Sanh et al., 2019) is the most popular distilled version of BERT, which leverages the knowledge distillation to learn knowledge from the BERT model. We use the 6-layer DistilBERT released by Hugging Face<sup>3</sup>. In addition, we use the same method to distill BERT with 3 layers to obtain DistilBERT<sub>3</sub>.
DeFormer (Cao et al., 2020) is designed for question answering, which encodes questions and passages separately in lower layers. It precomputes all the passage representation and reuses them to speed up the inference. In our experiments, we do not count DeFormer's pre-computation.
PoWER-BERT (Goyal et al., 2020) is mainly designed for text classification, which also decreases the length of a sequence as layer increases. It adopts the Attention Strategy to measure the significance of each token and always selects tokens with the highest attention. Given a length penalty, PoWER-BERT searches a fixed length pruning configuration for all examples.
DynaBERT (Hou et al., 2020) can not only adjust model's width by varying the number of attention heads, but also provide an adaptive layer depth to satisfy different requirements. For a given speed demand, we report its best performance with all the feasible width and depth combination options.
To verify the effectiveness of reducing the sequence length, we evaluate TR-BERT on several tasks with relatively long context, including question answering and text classification. Table 1 shows the context length of these datasets. We adopt 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). And we also evaluate models on four text classification datasets, including YELP.F (Zhang et al., 2015), IMDB (Maas et al., 2011), 20NewsGroups (20News.) (Lang, 1995), and Hyperpartisan (Hyperp.) (Kiesel et al., 2019). Among them, HotpotQA, TriviaQA and WikiHop possess abundant contexts for reading, while the performance of question answering (QA) models heavily relies on the amount of text they read. To fairly compare BERT and  $\mathrm{BERT_L}$ , we split the context into slices and apply a shared-normalization training objective (Clark and Gardner, 2018) to produce a global answer candidate comparison across different slices for the former two extractive QA datasets. And we average the candidate scores in all slices for WikiHop. Details of all datasets are shown in the Appendix.
We adopt a maximum input sequence length of 384 for SQuAD 2.0, 1,024 for long-text tasks and 512 for others. We use the Adam optimizer (Kingma and Ba, 2015) to train all models. The detailed training configuration is shown in the Appendix.
For the RL training, we sample 8 action sequences each time and average their rewards as the reward baseline. In the second training process which aims to warm up the policy network, we employ  $20\%$  imitation learning steps for question answering tasks and  $50\%$  steps for text classification tasks. We search the number of token reduction module  $T \in [1,2,3]$ . And we find the models with  $T = 2$  gets similar quality and speed trade-offs as the models with  $T = 3$ , and both of them perform better than models with  $T = 1$ . Thus we adopt  $T = 2$  for simplification. We denote the pruned models from BERT,  $\mathrm{BERT_L}$  and DistilBERT as TR-BERT $_{12}$ , TR-BERT $_{L}$ , TR-BERT $_{6}$ , respectively. For BERT and  $\mathrm{BERT_L}$ , we attach the token reduction modules before the second and the sixth layers.
$^{2}$ https://en.wikipedia.org/
# Tokens	SQuAD 2.0	NewsQA	NaturalQA	RACE	HotpotQA	TriviaQA	WikiHop	YELP.F	20News.	IMDB	Hyperp.
Average	152	656	248	381	1,988	3,117	1,499	179	551	264	755
95th percentile	275	878	1,088	542	2,737	4,004	2,137	498	1,769	679	2,005
Table 1: Average and 95th percentile of context length of datasets in wordpieces.
Model	SQuAD 2.0	NewsQA	NaturalQA	RACE	YELP.F	20NewsGroups	IMDB
F1	FLOPs	F1	FLOPs	F1	FLOPs	Acc.	FLOPs	Acc.	FLOPs	Acc.	FLOPs	Acc.	FLOPs
BERT	77.12	1.00x	66.82	1.00x	78.32	1.00x	66.30	1.00x	69.93	1.00x	86.66	1.00x	94.05	1.00x
DistilBERT6	68.17	2.00x	63.56	2.00x	76.21	2.00x	52.63	2.00x	69.42	2.00x	85.85	2.00x	93.11	2.00x
PoWER-BERT	-	-	-	-	-	-	-	-	67.37	2.75x	86.51	2.91x	92.02	3.05x
DeFormer	71.41	2.19x	60.68	2.01x	74.34	2.34x	64.27	2.19x	-	-	-	-	-	-
DynaBERT	74.83	2.00x	66.13	2.00x	78.14	2.00x	65.38	2.00x	69.96	2.00x	86.03	2.00x	94.00	2.00x
TR-BERT12	75.66	2.08x	67.18	2.05x	79.03	2.51x	65.00	2.16x	69.97	2.19x	87.44	4.22x	93.60	2.26x
TR-BERT6	71.75	3.07x	65.36	2.96x	78.11	3.74x	53.40	4.10x	70.04	2.83x	86.58	5.90x	92.64	4.13x
Table 2: Comparison of performance and FLOPs (speedup) between TR-BERT and baselines.
For DistilBERT $_6$ , we insert the token reduction modules before the second and the fourth layers.
To avoid the pseudo improvement by pruning padding for TR-BERT, we evaluate all models with input sequences without padding to the maximum length. For each dataset, we report the F1 scores or accuracy (Acc.), and the FLOPs speedup ratio compared to the BERT model. The model's FLOPs are consistent in the various operating environment. Therefore, it is convenient to estimate and compare the models' inference time by FLOPs.
The comparison between TR-BERT and the baselines are shown in Table 2 and Figure 3. We adjust the length penalty coefficient of TR-BERT for an intuitional comparison. From the experimental results, we have the following observations: