SlowGuess/NAACL_2021_OCR huggingface dataset:
Kaiyuan Liao†,* Yi Zhang‡,* Xuancheng Ren‡, Qi Su§§, Xu Sun†‡, Bin He* † Center for Data Science, Peking University
$\ddagger$  MOE Key Laboratory of Computational Linguistics, School of EECS, Peking University
§ School of Foreign Languages, Peking University
* Huawei Noah's Ark Lab
{kyliao,zhangyi16,renxc,sukia,xusun}@pktu.edu.cn
hebin.nlp@huawei.com
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 performance<sup>1</sup>.
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, Kaya et al. (2019) 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
The code is available at https://github.com/lancopku/Early-Exit
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2013-2023 June 6-11, 2021. ©2021 Association for Computational Linguistics
(a) DeeBERT with local early exit method
(b) The Global Past version of our early exit method
(c) The Global Past-Future version of our early exit method
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.
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.
- 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.
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.
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. (2020b) 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.
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.
3.1 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  $\pmb{s}_i$ . When forward propagation proceeds to the  $i$ -th intermediate layer, all the past states  $\pmb{s}_{1:i}$  are incorporated into a global past state  $\pmb{s}_p$ :
$$
\boldsymbol {s} _ {p} = G \left(\boldsymbol {s} _ {1: i}\right) \tag {1}
$$
where  $G(\cdot)$  refers to one of the state incorporation strategies.
3.2 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
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.
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  $\boldsymbol{s}_i$  as input and outputs an approximation  $\hat{\boldsymbol{s}}_j^i$  of the real state  $\boldsymbol{s}_j$ . Then  $\boldsymbol{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  $\boldsymbol{s}_j$  and the learned state  $\hat{\boldsymbol{s}}_j^i$ . Let  $\mathcal{L}_{\cos}^i$  denotes the imitation loss of the situation that the forward pass exists 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}_{\cos}^i$ , resulting the overall loss  $\mathcal{L}_{\cos}$ :
$$
\hat {\boldsymbol {s}} _ {j} ^ {i} = \operatorname {L e a r n e r} ^ {j} \left(\boldsymbol {s} _ {i}\right) \tag {2}
$$
$$
l _ {\cos} ^ {i, j} \left(\boldsymbol {s} _ {j}, \hat {\boldsymbol {s}} _ {j} ^ {i}\right) = 1 - \frac {\hat {\boldsymbol {s}} _ {j} ^ {i} \cdot \boldsymbol {s} _ {j}}{\| \hat {\boldsymbol {s}} _ {j} ^ {i} \| \| \boldsymbol {s} _ {j} \|} \tag {3}
$$
$$
\mathcal {L} _ {\cos} ^ {i} = \frac {1}{n - i} \sum_ {j = i + 1} ^ {n} l _ {\cos} ^ {i, j} \left(\boldsymbol {s} _ {j}, \hat {\boldsymbol {s}} _ {j} ^ {i}\right) \tag {4}
$$
$$
\mathcal {L} _ {\cos} = \frac {1}{n - 1} \sum_ {i = 2} ^ {n} \mathcal {L} _ {\cos} ^ {i} \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  $s_f$  with one of the global strategies introduced before:
$$
\boldsymbol {s} _ {f} = G \left(\hat {\boldsymbol {s}} _ {i + 1: n} ^ {i}\right) \tag {6}
$$
where  $\hat{s}_{i+1:n}^{i}$  denotes the approximations of the states from the  $(i+1)$ -th layer to the  $n$ -th layer.
3.3 Adaptive Merging Gate
We then explore how to adaptively merge the past information and future information. Intuitively, the past state  $s_p$  and the future state  $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  $s_p$  and the future state  $s_f$ . In our work, we design an adaptive merging gate to automatically fuse the past state  $s_p$  and the future state  $s_f$ . As the forward propagation proceeds to the  $i$ -th layer, we compute the reliability of the past state  $s_p$ , and the final merged representation is a trade-off between these two states:
$$
\alpha = \operatorname {s i g m o i d} \left(\operatorname {F F N} \left(s _ {p}\right)\right) \tag {7}
$$
$$
\boldsymbol {z} _ {i} = \alpha \boldsymbol {s} _ {p} + (1 - \alpha) \boldsymbol {s} _ {f} \tag {8}
$$
where  $z_{i}$  is the merged final state and  $\mathrm{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
depending on its depth  $i$  and get its weight  $w_{i}$ . The updating procedure is formalized as:
$$
w _ {i} = \frac {i}{\sum_ {j = 1} ^ {n} j} \tag {9}
$$
$$
p _ {i} = \operatorname {s o f t m a x} \left(\boldsymbol {z} _ {i}\right) \tag {10}
$$
$$
\mathcal {L} _ {\mathrm {c e}} ^ {i} = - \sum_ {l \in \text {l a b e l s}} y (l) \log \left(p _ {i} (l)\right) \tag {11}
$$
$$
\mathcal {L} _ {\mathrm {c e}} = \sum_ {i = 1} ^ {n} w _ {i} \mathcal {L} _ {\mathrm {c e}} ^ {i} \tag {12}
$$
The overall loss is computed as follows:
$$
\mathcal {L} = \mathcal {L} _ {\mathrm {c e}} + \mathcal {L} _ {\mathrm {c o s}} \tag {13}
$$
3.4 Fine-tuning and Inference
Here we introduce the fine-tuning technique and the exit condition at the inference stage.
Fine-tuning The 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.
Inference Following Xin et al. (2020), we quantify the prediction confidence  $e$  with the entropy of the output distribution  $p_i$  of  $i$ -th layer:
$$
e \left(p _ {i}\right) = \operatorname {E n t r o p y} \left(p _ {i}\right) \tag {14}
$$
The inference stops once the confidence  $e(p_i)$  is lower than a predefined threshold  $\tau$ . The hyperparameter  $\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.
Experimental Settings Following 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 \(\{1\mathrm{e} - 5,\) \(2\mathrm{e} - 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 Measurement Since 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 {s p e e d - u p r a t i o} = \frac {\sum_ {i = 1} ^ {n} n * m ^ {i}}{\sum_ {i = 1} ^ {n} i * m ^ {i}} \tag {15}
$$
where  $m^i$  is the number of examples that exit at the  $i$ -th layer of the model.
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- $mL$ .
4.3 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 DeeBERT (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
Model	MNLI-m	MNLI-mm	QQP	QNLI	SST-2	MRPC	RTE	Macro
Acc	Spd-up	Acc	Spd-up	F1/Acc	Spd-up	Acc	Spd-up	Acc	Spd-up	F1/Acc	Spd-up	Acc	Spd-up
BERT
BERT-base (Devlin et al., 2019)	84.6	1.00×	83.4	1.00×	71.2/ -	1.00×	90.5	1.00×	93.5	1.00×	88.9/ -	1.00×	66.4	1.00×	-
BERT-6L	80.8	2.00×	79.9	2.00×	69.7/88.3	2.00×	86.7	2.00×	91.0	2.00×	85.1/78.6	2.00×	63.9	2.00×	80.5
DecBERT (Xin et al., 2020)	-	-	-	69.4/ -	1.96×	87.9	1.79×	91.5	1.89×	85.2/ -	1.79×	-	-	-
DecBERT	74.4	1.87×	73.1	1.88×	70.4/88.8	2.13×	85.6	2.09×	90.2	2.00×	84.4/77.4	2.07×	64.3	1.95×	74.7
PABEE	79.8	2.07×	78.7	2.08×	70.4/88.6	2.09×	88.0	1.87×	89.3	1.95×	84.4/77.4	2.01×	64.0	1.81×	80.0
Ours	83.3	1.96×	82.7	1.96×	71.2/89.4	2.18×	89.8	1.97×	92.8	2.02×	87.0/81.8	1.98×	64.5	2.04×	82.5
RoBERTa
RoBERTa-base (Xin et al., 2020)	87.0	1.00×	86.3	1.00×	71.8/ -	1.00×	92.4	1.00×	94.3	1.00×	90.4/ -	1.00×	67.5	1.00×	-
RoBERTa-6L	84.4	2.00×	83.4	2.00×	71.6/89.2	2.00×	90.4	2.00×	93.5	2.00×	89.3/85.5	2.00×	58.0	2.00×	82.5
DecBERT	64.2	1.87×	64.7	1.87×	72.0/89.3	2.05×	83.8	2.01×	86.9	2.02×	88.7/84.3	1.86×	60.8	1.90×	75.4
Ours	86.6	1.92×	86.2	1.93×	72.0/89.3	2.54×	91.7	2.11×	94.5	1.98×	89.3/85.5	1.95×	58.0	2.11×	83.6
ALBERT
ALBERT-base	85.2	1.00×	84.7	1.00×	70.5/88.7	1.00×	92.0	1.00×	93.3	1.00×	89.0/84.8	1.00×	72.0	1.00×	84.8
ALBERT-6L	82.4	2.00×	81.7	2.00×	69.8/88.3	2.00×	90.0	2.00×	91.8	2.00×	87.0/82.4	2.00×	65.8	2.00×	82.2
PABEE	84.2	1.90×	83.5	1.81×	70.7/88.9	2.11×	90.9	1.98×	92.4	1.80×	87.6/82.6	1.91×	66.8	2.06×	83.2
Ours	84.8	1.94×	84.1	1.95×	70.4/88.6	2.35×	91.9	1.97×	92.8	2.13×	88.3/84.6	1.95×	72.0	1.93×	84.5