SlowGuess/NAACL_2022_OCR huggingface dataset:
Zhen Zhang $^{1,2}$ , Wei Zhu $^{3}$ , Jinfan Zhang $^{4}$ , Peng Wang $^{5}$ , Rize Jin $^{2}$ , Tae-Sun Chung $^{1*}$
$^{1}$  Dept. of Artificial Intelligence, Ajou University, Suwon 16499, South Korea
$^{2}$  School of Software, Tiangong University, Tianjin 300387, China
<sup>3</sup> Department of Computer Science, East China Normal University, Shanghai, China
<sup>4</sup> Harbin Engineering University, Harbin, China
<sup>5</sup> College of Computer Science, Northwest Normal University, Lanzhou, China
BERT and other pre-trained language models (PLMs) are ubiquitous in modern NLP. Even though PLMs are the state-of-the-art (SOTA) models for almost every NLP task (Qiu et al., 2020), the significant latency during inference prohibits wider industrial usage. In this work, we propose Patient and Confident Early Exiting BERT (PCEE-BERT), an off-the-shelf sample-dependent early exiting method that can work with different PLMs and can also work along with popular model compression methods. With a multi-exit BERT as the backbone model, PCEE-BERT will make the early exiting decision if enough numbers (patience parameter) of consecutive intermediate layers are confident about their predictions. The entropy value measures the confidence level of an intermediate layer's prediction. Experiments on the GLUE benchmark demonstrate that our method outperforms previous SOTA early exiting methods. Ablation studies show that: (a) our method performs consistently well on other PLMs, such as ALBERT and TinyBERT; (b) PCEE-BERT can achieve different speed-up ratios by adjusting the patience parameter and the confidence threshold. The code for PCEE-BERT can be found at https://github.com/michael-wzhu/PCEE-BERT.
Since BERT (Devlin et al., 2018), the pre-trained language models (PLMs) have become the default state-of-the-art (SOTA) models for natural language processing (NLP). The recent years have witnessed the rise of many PLMs, such as GPT (Radford et al., 2019), XLNet (Yang et al., 2019), ALBERT (Lan et al., 2020), and so forth. These BERT-style models achieved considerable improvements in many Natural Language Processing (NLP) tasks by pre-training on the unlabeled corpus and fine-tuning on labeled tasks, such as text classifi
cation, natural language inference (NLI) and sequence labeling. Despite their excellent performances, there are two issues for PLMs.
First, previous studies show that PLMs such as BERT suffer from the over-thinking problem. (Zhou et al., 2020; Zhu et al., 2021) shows that in the sentence classification task, BERT's last few layers may be too deep for some samples. For a sentence classification task, if we insert a classifier on a certain intermediate layer and drop the deeper layers, these intermediate layers may outperform the last layer.
Figure 1: This figure demonstrates the overthinking problem in BERT when it is applied to the sentence classification task, such as CoLA from the GLUE benchmark.
The second drawback of PLMs is their high latency. Sentence classification (CLS) tasks play a central role in many application scenarios, such as dialogue systems, document analysis, content recommendation, etc. However, these applications are time-sensitive. For example, if a task-oriented dialogue (TOD) system takes a lot of time to respond, users will no doubt stop using this system. User experience studies show that a response has to be made in between 0-100 ms. Thus, a CLS mod
*Corresponding author. Email: tschung@ajou.ac.kr.
Findings of the Association for Computational Linguistics: NAACL 2022, pages 327 - 338 July 10-15, 2022 ©2022 Association for Computational Linguistics
(a) PABEE
(b) PCEE-BERT
Figure 2: Comparison between PABEE (Zhou et al., 2020) and our PCEE-BERT, a novel early exiting method that combines the score-based early exiting with the patience-based early exiting.
ule should be efficient and accurate. In addition, a special feature of consumer queries is that there are times when the number of queries is extremely high. For example, during the flu season, online medical consultation will be used much more often than usual. Thus, it is important for deployed models to adjust their latency dynamically. During peak hours, it switches to a low-latency mode to deal with more queries. And in other hours, it makes the best of itself to provide accurate answers. So how can we make model inference dynamically? The answer is adaptive inference.
There exists a branch of literature focusing on making PLMs' inference more efficient via network pruning (Zhu and Gupta, 2018; Xu et al., 2020; Fan et al., 2020; Michel et al., 2019), knowledge distillation (Sun et al., 2019; Sanh et al., 2019; Jiao et al., 2020a), weight quantization (Zhang et al., 2020; Bai et al., 2020; Kim et al., 2021) and adaptive inference (Zhou et al., 2020; Xin et al., 2020; Liu et al., 2020). Adaptive inference has drawn much attention. The idea of adaptive inference is to deal with simple examples with only shallow layers of BERT and process more difficult queries with deeper layers, thus significantly speeding up the inference time on average while maintaining high accuracy. The speed-up ratio can be easily controlled with certain hyper-parameters to handle significant changes in query traffic without re-deploying the model services or maintaining a group of models.
Early exiting is one of the most important adaptive inference methods (Bolukbasi et al., 2017).
As depicted in Figure 2(b), it implements adaptive inference by installing an early exit, i.e., an intermediate prediction layer, at each layer of BERT and early exiting "easy" samples to speed up inference. At the training stage, all the exits are jointly optimized with BERT's parameters. At the inference stage, there are two different settings. First, in budgeted exiting mode, the model makes a prediction with a fixed exit for all queries. This mode deals with heavy traffic by assigning a shallower exit for prediction. The other one is dynamic exiting mode. That is, some strategies for early exiting are designed to decide whether to exit at each layer given the currently obtained predictions (from previous and current layers) (Teerapittayanon et al., 2016; Kaya et al., 2019; Xin et al., 2020; Zhou et al., 2020). In this mode, different samples can exit at different depths.
There are mainly three early exiting strategies for BERT dynamic exiting. The first one is score-based early exiting. BranchyNet (Teerapittayanon et al., 2016), FastBERT (Liu et al., 2020), and DeeBERT (Xin et al., 2020) calculated the entropy of the prediction probability distribution as an estimation for the confidence of exiting classifiers to enable dynamic early exiting. Shallow-Deep Nets (Kaya et al., 2019) and RightTool (Schwartz et al., 2020a) leveraged the maximum of the predicted distribution as the exiting signal. The second type is the learned exiting (Elbayad et al., 2020). In this type of work, an early exiting signal is generated by a learnable module in the neural network. For example, BERxiT (Xin et al., 2021) install a fully
connected layer right after each transformer block of BERT to output a score that is used to decide whether the BERT should stop inference and exit early. The third type is patience-based early exiting, which relies on cross-layer comparison to formulate the exiting signal. PABEE (Zhou et al., 2020) propose a dynamic exiting strategy analogous to early stopping model training. That is, if the exits' predictions remain unchanged for a pre-defined number of times (patience), the model will stop inference and exit early. PABEE achieves SOTAs results for BERT early exiting.
Despite its state-of-the-art performances during early exiting, PABEE is inflexible in adjusting the speedup ratios. On a given task, once the multi-exit BERT is fine-tuned and the patience parameter is fixed, PABEE can only achieve a fixed average speedup ratio. Thus, PABEE can not achieve speedup ratios of certain values. This drawback makes PABEE inconvenient to use in real industrial scenarios. Thus, it is of great importance to come up with a method that can flexibly adjust its speedup ratios and performs comparable to or better than PABEE.
In this work, we propose Patiently Confidently Early Exiting BERT (PCEE-BERT), a novel early exiting method that combines the advantage of score-based methods and the patience based early exiting method. A multi-exit BERT is adopted as the backbone model, and an intermediate classifier (i.e., an exit) is installed right after each transformer block. PCEE-BERT will early exit if there are enough numbers (i.e., the patience parameter) of consecutive exits being confident for their predicted distributions. We mainly use entropy as the confidence measure. Intuitively, our method requires patience and confidence. It will not rush into an early exiting if we only see a couple of intermediate layers being confident. In addition, it allows the next layer to modify the predictions. In this way, our PCEE-BERT can exit with higher accuracy while maintaining flexibility.
Extensive experiments are conducted on the GLUE benchmark (Wang et al., 2018). The results show that our method outperforms the previous SOTA early exiting methods, especially in cases where the speedup ratio is large. In addition, one can adjust the patience and confidence threshold so that PCEE-BERT can arrive at different speedup ratios. A series of ablation studies are conducted, resulting in the following observations: (a)
PCEE-BERT can work with different confidence measures; (b) our method performs consistently well on different PLMs, and can work alongside model compression methods to further speed up the BERT's inference; (c) our PCEE-BERT can also be applied to computer vision tasks.
The rest of the paper is organized as follows. First, we introduce the preliminaries for multi-exit BERT and early exiting. Second, we elaborate on our PCEE-BERT method. Third, we conduct experiments on the GLUE benchmark and conduct a series of ablation studies. Finally, we conclude with possible future works.
In this section, we introduce the necessary background for BERT early exiting. Throughout this work, we consider the case of multi-class classification with samples  $\{(x,y), x \in \mathcal{X}, y \in \mathcal{Y}, i = 1,2,\dots,N\}$ , e.g., sentences, and the number of classes is  $K$ .
2.1 Backbone models
In this work, we adopt BERT as the backbone model. BERT is a multi-layer Transformer (Vaswani et al., 2017) network, which is pre-trained in a self-supervised manner on a large corpus. The number of transformer layers of our backbone is denoted as  $M$ , and the hidden dimension is  $d$ .
2.2 Early-exiting Architecture
As depicted in Figure 2, early exiting architectures are networks with exits at each transformer layer. With  $M$  exits, M classifiers  $f^{(m)}(x;\theta^{(m)}):$ $\mathcal{X}\to \Delta^K (m = 1,2,\dots,M)$  are designated at  $M$  layers of BERT, each of which maps its input to  $p^{(m)}(x;\theta^{(m)})$ , a probability distribution over the K classes. All the parameters of the transformer layers and exits are denoted as  $\Theta$ .
2.2.1 Training
At the training stage, all the exits are jointly optimized with a summed loss function. Following Huang et al. (2017) and Zhou et al. (2020), the loss function is the weighted average of the cross-entropy (CE) losses given by
$$
\mathcal {L} = \frac {\sum_ {m = 1} ^ {M} m * \mathcal {L} ^ {(m)}}{\sum_ {m = 1} ^ {M} m}, \tag {1}
$$
where  $\mathcal{L}^{(m)} = \mathcal{C}\mathcal{E}(y,p^{(m)}(x;\theta^{(m)}))$  denotes the cross-entropy loss of the  $m$ -th exit. Note that the
weight  $m$  corresponds to the relative inference cost of exit  $m$ .
2.2.2 Inference
During inference, the multi-exit BERT can exit early in two different modes, depending on whether the computational budget to classify an example is known or not.
Budgeted Exiting. If the computational budget is known, we can directly appoint a suitable exit  $m^{*}$  of BERT,  $f^{(m^{*})}(x; \theta^{(m^{*})})$ , to predict all queries.
Dynamic Exiting. Under this mode, after receiving a query input  $x$ , the model starts to predict on the classifiers  $f^{(1)}(x; \theta^{(1)})$ ,  $f^{(2)}(x; \theta^{(2)})$ , ..., in turn in a forward pass, reusing computation where possible. It will continue to do so until it receives a signal to stop early at an exit  $m^* < M$ , or arrives at the last exit  $M$ . At this point, it will output the final predictions based on the current and previous predictions. Note that under this early exit setting, different samples might exit at different layers.
3 PCEE-BERT
PABEE achieves the SOTA performances for BERT early exiting by applying an early exiting decision-making process that mimics the early stopping of model training. However, one drawback of PABEE is that it can not flexibly adjust the average inference layers (i.e., speed-ups) for a given dataset once its patience parameter is set. Table 1 shows PABEE can not achieve certain values for average inference layers, such as around 4.0, 6.0, or 9.0 on RTE. This drawback may limit the industrial usage of early exiting techniques. Thus, it is of great importance to develop a new method that performs comparably with PABEE and is more flexible than PABEE.
3.2 PCEE-BERT: a novel dynamic exiting method
The inference process of PCEE-BERT is illustrated in Figure 2(b). Assume the feed forward process for predicting sample  $x$  has gone through layers 1, ...,  $m - 1$ , and we are now at layer  $m$ . After going through the transformer layer  $m$ , the intermediate classifier  $f^{(m)}(x; \theta^{(m)})$  predicts a class label distribution  $p^{(m)}(x; \theta^{(m)})$ . The confidence level of layer  $m$  is measured by the entropy value of distribution  $p^{(m)}(x; \theta^{(m)})$ :
$$
C ^ {(m)} = \frac {\sum_ {k = 1} ^ {K} p _ {k} ^ {(m)} \log p _ {k} ^ {(m)}}{\log (1 / K)}, \tag {2}
$$
RTE	QNLI	MRPC
patience=1	3.24	2.25	2.00
patience=2	4.96	3.87	3.00
patience=3	6.69	5.32	4.18
patience=4	7.77	6.50	5.60
patience=5	8.78	7.61	6.81
patience=6	9.75	8.64	7.91
patience=7	10.68	9.54	8.83
patience=8	11.47	10.36	9.72
patience=9	11.79	11.04	10.51
patience=10	11.92	11.57	11.26
patience=11	12.00	12.00	12.00
Table 1: Average inference layers of PABEE on the RTE, QNLI and MRPC tasks.
where  $p_k^{(m)}$  is the probability mass for  $k$ -th class label. If  $C^{(m)}$  is smaller than a pre-defined threshold  $\tau$ , the predictions of layer  $m$  is considered confident. Otherwise, it is considered in-confident.
We use a patience counter  $pct$  to store the number of times that the predictions remain confident in consecutive layers. Formally, at layer  $m$ ,  $pct^{(m)}$  is calculated as
$$
p c t ^ {(m)} = \left\{ \begin{array}{l} p c t ^ {(m - 1)} + 1, \text {i f} C ^ {(m)} <   \tau , \\ 0, \text {o t h e r w i s e .} \end{array} \right. \tag {3}
$$
We stop inference early at layer  $m$  when  $pct^{(m)}$  reaches a predefined integer number  $t$  (the patience parameter). If this condition is never fulfilled, we use the final classifier  $M$  for prediction. In this way, the model can make an early exit without passing through all layers to make a prediction.
Our method draws advantages from the previous score-based early exiting method (Teerapittayanon et al., 2016) and patience-based method (Zhou et al., 2020) and overcomes their shortcomings. First, the score-based early exiting method relies on the confidence score from only the current layer. However, as revealed by Szegedy et al. (2014); Jiang et al. (2018), prediction of probability distributions (i.e., softmax scores) suffers from being over-confident to one class, making it an unreliable metric to represent confidence. In our method, early exiting occurs when a group of consecutive layers is confident, thus making the early exiting decision more reliable. Second, with a patience-based early exiting method like PABEE, when a deeper layer tries to correct the predictions, the patience count resets to zero. As a result, PABEE is less efficient than our PCEE-BERT. Third, since our method is a combination of PABEE and the
score-based method, one can conveniently adjust the threshold and patience parameters to control the speed-up ratios, which makes our method more flexible than PABEE.
We evaluate our proposed approach to the classification tasks on the GLUE benchmark. We only exclude the STS-B task since it is a regression task, and we exclude the WNLI task following previous work (Devlin et al., 2018; Xu et al., 2020).
We compare our approaches with three groups of baselines.
Backbone models: We mainly choose the BERT-base model open-sourced by Devlin et al. (2019) as the backbone model. We also investigate whether our method is applicable across different backbones, so we also run ablation experiments with ALBERT base (Lan et al., 2020) and TinyBERT $_6$  (Jiao et al., 2020b).
Budgeted exiting: In the section 2.2 we have introduced how to train a multi-exit BERT. Once the multi-exit BERT, we can conduct budgeted early exiting, that is, asking a designated intermediate layer to encode and predict all the samples. Budgeted exiting is a direct way to speed up BERT's inference, but it is not instance adaptive. Some of the samples may not need to go through many of the BERT's layers, and the others may be more difficult and require deeper feature encoding from the deeper layers of BERT.
Dynamic exiting: In this part, we compare our methods with a series of strong baselines, including BranchyNet (Teerapittayanon et al., 2016), Shallow-Deep (Kaya et al., 2019), BERxiT (Xin et al., 2021), and PABEE (Zhou et al., 2020). Note that PABEE can not flexibly adjust the average inference layers on a task once the patience parameter is set. So we will adjust the thresholds in the other baselines and our PCEE-BERT so that all methods' number of average inference layers are close.
4.3 Evaluation of early exiting method
In this work, we strictly follow the GLUE benchmark to report the performances metrics on each task. Note that this work focuses on investigating the early exiting of PLMs. Thus we have to consider the trade-offs between performance and
efficiency. Following PABEE (Zhou et al., 2020), we mainly report the speedup ratio as the efficiency metric. Assume the PLM backbone has  $N$  layers in total. For each test sample  $x_{i}$ $(i\in \{0,1,\dots,N\})$ , the early exiting layer is  $m_{i}$ , then the average speedup ratio on the test set is calculated by
$$
\mathbf {S p e e d u p} = 1 - \frac {\sum_ {1} ^ {N} m _ {i}}{\sum_ {1} ^ {N} M}. \tag {4}
$$
We choose this efficiency metric for the following reason: (1) it is linear w.r.t. the actual amount of computation; (2) according to our experiments, it is proportional to actual wall-clock runtime and is also more stable across different runs compared with actual runtime due to randomness by other processes on the same machine.
4.4 Experimental settings
Training We add a linear output layer after each intermediate layer of the pre-trained BERT or other backbone models as the internal classifiers. We perform grid search over batch sizes of 16, 32, 128, and learning rates of 1e-5, 2e-5, 3e-5, 5e-5 with an Adam optimizer. The hyper-parameters are selected via the 5-fold cross validation on the train set of GLUE tasks. We implement PCEE-BERT on the base of Hugging Face's Transformers (Wolf et al., 2020). Experiments are conducted on a single Nvidia V100 16GB GPU.
Inference Following prior work on input-adaptive inference (Teerapittayanon et al., 2016; Kaya et al., 2019), inference is on a per-instance basis, i.e., the batch size for inference is set to 1. This is a common scenario in the industry where individual requests from different users (Schwartz et al., 2020b) come at different time points. We report the median performance over five runs with different random seeds.
4.5 Main results
In Table 2, we report the performance comparisons of each method on the GLUE benchmark under three different speedup settings. The three speedup settings are: (1)  $74\%$  to  $82\%$  speedup; (2)  $46\%$  to  $54\%$  speedup; (3)  $23\%$  to  $28\%$  speedup. Since PABEE can not flexibly adjust the speedup ratios for a given patience parameter and a given task, we adjust the hyper-parameters (such as entropy threshold) of our PCEE-BERT and the other baselines to achieve similar speedups with PABEE. The results in table 2 clearly show that our PCEE-BERT
(a) MRPC task
(b) QNLI task
(c) MNLI task
(d) QQP task
Figure 3: Performance-efficiency trade-offs using different exiting strategies. We can see that our PCEE-BERT consistently outperforms the strong baseline, PABEE, especially when the speed-up ratio is large.
CoLA	MNLI	MRPC	QNLI	QQP	RTE	SST-2
score	speedup	score	speedup	score	speedup	score	speedup	score	speedup	score	speedup	score	speedup
BERT base	0.542	0%	0.831	0%	0.868	0%	0.898	0%	0.892	0%	0.691	0%	0.913	0%
Budgeted-Exit-3L	0.0	75%	0.700	75%	0.758	75%	0.774	75%	0.818	75%	0.547	75%	0.810	75%
Budgeted-Exit-6L	0.0	50%	0.796	50%	0.847	50%	0.853	50%	0.893	50%	0.681	50%	0.886	50%
Budgeted-Exit-9L	0.519	25%	0.830	25%	0.870	25%	0.884	25%	0.903	25%	0.690	25%	0.912	25%
BranchyNet	0.0,	74%	0.638	76%	0.757	76%	0.742	80%	0.716	80%	0.547	76%	0.799	76%
0.0	51%	0.783	53%	0.830	52%	0.871	47%	0.893	50%	0.674	47%	0.883	49%
0.521	27%	0.830	25%	0.858	24%	0.893	27%	0.901	26%	0.680	26%	0.912	24%
Shallow-Deep	0.0	75%	0.641	77%	0.756	76%	0.743	78%	0.714	79%	0.547	76%	0.795	77%
0.0	52%	0.782	51%	0.828	51%	0.872	49%	0.896	51%	0.672	48%	0.884	48%
0.523	26%	0.829	26%	0.857	25%	0.893	26%	0.901	27%	0.678	26%	0.912	25%
BERxiT	0.0	76%	0.635	76%	0.756	76%	0.733	78%	0.682	80%	0.553	77%	0.795	76%
0.1232	52%	0.784	51%	0.829	51%	0.870	48%	0.891	49%	0.673	47%	0.883	49%
0.522	25%	0.832	26%	0.862	26%	0.896	27%	0.901	26%	0.681	27%	0.914	24%
PABEE	0.0	75%	0.639	77%	0.758	75%	0.736	81%	0.686	82%	0.558	75%	0.799	77%
0.0	50%	0.789	52%	0.831	53%	0.872	46%	0.896	49%	0.677	46%	0.887	48%
0.524	26%	0.834	24%	0.861	26%	0.898	28%	0.904	24%	0.683	28%	0.917	22%
PCEE-BERT	0.098	79%	0.734	72%	0.788	77%	0.804	75%	0.796	82%	0.584	76%	0.836	76%
0.232	57%	0.801	53%	0.848	53%	0.871	54%	0.908	49%	0.694	47%	0.904	48%
0.528	27%	0.834	28%	0.868	26%	0.905	27%	0.912	25%	0.697	30%	0.918	23%
Table 2: Experimental results of different early exiting methods with the same fine-tuned BERT backbone on the GLUE benchmark. The results show that PCEE-BERT is effective in accelerating BERT's inference with less performance loss compared with the baseline methods.
method outperforms the baseline methods under different speedup ratios. Table 2 also shows that
the PABEE method is the best performing baseline. Thus, in order to further analyze and better visualize the results, we draw the score-speedup curves (in Figure 3) for budgeted early exiting, PABEE and PCEE-BERT, on the QNLI and MRPC tasks. With Table 2 and Figure 3, we can make the following observations:
- Although it is clear that PABEE performs better than the other baselines when the speedup ratio is around  $50\%$  or  $25\%$ , its advantages over the other baselines with the  $75\%$  speedup ratio is relatively small. With the  $75\%$  speedup ratio for seven GLUE tasks, it performs better than the score-based methods only on three tasks. This observation motivates us to improve PABEE by combining its patience-based early exiting mechanism with the score-based ones.
- Our PCEE-BERT consistently performs better than the baseline methods, especially when the speedup ratio is large. Note that our PCEE-BERT also consistently outperforms the budgeted exiting speedup ratios, which the other baselines do not achieve. Figure 3(b) and 3(a) show that score-speedup curve for PABEE is interleaving with that of the budgeted exiting. However, the score-speedup curve for PCEE-BERT distances itself from the others for most of the GLUE tasks.
- The overthinking problem is prevailing in the GLUE benchmark, and our PCEE-BERT early exiting can effectively take advantage of this phenomenon. For 6 of the GLUE tasks, PCEE-BERT can outperform BERT-base with a  $25\%$  (or more than) speedup ratio. And for 2 of the GLUE tasks, PCEE-BERT can outperform BERT-base with a  $50\%$  (or more than) speedup ratio.