Venue: NAACL 2021

# A Global Past-Future Early Exit Method for Accelerating Inference

of Pre-trained Language Models

 Kaiyuan Liao\({}^{\dagger}\), Yi Zhang\({}^{\ddagger}\), Xuancheng Ren\({}^{\sharp}\), Qi Su\({}^{\natural\$}\), Xu Sun\({}^{\dagger\dagger}\), Bin He\({}^{*}\)

\({}^{\dagger}\) Center for Data Science, Peking University

\({}^{\ddagger}\) MOE Key Laboratory of Computational Linguistics, School of EECS, Peking University

\({}^{\lx@sectionsign}\) School of Foreign Languages, Peking University

\({}^{\star}\) Huawei Noah's Ark Lab

{kyliao,zhangyi16,renxc,sukia,xusun}@pku.edu.cn hebin.nlp@huawei.com

Equal contribution

###### Abstract

Early exit mechanism aims to accelerate the inference speed of large-scale pre-trained language models. The essential idea is to exit early without passing through all the inference layers at the inference stage. To make accurate predictions for downstream tasks, the hierarchical linguistic information embedded in all layers should be jointly considered. However, much of the research up to now has been limited to use local representations of the exit layer. Such treatment inevitably loses information of the unused past layers as well as the high-level features embedded in future layers, leading to sub-optimal performance. To address this issue, we propose a novel Past-Future method to make comprehensive predictions from a _global_ perspective. We first take into consideration all the linguistic information embedded in the past layers and further engage the future information which is originally inaccessible for predictions. Extensive experiments demonstrate that our method outperforms previous early exit methods by a large margin, yielding better and robust performance1.

Footnote 1: The code is available at [https://github.com/lancopku/Early-Exit](https://github.com/lancopku/Early-Exit)

+
Footnote †: The code is available at [https://github.com/lancopku/Early-Exit](https://github.com/lancopku/Early-Exit)

## 1 Introduction

Pre-trained language models (PLMs), e.g., BERT Devlin et al. (2019), RoBERTa Liu et al. (2019) and XLNet Yang et al. (2019), have obtained remarkable success in a wide range of NLP tasks. Despite their impressive performance, PLMs are usually associated with large memory requirement and high computational cost. Such drawbacks slow down the inference and further encumber the application of PLMs in the scenarios where inference time and computation budget are restricted.

To address this issue, a growing number of studies focusing on improving model efficiency have emerged recently. Particularly, kaya2019conceptual point out that the current over-parameterized models conduct excessive computation for simple instances, which is actually undesirable and computationally wasteful. In light of this observation, an increasing amount of work seeks various early exit methods, of which the basic idea is to exit early without passing through the entire model during inference. Concretely, for NLP tasks, they couple branch classifiers with each layer of the pre-trained language models and stop forward propagation at an intermediate layer. Then the current branch classifier makes a prediction based on the representation of the token that is used as the aggregated sequence representation for classification tasks and is referred to as the state of the layer in this work.

However, existing work on early exit has two major drawbacks. First, existing work Xin et al. (2020); Zhou et al. (2020) uses only local states in the early exit framework. They inevitably lose valuable features that are captured by passed layers but are ignored for prediction, leading to less reliable prediction results. Moreover, these methods abandon the potentially useful features captured by the future layers that have not been passed, which may hurt the performance of the instances requiring high-level features embedded in the deep layers. Consequently, their performance dramatically declines when the inference exits earlier for a higher speed-up ratio.

These two major drawbacks hinder the progress of early exit research and motivate us to develop a new mechanism using the hierarchical linguistic information embedded in all layers Jawahar et al. (2019) from a global perspective. However, up to now, a global early exit mechanism remains a under-explored challenging problem. We extend the existing methods to their corresponding global versions and find that naive global strategies only result in poor performance. Meanwhile, the future states are originally inaccessible in the early exit

framework, which also remains a bottleneck for a global prediction considering both past and future states.

In this paper, we focus on the aforementioned problems and first put into practice a global Past-Future early exit mechanism. The term _global_ is two-fold: (1) instead of using one or several local state(s) for prediction in previous work, all the available past states are effectively incorporated in our method; (2) furthermore, to grasp the features embedded in the deep layers, the originally inaccessible future states are approximated by imitation learning and are also engaged for prediction. The comparison of the previous method and our method is illustrated in Figure 1. By combining both past and future states, our model is able to make more accurate predictions for downstream tasks.

Extensive experiments reveal that the proposal significantly outperforms previous early exit methods. Particularly, it surpasses the previous methods by a large margin when the speed-up ratio is relatively high. In addition, extensive experiments with different pre-trained language models as backbones demonstrate consistent improvement over the baseline methods, which verifies the generality of our method.

To summarize, our contributions are as follows:

* We propose a set of global strategies which effectively incorporate all available states and they achieve better performance compared to the existing naive global strategies.
* Our early exit method first utilizes the future states which are originally inaccessible at the inference stage, enabling more comprehensive global predictions.
* Experiments show that our proposal achieves better performance compared to the previous state-of-the-art early exit methods.

## 2 Related Work

Large-scale pre-trained language models Devlin et al. (2019); Liu et al. (2019) based on the Transformer Vaswani et al. (2017) architecture demonstrate superior performance in various NLP tasks. However, the impressive performance is on the basis of massive parameters, leading to large memory requirement and computational cost during inference. To overcome this bottleneck, increasing studies work on improving the efficiency of overparameterized pre-trained language models.

Knowledge distillation Hinton et al. (2015); Turc et al. (2019); Jiao et al. (2019); Li et al. (2020a) compacts the model architecture to obtain a smaller model that remains static for all instances at the inference stage. Sanh et al. (2019) focus on reducing the number of layers since their investigation reveals variations on hidden size dimension have a smaller impact on computation efficiency. Sun et al. (2019) learn from multiple intermediate layers of the teacher model for incremental knowledge extraction instead of only learning from the last hidden representations. Further, Wang et al. (2020) design elaborate techniques to drive the student model to mimic the self-attention module of teacher models. Xu et al. (2020) compress model by progressive module replacing, showing a new perspective of model compression. However, these static model compression methods treat the instances requiring different computational cost without distinction. Moreover, they have to distill a model from scratch to meet the varying speed-up ratio requirements.

Figure 1: Comparison of the local early exit method and our proposed method. The red rectangles highlight the states that the models rely on to make predictions.

 

To meet different constraints for acceleration, another line of work studies instance-adaptive methods to adjust the number of executed layers for different instances. Li et al. (2020) select models in different sizes depending on the difficulty of input instance. Besides, early exit is a practical method to adaptively accelerate inference and is first proposed for computer vision tasks Kaya et al. (2019); Teerapittayanon et al. (2016). Elbayad et al. (2020); Xin et al. (2020); Schwartz et al. (2020) follow the essential idea and leverage the method in NLP tasks. To prevent the error from one single classifier, Zhou et al. (2020) make the model stop inference when a cross-layer consistent prediction is achieved. However, researches on the subject has been mostly restricted to only use the local states around the exit layer.

## 3 Method

We first introduce the strategies to incorporate multiple states and the imitation learning method for generating approximations of future states. Then we introduce the merging gate to adaptively fuse past and future states. At last, we show the training process and the exit condition during inference.

### Incorporation of Past States

Existing work Xin et al. (2020) focuses on making exit decision based on a single branch classifier. The consequent unreliable result motivates the recent advance Zhou et al. (2020) that uses consecutive states to improve the accuracy and robustness. However, the model prediction is still limited to use several local states. In contrast, we investigate how to incorporate all the past states from a global perspective. The existing strategy using consecutive consistent prediction labels can be easily extended to a global version that counts the majority of the predicted labels which is regarded as a voting strategy. Another alternative is the commonly-used ensemble strategy that averages the output probabilities for prediction. Besides these naive solutions, we explore the following strategies to integrate multiple states into a single one:

* Max-Pooling: The max-pooling operation is performed on all available states, resulting in the integrated state.
* Avg-Pooling: The average-pooling operation is performed on all available states, resulting in the integrated state.
* Attn-Pooling: The attentive-pooling takes the weighted summation of all available states as the integrated state. The attention weights are computed with the last state as the query.
* Concatenation: All available states are concatenated and then fed into a linear transformation layer to obtain the compressed state.
* Sequential Neural Network: All available states are sequentially fed into an LSTM and the hidden output of the last time-step is regarded as the integrated state.

Formally, the state of the \(i\)-th layer is denoted as \({\mathbf{s}_{i}}\). When forward propagation proceeds to the \(i\)-th intermediate layer, all the past states \({\mathbf{s}_{1:i}}\) are incorporated into a global past state \({\mathbf{s}_{p}}\):

\[{\mathbf{s}_{p}}=G({\mathbf{s}_{1:i}}) \tag{1}\]

where \(G(\cdot)\) refers to one of the state incorporation strategies.

### Imitation of Future States

Existing work for early exit stops inference at an intermediate layer and ignores the underlying valuable features captured by the future layers. Such treatment is partly rationalized by the recent claim Kaya et al. (2019) that shallow layers are adequate to make a correct prediction. However, Jawahar et al. (2019) reveal that the pre-trained language models capture a hierarchy of linguistic information from the lower to the upper layers, e.g., the lower layers learn the surface or syntactic features while the upper layers capture high-level information like the semantic features. We hypothesize that some instances not only rely on syntactic features but also require semantic features. It is actually undesirable to only consider features captured by shallow layers. Therefore, we propose to take advantage of both past and future states.

Normally, we can directly fetch the past states, while using future information is intractable how since the future states are inaccessible before passing through the future layers. To bridge this gap, we propose a simple method to approximate the future states in light of imitation learning Ross et al. (2011); Nguyen (2016); Ho and Ermon (2016). We couple each layer with an imitation learner. During training, the imitation learner is encouraged to mimic the representation of the real state of that layer. Through this layer-wise imitation, we can 

obtain approximations of the future states with minimum cost. The illustration of the future imitation learning during inference is shown in Figure 2.

To be precise, we intend to obtain a state approximation of the \(j\)-th layer if the forward pass exits at the intermediate \(i\)-th layer for any \(j>i\). During training, we pass through the entire \(n\)-layer model but we simulate the situation that the forward pass ends up at the \(i\)-th layer for any \(i<n\). The \(j\)-th learner corresponding to the \(j\)-th layer takes \({\mathbf{s}_{i}}\) as input and outputs an approximation \(\hat{\mathbf{s}^{i}_{j}}\) of the real state \({\mathbf{s}_{j}}\). Then \({\mathbf{s}_{j}}\) serves as a teacher to guide the \(j\)-th imitation learner. We adopt cosine similarity as the distance measurement and penalize the discrepancy between the real state \({\mathbf{s}_{j}}\) and the learned state \(\hat{\mathbf{s}^{i}_{j}}\). Let \(\mathcal{L}^{i}_{\text{cos}}\) denotes the imitation loss of the situation that the forward pass exits at the \(i\)-th layer, it is computed as the average of the similarity loss for any \(j>i\). Since the exit layer \(i\) can be any number between 2 to \(n\) during inference, we go through all possible number \(i\) and average the corresponding \(\mathcal{L}^{i}_{\text{cos}}\), resulting the overall loss \(\mathcal{L}_{\text{cos}}\):

\[\hat{\mathbf{s}^{i}_{j}} =\text{Learner}^{j}({\mathbf{s}_{i}}) \tag{2}\] \[l^{i,j}_{\text{cos}}({\mathbf{s}_{j}},\hat{\mathbf{s}^{i}_{j}}) =1-\frac{\hat{\mathbf{s}^{i}_{j}}\cdot{\mathbf{s}_{j}}}{\|\hat{\mathbf{s}^{ i}_{j}}\||{\mathbf{s}_{j}}\|}\] (3) \[\mathcal{L}^{i}_{\text{cos}} =\frac{1}{n-i}\sum\nolimits_{j=i+1}^{n}l^{i,j}_{\text{cos}}({\bm {s}_{j}},\hat{\mathbf{s}^{i}_{j}})\] (4) \[\mathcal{L}_{\text{cos}} =\frac{1}{n-1}\sum\nolimits_{i=2}^{n}\mathcal{L}^{i}_{\text{cos}} \tag{5}\]

where \(\|\cdot\|\) denotes the \(L_{2}\) norm. Learner\({}^{j}(\cdot)\) is a simple feed-forward layer with learnable parameters \({\mathbf{W}^{i}}\) and \({\mathbf{b}^{i}}\).

During training, the forward propagation is computed on all layers and all imitation learners are encouraged to generate representations close to the real states. During inference, the forward propagation proceeds to the \(i\)-th intermediate layer and the subsequent imitation learners take the \(i\)-th real state as input to generate the approximations of future states. Then the approximations are incorporated into a comprehensive future state \({\mathbf{s}_{f}}\) with one of the global strategies introduced before:

\[{\mathbf{s}_{f}}=G(\hat{\mathbf{s}^{i}_{i+1:n}}) \tag{6}\]

where \(\hat{\mathbf{s}^{i}_{i+1:n}}\) denotes the approximations of the states from the (\(i\)+1)-th layer to the \(n\)-th layer.

### Adaptive Merging Gate

We then explore how to adaptively merge the past information and future information. Intuitively, the past state \({\mathbf{s}_{p}}\) and the future state \({\mathbf{s}_{f}}\) are of different importance since the authentic past states are more reliable than our imitated future states. In addition, different instances depend differently on high-level features learned by future layers. Therefore, it is indispensable to develop an adaptive method to automatically combine the past state \({\mathbf{s}_{p}}\) and the future state \({\mathbf{s}_{f}}\). In our work, we design an adaptive merging gate to automatically fuse the past state \({\mathbf{s}_{p}}\) and the future state \({\mathbf{s}_{f}}\). As the forward propagation proceeds to the \(i\)-th layer, we compute the reliability of the past state \({\mathbf{s}_{p}}\), and the final merged representation is a trade-off between these two states:

\[\alpha =\text{sigmoid}(\text{FFN}({\mathbf{s}_{p}})) \tag{7}\] \[{\mathbf{z}_{i}} =\alpha{\mathbf{s}_{p}}+(1-\alpha){\mathbf{s}_{f}} \tag{8}\]

where \({\mathbf{z}_{i}}\) is the merged final state and FFN\((\cdot)\) is a linear feed forward layer of the merging gate.

During training, each layer can generate the approximated states of future and obtain a merged final state which is used for prediction. Then the model will be updated with the layer-wise cross-entropy loss against the ground-truth label \(y\). The merging gate adaptively learns to adjust the balance under the supervision signal given by ground-truth labels. However, with the layer-wise optimization objectives, the shallow layers will be updated more frequently since they receive more updating signals from higher layers. To address this issue, we heuristically re-weight the cross entropy loss of each layer

Figure 2: The illustration of the future imitation learning. \(m\) is the exit layer and the dashed line denotes the unused modules at inference stage.

 

depending on its depth \(i\) and get its weight \(w_{i}\). The updating procedure is formalized as:

\[w_{i}=\frac{i}{\sum^{n}_{j=1}j} \tag{9}\] \[p_{i}=\text{softmax}(\mathbf{z}_{i})\] (10) \[\mathcal{L}^{i}_{\text{ce}}=-\sum_{l\in\text{labels}}y(l)\text{ log}(p_{i}(l))\] (11) \[\mathcal{L}_{\text{ce}}=\sum^{n}_{i=1}w_{i}\mathcal{L}^{i}_{\text {ce}} \tag{12}\]

The overall loss is computed as follows:

\[\mathcal{L}=\mathcal{L}_{\text{ce}}+\mathcal{L}_{\text{cos}} \tag{13}\]

### Fine-tuning and Inference

Here we introduce the fine-tuning technique and the exit condition at the inference stage.

Fine-tuningThe representations learned by shallow layers have a big impact on performance in the early exit framework since the prediction largely depends on the states of shallow layers. Most existing work updates all of the model layers at each step during fine-tuning to adapt to the data of downstream tasks. However, we argue that such an aggressive updating strategy may undermine the well-generalized features learned in the pretraining stage. In our work, we try to balance the requirements of maintaining features learned in pre-training and adapting to data at the fine-tuning stage. Specifically, the parameters of a layer will be frozen with a probability \(p\) and the probability \(p\) linearly decreases from the first layer to the \(L\)-th layer in a range of \(1\) to \(0\).

InferenceFollowing Xin et al. (2020), we quantify the prediction confidence \(e\) with the entropy of the output distribution \(p_{i}\) of \(i\)-th layer:

\[e(p_{i})=\text{Entropy}(p_{i}) \tag{14}\]

The inference stops once the confidence \(e(p_{i})\) is lower than a predefined threshold \(\tau\). The hyper-parameter \(\tau\) is adjusted according to the required speed-up ratios. If the exit condition is never reached, our model degrades into the common case of inference that the complete forward propagation is accomplished.

## 4 Experiments

### Experimental Setup

Experimental SettingsFollowing previous work (Xin et al., 2020), we evaluate our proposed method on six classification datasets from the GLUE benchmark (Wang et al., 2019): SST-2, MRPC, QNLI, RTE, QQP, and MNLI. We perform a grid search over the sets of learning rate as {1e-5, 2e-5, 3e-5, 5e-5}, batch size as {16, 32, 128} and number of frozen layers during fine-tuning as {0,1,2,3}. The maximum sequence length is fixed to 128. We employ a linear decay learning rate scheduler and the AdamW optimizer. In addition, we use the concatenation strategy to incorporate all available states for its best performance on the GLUE dev set.

Speed MeasurementSince the measurement of runtime might not be stable, following Xin et al. (2020); Zhou et al. (2020), we manually adjust the exit threshold \(\tau\) and calculate the speed-up ratio by comparing the actually executed layers in forward propagation and the required complete layers. For a \(n\)-layer model, the speed-up ratio is:

\[\text{speed-up ratio}=\frac{\sum^{n}_{i=1}n*m^{i}}{\sum^{n}_{i=1}i*m^{ i}} \tag{15}\]

where \(m^{i}\) is the number of examples that exit at the \(i\)-th layer of the model.

### Baselines

The proposed method can be practical for a range of existing pre-trained language models. Without losing generality, we conduct experiments with several well-known PLMs as backbones, namely, BERT, RoBERTa, and ALBERT Lan et al. (2019). Both BERT and RoBERTa suffer from the problem of over-parameterization. ALBERT largely alleviates this problem and is very efficient in terms of model size, the results on which verify the effectiveness on such parameter-efficient models. We mainly compare our method with other methods targeting on reducing the depth of models, including the recent early exit methods and the method directly reducing model depth to \(m\) layers which is denoted as (AL)BERT-\(m\)L.

### Overall Comparison

We compare our model performance with the baseline methods when different backbone models are adopted and show the result in Table 1 and Table 2. Both PABEE Zhou et al. (2020) and DeBERT Xin et al. (2020) accelerate inference with a highest \(2\times\) speed-up ratio. To be consistent, we adjust the exit threshold to obtain a \(2\times\) speed-up ratio and report the results in Table 1. As shown, our 

method maintains a comparable result with the original models on most datasets. We also notice that directly reducing layers performs well and serves as a strong baseline. Nevertheless, our proposal significantly outperforms such a method as well as the other two early exit methods.

We then adopt a more aggressive 3.00\(\times\) speed-up ratio to verify the effectiveness of our method. According to Table 2, the performance of PABEE and DeeBERT deteriorates badly. In contrast, our model exhibits more robust and stable performance, showing its superiority over previous early exit methods. Particularly, ALBERT is already very efficient in model size owing to its layer-sharing mechanism. Results shown in the bottom of Table 2 suggest that our model can obtain a good result with minimum performance loss on such a parameter-efficient model.

The success of our proposal might be attributed to the global perspective for prediction. DeeBERT makes prediction with the help of the state of a single branch classifier, leading to less reliable results. Although PABEE employs cross-layer prediction to prevent error from one single classifier, they ignore much available information of past states as well as the high-level semantic features captured by future layers. Different from those methods, our method jointly takes into consideration the hierarchical linguistic information embedded in all layers and thus is able to produce more accurate results.

### Performance-Efficiency Trade-Off

To further verify the robustness and efficiency of our method, we visualize the performance-efficiency trade-off curves in Figure 3 on a representative subset of the GLUE dev set. The backbone

\begin{table}
\begin{tabular}{l|c c c c|c c c|c c c|c c c|c} \hline \multirow{2}{*}{**Model**} & \multicolumn{2}{c|}{**MNLI-m**} & \multicolumn{2}{c|}{**MNLI-mm**} & \multicolumn{2}{c|}{**QQP**} & \multicolumn{2}{c|}{**QNLI**} & \multicolumn{2}{c|}{**SSTF2-**} & \multicolumn{2}{c|}{**MRPC**} & \multicolumn{2}{c|}{**RTE**} \\  & Acc & Spd-up & Acc & Spd-up & FI/Acc & Spd-up & Acc & Spd-up & Acc & Spd-up & FI/Acc & Spd-up & Acc & Spd-up \\ \hline BERTbase [18] & 84.6 & 1.00\(\times\) & 83.4 & 1.00\(\times\) & 71.2\(\cdot\) & 1.00\(\times\) & 90.5 & 1.00\(\times\) & 93.5 & 1.00\(\times\) & 88.9\(\cdot\) & - & 1.00\(\times\) & 66.4 & 1.00\(\times\) & - \\ \hline BERT-4L & 77.6 & 3.00\(\times\) & 77.2 & 3.00\(\times\) & 67.78\(\cdot\) & 3.00\(\times\) & 85.4 & 3.00\(\times\) & 88.7 & 3.00\(\times\) & 82.974\(\cdot\) & 9.30\(\times\) & **63.0**\(\times\) & 3.00\(\times\) & 78.4 \\ DeepBERT & 61.0 & 2.80\(\times\) & 59.8 & 2.84\(\times\) & 66.16\(\%\) & 3.9 & 80.8 & 2.88\(\times\) & 84.7 & 27.1\(\times\) & 83.575\(\cdot\) & 2.61\(\times\) & 60.5 & 2.90\(\times\) & 71.8 \\ PABEET & 75.9 & 2.70\(\times\) & 75.3 & 2.71\(\times\) & 69.58\(\pm\) & 2.57\(\times\) & 82.6 & 3.04\(\times\) & 85.2 & 3.04\(\times\) & 82.3 & 85.2 & 62.73\(\times\) & 2.72\(\times\) & 63.0 & 2.83\(\times\) & 76.6 \\ Ours & **78.4** & 2.99\(\times\) & **77.4** & 3.02\(\times\) & **70.49\(\times\)** & **3.16\(\times\)** & **87.3** & 2.78\(\uparrow\)** & **91.1** & 2.97\(\times\) & **84.577\(\cdot\)** & **2.87\(\times\)** & **63.0**\(\times\)** & 2.88\(\times\)** & **79.7** \\ \hline RoBERTa & & & & & & & & & & & & & & & \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 86.3 & 1.00\(\times\) & 71.8\(\cdot\) & - & 1.00\(\times\) & 92.4 & 1.00\(\times\) & 94.3 & 1.00\(\times\) & 90.4\(\cdot\) & - & 1.00\(\times\) & 67.5 & 1.00\(\times\) & - \\ \hline RoBERTa & & & & & & & & & & & & & & & & \\ \hline RoBERTa & & & & & & & & & & & & & & & & \\ \hline ALBERTbae [18] & 87.0 & 1.00\(\times\) & 86.3 & 1.00\(\times\) & 71.8\(\cdot\) & - & 1.00\(\times\) & 92.4 & 1.00\(\times\) & 93.3 & 1.00\(\times\) & 93.68\(\bullet\) & 1.00\(\times\) & 72.0 & 1.00\(\times\) & 84.8 \\ \hline ResNet & & & & & & & & & & & & & & & \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 86.3 & 1.00\(\times\) & 71.8\(\cdot\) & - & 1.00\(\times\) & 92.4 & 1.00\(\times\) & 93.3 & 1.00\(\times\) & 93.68\(\bullet\) & **1.00\(\times\)** & 72.0 & 1.00\(\times\) & 84.8 \\ \hline ResMatch & & & & & & & & & & & & & & & \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 86.3 & 1.00\(\times\) & 71.8\(\cdot\) & - & 1.00\(\times\) & 92.4 & 1.00\(\times\) & 93.3 & 1.00\(\times\) & 93.68\(\bullet\) & **1.00\(\times\)** & 72.0 & 1.00\(\times\) & 84.8 \\ \hline ResNet & & & & & & & & & & & & & & & \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 86.3 & 1.00\(\times\) & 71.8\(\cdot\) & - & 1.00\(\times\) & 92.4 & 1.00\(\times\) & 93.3 & 1.00\(\times\) & 90.84\(\bullet\) & 1.00\(\times\) & 67.5 & 1.00\(\times\) & - \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 83.4 & 2.00\(\times\) & 71.6\(\otimes\) & **72.03\(\times\)** & **70.49\(\times\)** & **2.05\(\times\)** & **86.7** & **2.03\(\times\)** & **87.0** & **2.6**\(\times\)** & **64.6** & 2.69\(\times\)** & 80.3 \\ \hline ResNet & & & & & & & & & & & & & & \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 86.3 & 1.00\(\times\) & 71.8\(\cdot\) & - & 1.00\(\times\) & 92.4 & 1.00\(\times\) & 93.3 & 1.00\(\times\) & 90.84\(\bullet\) & 1.00\(\times\) & 67.5 & 1.00\(\times\) & - \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 83.4 & 2.00\(\times\) & 71.6\(\otimes\) & **72.03\(\times\)** & **70.49\(\times\)** & **2.05\(\times\)** & **86.7** & **2.03\(\times\)** & **87.0** & **2.6**\(\times\)** & **64.6** & 2.69\(\times\)** & 80.3 \\ \hline ResNet & & & & & & & & & & & & & & \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 86.3 & 1.00\(\times\) & 71.8\(\cdot\) & - & 1.00\(\times\) & 92.0 & 1.00\(\times\) & 93.3 & 1.00\(\times\) & 93.68\(\bullet\) & **1.00\(\times\)** & 72.0 & 1.00\(\times\) & 84.8 \\ \hline ResNet & & & & & & & & & & & & & & \\ \hline RoBERTa-base [18] & 87.0 & 1.00\(\times\) & 86.3 

model is BERT. Please refer to the Appendix A for results of RoBERTa and ALBERT. As can be seen from Figure 3, the performance of previous state-of-the-art early exit methods drops dramatically when the speed-up ratio increases, which limits their practicality for higher acceleration requirements. By comparison, our method demonstrates more tolerance of speed-up ratio. It significantly improves performance compared to previous best-performing early exit models under the same speed-up ratio, especially in the case that the speed-up ratio is high, indicating that it can be applied in a wider range of acceleration scenarios.

### Analysis

#### 4.5.1 Effect of Global Strategies

The results of different global strategies on a representative subset of GLUE dev are shown in Table 3. The naive global strategies including voting and ensemble perform poorly, which demonstrates that existing global strategies can only achieve sub-optimal performance. In contrast, we design simple yet effective global strategies to incorporate past states which bring significant improvement compared to baselines. In addition, we empirically find that the concatenation strategy works best from an overall point of view. We assume that such a strategy allows interaction among different states, yielding better performance. In addition, the effect of the merging gate can be found in Appendix B.

#### 4.5.2 Analysis of Future Information

To assess whether and how future information contributes to the prediction, we first evaluate the Global Future version of our early exit method where all the approximations of futures states are incorporated through the concatenation strategy. Effect of future information is backed with the results shown in Table 4. We observe that the Global Future mechanism brings improvement on most datasets for both \(2\times\) speed-up ratio and \(3\times\) speed-up ratio, which confirms that the approximations of future states help enhance the model ability in prediction. Beyond that, the future states can be especially advantageous for the models with a higher speed-up ratio. Recall that approximations of future states complement the high-level semantic information and the exit at shallow layers loses more semantic information in comparison with the exit at deep layers. Therefore, the benefit of future information is more significant compared to the exit at shallow layers, which is validated by the larger improvement gap with a \(3\times\) speed-up ratio.

We also investigate the effect of future information on exit time. Figure 4 demonstrates the distribution of exit layers with and without future information. When future information is engaged, we observe that the proportion of exit at shallow layers increases. The observation conforms with our intuition: with the approximations of future states supplemented for prediction, the merged state at a shallow layer is able to make a confident and correct prediction. Thus the exit time is earlier compared to situations without future states, result

Figure 3: Performance and efficiency trade-off for early exit methods with BERT as backbone. Our method outperforms previous early exit methods by a large margin especially under high speed-up ratios.

 

ing in a higher speed-up ratio. To be more specific, for MRPC, the speed-up ratios with and without future states are \(1.69\) and \(1.99\), and are \(1.92\) and \(2.04\) for MNLI, respectively. Meanwhile, we observe a performance boost with future states involved. It confirms our assumption that the high-level semantic features embedded in future states help improve performance in early exit framework.

#### 4.5.3 Comparison with Distillation Methods

As an alternative method to accelerate inference, knowledge distillation also exhibits promising performance for NLP tasks. We provide comparison with typical knowledge distillation methods in Table 5. Existing model TinyBERT Jiao et al. (2019) exerts multiple elaborate strategies to achieve the state-of-the-art results, including the expensive general distillation process and a vast amount of augmented data for fine-tuning. We remove these two techniques to exclude the effect of extra training data. Under the same settings, we observe that our method outperforms the distillation methods with the same speed-up ratio.

In general, early exit and distillation methods improve inference efficiency from different perspectives. The distillation methods are more efficient in saving memory usage, but the downside is that such static methods suffer from high computation cost to adapt to different speed-up ratios. A new student model has to be trained from scratch if the speed-up requirement changes. By contrast, dynamic methods are more flexible to meet different acceleration requirements. Concretely, simple instances will be processed by passing through fewer layers and complex instances may require more layers. Moreover, the speed-up ratio can be easily adjusted depending on the acceleration requests. Nevertheless, early exit and distillation accelerate inference from different perspectives and these two kinds of techniques can be integrated to further compress the model size and accelerate the inference time.

## 5 Conclusions

We propose a novel Past-Future early exit method from a global perspective. Unlike previous work using only local states for prediction, our model employs all available past states for prediction and propose a novel approach to engage the future states which are originally inaccessible for prediction. Experiments illustrate that our method achieves significant improvement over baseline methods with different models as backbones, suggesting the superiority of our early exit method.

Figure 4: The distribution of exit layers with and without future states on the MNLI and MRPC tasks. The exit threshold for the same task is fixed. When future states are engaged for prediction, we observe an increase of exit at shallow (1-4) layers as well as a performance boost.

 



## Acknowledgements

We thank all the anonymous reviewers for their constructive comments. This work is partly supported by National Key R&D Program of China No. 2019YFC1521200 and Beijing Academy of Artificial Intelligence (BAAI). Xu Sun is the corresponding author.

## References

* J. Devlin, M. Chang, K. Lee, and K. Toutanova (2019)BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), Minneapolis, Minnesota, pp. 4171-4186. External Links: Link Cited by: SS3.
* M. Elbayad, I. Gu, E. Grave, and M. Auli (2020)Depth-adaptive transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, External Links: Link Cited by: SS1, SS2.
* G. 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, SS2.
* J. Ho and S. Ermon (2016)Generative adversarial imitation learning. In Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 4565-4573. External Links: Link, Document Cited by: SS1, SS2.
* G. Jawahar, B. Sagot, and D. Seddah (2019)What does BERT learn about the structure of language?. In Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pp. 3651-3657. External Links: Link, Document Cited by: SS1, SS2.
* 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, SS2.
* Y. Kaya, S. Hong, and T. Dumitras (2019)Shallow-deep networks: understanding and mitigating network overthinking. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, Vol. 97, pp. 3301-3310. External Links: Link, 1909.10351 Cited by: SS1, SS2.
* Z. Lan, M. Chen, S. Goodman, K. Gimpel, P. Sharma, and R. Soricut (2019)ALBERT: a lite BERT for self-supervised learning of language representations. CoRRabs/1909.11942. External Links: Link, 1909.11942 Cited by: SS1, SS2.
* J. Li, X. Liu, H. Zhao, R. Xu, M. Yang, and Y. Jin (2020a)BERT-EMD: many-to-many layer mapping for BERT compression with earth mover's distance. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pp. 3009-3018. External Links: Link, 2009.11942 Cited by: SS1, SS2.
* L. Li, Y. Lin, S. Ren, D. Chen, X. Ren, P. Li, J. Zhou, and X. Sun (2020b)Accelerating pre-trained language models via calibrated cascade. CoRRabs/2012.14682. External Links: Link, 2012.14682 Cited by: SS1, SS2.
* Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov (2019)Roberta: a robustly optimized BERT pretraining approach. CoRRabs/1907.11692. External Links: Link, 1907.11692 Cited by: SS1, SS2.
* K. Nguyen (2016)Imitation learning with recurrent neural networks. CoRRabs/1607.05241. External Links: Link, 1607.05241 Cited by: SS1, SS2.
* S. Ross, G. J. Gordon, and D. Bagnell (2011)A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, AISTATS 2011, Fort Lauderdale, USA, April 11-13, 2011, volume 15 of _JMLR Proceedings_, pages 627-635. External Links: Link, 1910.01108 Cited by: SS1, SS2.
* V. Sanh, L. Debut, J. Chaumond, and T. Wolf (2019)DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. CoRRabs/1910.01108. External Links: Link, 1910.01108 Cited by: SS1, SS2.
* R. Schwartz, G. Stanovsky, S. Swabha, W. Jayamdipta, J. Dodge, and N. A. Smith (2020)The right tool for the job: matching model and instance complexities. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pp. 6640-6651. External Links: Link, 2009.11942 Cited by: SS1, SS2.
* S. Sun, Y. Cheng, Z. Gan, and J. Liu (2019)Patient knowledge distillation for BERT model compression. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pp. 4322-4331. External Links: Link, 1910.01108 Cited by: SS1, SS2.
* S. Teerapittayanon, B. McDanel, and H. T. Kung (2016)Branchynet: fast inference via early exiting from deep neural networks. In 23rd International Conference on Pattern Recognition, ICPR 2016, Cancun, Mexico, December 4-8, 2016, pp. 2464-2469. External Links: Link, 1607.05241 Cited by: SS1, SS2.
* I. Turc, M. Chang, K. Lee, and K. Toutanova (2019)Well-Read Students Learn Better:

On the Importance of Pre-training Compact Models. _arXiv e-prints_, page arXiv:1908.08962.
* 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 _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA_, pages 5998-6008.
* Wang et al. (2019) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. 2019. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net.
* Wang et al. (2020) Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. 2020. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. _CoRR_, abs/2002.10957.
* 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 the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020_, pages 2246-2251. Association for Computational Linguistics.
* Xu et al. (2020) Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, and Ming Zhou. 2020. Bert-of-theseus: Compressing BERT by progressive module replacing. _CoRR_, abs/2002.02925.
* 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 _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada_, pages 5754-5764.
* Zhou et al. (2020) Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. 2020. BERT loses patience: Fast and robust inference with early exit. _CoRR_, abs/2006.04152.

## Appendix A More Performance-Efficiency Trade-Off Curves

Performance-efficiency curves with RoBERTa and ALBERT as backbones are shown in Figure 5 and Figure 6 respectively. Similar to the observation with BERT as backbone, the performance of DeeBERT and PABEE becomes progressively worse as the speed-up ratio increases. In contrast, our past-future early exit method shows more robust results.

## Appendix B Effect of Merging Gate

\begin{table}
\begin{tabular}{p{56.9pt}|c|c|c|c} \hline \multirow{2}{*}{**Method**} & **MNLI-nn** & **QNLI** & **SST-2** & **MRPC** \\  & Acc & Acc & FC1/Acc & **Macro** \\ \hline Ours & 83.30 & 90.46 & 92.89 & 88.44/83.08 & 88.10 \\ -merging gate & 83.15 & 90.61 & 92.43 & 86.86/80.64 & 87.49 \\ \hline \end{tabular}
\end{table}
Table 6: Ablation study of the merging gate. The speedup ration is approximately 2.00\(\times\) and the model implementation is based on BERT.

Figure 5: Performance-efficiency trade-off for early exit method DeeBERT with RoBERTa as backbone.

Figure 6: Performance-efficiency trade-off for early exit method PABEE with ALBERT as backbone.

 

We conduct ablation study to show the effect of the merging gate and report the result in Table 6. We can see that the performance drops when we remove the merging gate from our model, suggesting that the merging gate plays an important role in keeping the balance between past information and future information.