Venue: NAACL 2021

# TR-BERT: Dynamic Token Reduction for Accelerating BERT Inference

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

Institute for Artificial Intelligence, Tsinghua University, Beijing, China

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

Corresponding author: M. Sun (sms@tsinghua.edu.cn)

###### 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\cdot 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](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 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. (2019) 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. (2019), 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. (2019) 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

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 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 Transformer 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, \(\mathbf{H}_{i}=(h_{1},h_{2},\ldots,h_{n})\), is computed from the previous layer state:

\[\mathbf{M}_{i-1} =\text{LN}(\mathbf{H}_{i-1}+\text{Self-ATT}(\mathbf{H}_{i-1})),\] \[\mathbf{H}_{i} =\text{LN}(\mathbf{M}_{i-1}+\text{FFN}(\mathbf{M}_{i-1})), \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 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.

 

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 \({\mathbf{H}}_{r}\) with the \(l\)-th layer representation \({\mathbf{H}}_{l}\) (\(r>l\)) , we define the approximate variation to model loss as the token importance: \({\mathbf{I}}=\frac{\partial loss}{\partial{\mathbf{H}}_{r}}({\mathbf{H}}_{r}-{\mathbf{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.

## 3 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 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.

### Model Architecture

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

Figure 2: Performance under different selected proportion of tokens and different heuristic strategies.

 

final layer. We implement the policy network as a two-layer feed-forward network with GeLU activation [17]:

\[\pi(a_{t}|s_{t};\mathbf{\theta})=\sigma(\mathbf{W}_{2}(\text{GeLU}(\mathbf{W}_{1}\mathbf{H}_{s_ {t}}+b_{1}))+b_{2}), \tag{2}\]

where \(a_{t}\) denotes the action at state \(s_{t}\) for sequence representation \(\mathbf{H}_{s_{t}}=\{h_{1},h_{2},...,h_{n}\}\) at \(t\)-th reduction, \(\mathbf{\theta}=\{\mathbf{W}_{1},\mathbf{W}_{2},b_{1},b_{2}\}\) are trainable parameters, and \(\sigma(.)\) is sigmoid activation function. For the selected token set \(\{t_{1},t_{2},...,t_{n^{*}}\}\), where \(n^{*}\leq n\), we conduct a Transformer layer operation on their corresponding representations:

\[\mathbf{H}^{\prime}=\text{Transformer}([h_{t_{1}},h_{t_{2}},\ldots,h_{t_{n^{*} }}]). \tag{3}\]

For the selected tokens, their representation \(\mathbf{H}^{\prime}\) 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}|\{a_{t}=\textit{Select}\}|, \tag{4}\]

where \(\sum_{t}|\{a_{t}=\textit{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:

\[J(\mathbf{\theta}) = \mathbb{E}_{(s_{t},a_{t})\sim\pi(a_{t}|s_{t};\theta)}r[(s_{1},a_ {1})...(s_{T},a_{T})] \tag{5}\] \[= \sum_{(s_{1},a_{1})...(s_{T},a_{T})}\prod_{t}\pi_{\theta}(a_{t}| s_{t})\cdot R,\]

where \(T\) is the number of states. According to the REINFORCE algorithm [8] and policy gradient method [9], we update network with the policy gradient as below:

\[\nabla_{\theta}J(\theta)=\sum_{t=1}^{T}R\cdot\nabla_{\theta}\log\pi_{\theta} (a_{t}|s_{t}). \tag{6}\]

### Model Training

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 [8] 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) [10] to transfer knowledge from the intact original fine-tuned model.

### Complexity Analysis

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 

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.

### Baselines

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 BERT\({}_{\text{BASE}}\) model1, which consists of \(12\) Transformer layers and supports a maximum sequence length of \(512\).

Footnote 1: [https://github.com/google-research/bert](https://github.com/google-research/bert)

**BERTL** is our implemented BERT, which can support input sequences with up to \(1{,}024\) tokens. We initialize the parameters of BERTL with that of BERT, where the additional position embedding is initialized with the first \(512\) ones. After that, we continue to train it on Wikipedia2 for \(22\)k steps.

Footnote 2: [https://en.wikipedia.org/](https://en.wikipedia.org/)

**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 Face3. In addition, we use the same method to distill BERT with \(3\) layers to obtain DistilBERT\({}_{3}\).

Footnote 3: [https://github.com/huggingface/transformers](https://github.com/huggingface/transformers)

**DeFormer**Cao et al. (2020) is designed for question answering, which encodes questions and passages separately in lower layers. It prepomutes 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 search s 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.

### Datasets

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 relys on the amount of text they read. To fairly compare BERT and BERTL, 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.

### Experimental Settings

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, BERTL and DistilBERT\({}_{6}\) as TR-BERT\({}_{12}\), TR-BERTL, TR-BERT\({}_{6}\), respectively. For BERT and BERTL, we attach the token reduction modules before the second and the sixth layers.

 

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.

### Overall Results

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:

(1) TR-BERT\({}_{12}\) achieves higher performance while using less computation on all span-extraction QA datasets compared to all the baselines. For example, TR-BERT\({}_{12}\) outperforms DynaBERT by \(1.8\) F1 with faster speed. TR-BER\({}_{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, TR-BERT\({}_{12}\) achieves better performance than DeFormer while doesn't need to pre-compute the passage representation.

(2) 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, 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.

### 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

Figure 3: Quality and efficiency trade-offs for TR-BERT\({}_{12}\) and TR-BERT\({}_{6}\).

 

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 decrease, 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, TR-BERT\({}_{12}\) performs better than 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.

### 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 BERTL, which can process sequence with up to \(1{,}024\) tokens, to obtain TR-BERT\({}_{\text{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) BERTL 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\({}_{\text{L}}\) achieves \(8.2\%\) F1 improvement with \(1.56\)x speedup on HotpotQA, obtains \(1.7\%\) F1 improvement with \(1.24\)x speedup on TriviaQA, gains \(4.65\)x speedup on WikiHop and \(1.96\)x speedup on Hyperparisan without performance drops. Compared to BERT which can only deal with up to \(512\) tokens at a time, BERTL considers a longer-range token interaction and obtains a more complete reasoning chain. However, the running time of BERTL also increase as the input sequence's length extends, which poses a challenge to the utilization of longer text. TR-BERTLIN inherits the broader view from BERTL to get a better performance with a faster inference. Moreover, the inference acceleration effect of TR-BERTLIN 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.

### Case Study

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 sentiment 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 groundtruth probability is an effective signal to facilitate the reinforcement learning for token selection.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Model & \multicolumn{2}{c}{HotpotQA} & \multicolumn{2}{c}{TriviaQA} \\  & F1 & FLOPs & F1 & FLOPs \\ \hline BERT & 57.33 & 1.00x & 68.75 & 1.00x \\ BERTTL & 65.45 & 0.91x & 69.69 & 0.92x \\
**TR-BERTTL** & **65.57** & **1.56x** & **70.41** & **1.20x** \\ \hline \multirow{2}{*}{Model} & \multicolumn{2}{c}{WikiHop} & \multicolumn{2}{c}{Hyperparisan} \\  & Acc. & FLOPs & Acc. & FLOPs \\ \hline BERT & 67.67 & 1.00x & 74.39 & 1.00x \\ BERTTL & 68.49 & 0.92x & 76.83 & 0.92x \\
**TR-BERTTL** & **67.78** & **4.65x** & **74.39** & **1.96x** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Comparison performance and FLOPs (speedup) between TR-BERTLI and BERTTs with different maximum input sequence.

 



## 5 Related Work

Researchers have made various attempts to accelerate the inference of PLMs, such as quantization Shen et al. (2020); Zhang et al. (2020), 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. (2019); 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) applies 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.

## 6 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\text{-}5\) times for various performance demand. Besides, TR-BERT achieves a better quality

\begin{table}
\begin{tabular}{l|p{347.0pt}} \hline \hline Dataset & Example \\ \hline 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 \\ \hline 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 \\ \hline 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 . \\ \hline \hline \end{tabular}
\end{table}
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.



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.

## Acknowledgement

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.

## References

* L. Beltagy, M. E. Peters, and A. Cohan (2020)Longformer: The long-document transformer. CoRRabs/2004.05150. External Links: Link Cited by: SS1.
* Q. Cao, H. Trivedi, A. Balasubramanian, and N. Balasubramanian (2020)Deformer: decomposing pre-trained transformers for faster question answering. In Proceedings of ACL, pp. 4487-4497. Cited by: SS1.
* T. Chen, J. Frankle, S. Chang, S. Liu, Y. Zhang, Z. Wang, and M. 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, Cited by: SS1.
* R. Child, S. Gray, A. Radford, and I. Sutskever (2019)Generating long sequences with sparse transformers. CoRRabs/1904.10509. External Links: Link, 1904.10509 Cited by: SS1.
* C. Clark and M. Gardner (2018)Simple and effective multi-paragraph reading comprehension. In Proceedings of ACL, pp. 845-855. Cited by: SS1.
* Z. Dai, G. Lai, Y. Yang, and Q. V. Le (2020)Funnel-transformer: filtering out sequential redundancy for efficient language processing. CoRRabs/2006.03236. External Links: Link, 2006.03236 Cited by: SS1.
* M. Dehghani, S. Gouws, O. Vinyals, J. Uszkoreit, and L. Kaiser (2019)Universal transformers. In Proceedings of ICLR, Cited by: SS1.
* J. Devlin, M. Chang, K. Lee, and K. Toutanova (2019)BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, pp. 4171-4186. Cited by: SS1.
* A. Fan, E. Grave, and A. Joulin (2020)Reducing transformer depth on demand with structured dropout. In Proceedings of ICLR, pp. vieu.net. Cited by: SS1.
* A. Fisch, A. Talmor, R. Jia, M. Seo, E. Choi, and D. Chen (2019)MRQA 2019 shared task: evaluating generalization in reading comprehension. pp. 1-13. Cited by: SS1.
* S. Goyal, A. Roy Choudhury, P. T. Chakaravarthy, S. ManishRaje, Y. Sabharwal, and A. Verma (2020)Powerbert: accelerating BERT inference for classification tasks. Cited by: SS1.
* D. Hendrycks and K. Gimpel (2017)A baseline for detecting misclassified and out-of-distribution examples in neural networks. In Proceedings of ICLR, Cited by: SS1.
* J. Hewitt and C. D. Manning (2019)A structural probe for finding syntax in word representations. In Proceedings of NAACL-HLT, pp. 4129-4138. Cited by: SS1.
* Geoffrey E. Hinton, O. Vinyals, and J. Dean (2015)Distilling the knowledge in a neural network. CoRRabs/1503.02531. External Links: Link, 1503.02531 Cited by: SS1.
* L. Hou, L. Shang, X. Jiang, and Q. Liu (2020)Dynabert: dynamic BERT with adaptive width and depth. CoRRabs/2004.04037. External Links: Link, 2004.04037 Cited by: SS1.
* A. Hussein, M. Medhat Gaber, E. Elyan, and C. Jayne (2017)Imitation learning: a survey of learning methods. ACM Computing Survey50 (2), pp. 21:1-21:35. Cited by: SS1.
* X. Jiao, Y. Yin, L. Shang, X. Jiang, X. Chen, L. Li, F. Wang, and Q. Liu (2019)TinyBert: distilling BERT for natural language understanding. CoRRabs/1909.10351. External Links: Link, 1909.10351 Cited by: SS1.
* M. Joshi, E. Choi, D. S. Weld, and L. Zettlemoyer (2017)Triviaqa: a large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of ACL, pp. 1601-1611. Cited by: SS1.
* J. Kiesel, M. Mestre, R. Shukla, E. Vincent, P. Adineh, D. Corney, B. Stein, and M. Potthast (2019)Semeval-2019 task 4: hyperpartisan news detection. In Proceedings of Workshop on NAACL-HLT, pp. 829-839. Cited by: SS1.
* Gyuwan Kim and Kyunghyun Cho (2020)Length-adaptive transformer: train once with length drop, use anytime with search. CoRRabs/2010.07003. External Links: Link, 2010.07003 Cited by: SS1.
* D. 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.
* Lai et al. (2017) 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.
* Lang (1995) 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.
* Lin et al. (2019) Yongjie Lin, Yi Chern Tan, and Robert Frank. 2019. Open sesame: Getting inside bert's linguistic knowledge. _CoRR_, abs/1906.01698.
* Liu et al. (2020) 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.
* Liu et al. (2019) 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.
* Maas et al. (2011) 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.
* Michel et al. (2019) Paul Michel, Omer Levy, and Graham Neubig. 2019. Are sixteen heads really better than one? In _Proceedings of NeurIPS_, pages 14014-14024.
* Nie et al. (2019) 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.
* Rajpurkar et al. (2018) 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.
* Rajpurkar et al. (2016) 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.
* Rogers et al. (2020) 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.
* Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. Distilbert, a distilled version of BERT: smaller, faster, cheaper and lighter. _CoRR_, abs/1910.01108.
* Shen et al. (2020) 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.
* Sun et al. (2019a) 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.
* Sun et al. (2019b) Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019b. Patient knowledge distillation for BERT model compression. In _Proceedings of ACL_, pages 4322-4331.
* Sun et al. (2020) 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.
* Sutton et al. (1999) 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.
* Talmor and Berant (2019) Alon Talmor and Jonathan Berant. 2019. MultiQA: An empirical investigation of generalization and transfer in reading comprehension. In _Proceedings of ACL_, pages 4911-4921.
* Trischler et al. (2017) 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.
* Vaswani et al. (2017) 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.
* Welbl et al. (2018) Johannes Welbl, Pontus Steentorp, and Sebastian Riedel. 2018. Constructing datasets for multi-hop reading comprehension across documents. _Trans. Assoc. Comput. Linguistics_, 6:287-302.
* Williams (1992) Ronald J. Williams. 1992. Simple statistical gradient-following algorithms for connectionist reinforcement learning. _Mach. Learn._, 8:229-256.
* Xin et al. (2020) 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.

 

* Yang et al. (2019) 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.
* Yang et al. (2018) Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. 2018. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In _Proceedings of EMNLP_, pages 2369-2380.
* Zaheer et al. (2020) Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Anislie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. 2020. Big bird: Transformers for longer sequences. _CoRR_, abs/2007.14062.
* Zhang et al. (2018) 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.
* Zhang et al. (2020a) 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.
* Zhang et al. (2015) Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. 2015. Character-level convolutional networks for text classification. In _Proceeding of NeurIPS_, pages 649-657.
* Zhang et al. (2020b) Zhuosheng Zhang, Yuwei Wu, Hai Zhao, Zuchao Li, Shuailiang Zhang, Xi Zhou, and Xiang Zhou. 2020b. Semantics-aware BERT for language understanding. In _Proceedings of AAAI_, pages 9628-9635.
* Zhao et al. (2020) Chen Zhao, Chenyan Xiong, Corby Rosset, Xia Song, Paul N. Bennett, and Saurabh Tiwary. 2020. Transformer-xhi: Multi-evidence reasoning with extra hop attention. In _Proceedings of ICLR_.
* Zhuang and Wang (2019)Yimeng Zhuang and Huadong Wang. 2019. Token-level dynamic self-attention network for multi-passage reading comprehension. In _Proceedings of ACL_, pages 2252-2262.

## Appendix A Details of Datasets

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), HotopQA 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 \(54\)k 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 \(98\)k 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.

**HotpotQA**Yang et al. (2018) is an extractive question answering dataset, which requires multi-hop reasoning over multiple supporting documents for answering \(113\)k questions. We adopt the full-wiki setting for HotpotQA, which requires models to find answers from a large-scale corpus. We employ the retriever and re-ranker in Transformer-XH Zhao et al. (2020) to obtain question-related paragraphs and merge them into one document. In HotpotQA, models are required to reason over bridge entities or check multiple properties in different positions, which brings challenges to the token selection of TR-BERT in considering the global information. We concatenate two positive paragraphs and several negative paragraphs to make the reading material for training, which contains up to \(1{,}024\) tokens. And we concatenate the re-rank passages in order for evaluating. For the vanilla BERT, we apply a shared-normalization training objective Clark and Gardner (2018) to produce a global answer candidate comparison across two parts of the document 



**TriviaQA**Joshi et al. (2017) has more than \(95\)k question-answer pairs authored by Trivia enthusiasts. We use the Wikipedia setting of TriviaQA, which provides question-retrieved paragraphs from Wikipedia. We use the linear passage re-ranker in DocQA Clark and Gardner (2018) to re-rank these retrieved paragraphs and finally concatenate the first \(1{,}024\) tokens as a new reading document. We also employ the shared-normalization training objective Clark and Gardner (2018) for the vanilla BERT.

**WikiHop**Welbl et al. (2018) consists of \(51\)k questions, candidate answers, and supporting paragraphs triples. It requires models to find multi-hop reasoning chains for choosing the correct answer. Due to the long length of the concatenation of supporting paragraphs, we follow the processing tactic in Longformer Beltagy et al. (2020), which splits the document into several parts and then averages their candidate scores.

**YELP.F**Zhang et al. (2015) contains \(1,569\)k samples with review texts, which are obtained from the YELP Dataset Challenge in 2015. Yelp Review Full (YELP.F) contains five star classes.

**20NewsGroups**Lang (1995) comprises around \(18\)k newsgroups posts on \(20\) topics.

**IMDB**Maas et al. (2011) consists of \(50\)k informal movie reviews from the Internet Movie Database. Each review is annotated as _positive_ or _negative_.

**Hyperpartisan**Kiesel et al. (2019) aims to decide whether a news article text follows a hyperpartisan argumentation. Hyperparisan only contains \(645\) documents, which makes it a good testbed in a low-resource scenario. For the vanilla BERT, we adopt the max-pooling results of BERT sliding window. We split the data into training / validation / test set with a ratio of 9:1:1, run each model five times, and report the median performance.

## Appendix B Training Configuration

We follow the configuration from previous work Devlin et al. (2019); Beltagy et al. (2020); Liu et al. (2020) for fine-tuning BERT, DistilBERT and BERTL. Hyperparameters are shown in Table 5.

To train TR-BERT, we first initialize TR-BERT with corresponding fine-tuned models, which are trained with a task-specific objective for \(N\) epochs. After that, we maintain the same learning rate, warmup proportion and batch size for the latter two-step training: (1) Freeze all the parameters except that of the policy network and conduct reinforcement learning to update the policy network for \(\lceil(N+1)/2\rceil\) epochs; (2) Unfreeze all parameters and train the entire network with the task-specific knowledge distillation objective and the reinforcement learning objective simultaneously for \(N\) epochs.

## Appendix C Actual Wall Time

In practical applications, the wall time acceleration of TR-BERT is similar to the FLOPs acceleration. We evaluate our model on a single V100 GPU with \(32\) batch size on SQuAD. TR-BERT in Table 2 with \(2.08\)x FLOPs speedup achieves \(2.01\)x actual inference time speedup.