F-PABEE: Flexible-patience-based Early Exiting for Single-label and Multi-label text Classification Tasks

The static inference approach compresses the heavy model into a smaller one, including pruning, knowledge distillation, quantization, and weight sharing [14, 15, 16]. For example, HeadPrune [17] ranks the attention heads and prunes them to reduce inference latency. PKD [18] investigates the best practices of distilling knowledge from BERT into smaller-sized models. I-BERT [19] performs an end-to-end BERT inference without any floating point calculation. ALBERT [8] shares the cross-layer parameters. [20], [21] and [22] distills knowledge from the larger BERT teacher model for improving the performances of student networks which are learned with neural architecture search. Note that the static models are still in the form of deep neural networks with multiple stacked layers. The computational path is invariable for all examples in the inference process, which is not flexible.

Orthogonal to the static inference approach, early exiting dynamically adjusts hyper-parameters in response to changes in request traffic. It does not need to make significant changes to the original model structure or weight bits, nor does it need to train different teacher-student learning networks, which saves computing resources [23].

There are mainly three groups of dynamic early exiting strategies. The first type is confidence-based early exiting [24]. For example, BranchyNet [25], FastBERT [26], and DeeBERT [27] calculate the entropy of the prediction probability distribution to estimate the confidence of classifiers to enable dynamic early exiting. Shallow-deep [28] and RightTool [29] leverage the maximum of the predicted distribution as the exiting signal. The second type is the learned-based exiting, such as BERxiT [30] and CAT [31]. They learn a criterion for early exiting. The third type is patience-based early exiting, such as PABEE [13], which stops inference and exits early if the classifiers’ predictions remain unchanged for pre-defined times. Among them, patience-based PABEE achieves SOTA performance. However, PABEE suffers from too strict cross-layer comparison, and the applications on MLC tasks are neglected. There are also literature focusing on improving the training of multi-exit BERT, like LeeBERT [32] and GAML-BERT [33].

F-PABEE is a more flexible extension to PABEE, which can simultaneously adjust the confidence thresholds and patience parameters to meet different requirements. In addition, it outperforms other existing early exiting strategies on both SLC and MLC tasks.

The literature on early exiting focuses more on the design of early exiting strategies, thus neglect the advances of multi-exit backbones’ training methods. LeeBERT [32] employs an adaptive learning method for training multiple exits. GAML-BERT [33] enhances the training of multi-exit backbones by a mutual learning approach.

The inference procedure of F-PABEE is shown in Fig 1(b), which is an improved version of PABEE (Fig 1(a)), where $L_{i}$ is the transformer block of the model, $n$ is the number of transformer layers, $C_{i}$ is the inserted classifier layer, $s$ is the cross-layer similarity score, $thre$ is the similarity score threshold, $P_{0}$ is the pre-defined patience value in the model.

The input sentences are first embedded as the vector:

$$h_{0}=\text{Embedding}(x).$$

(1)

The vector is then passed through transformer layers ($L_{1}...L_{n}$) to extract features and compute its hidden state $h$. After which, we use internal classifiers ($C_{1}...C_{n}$), which are connected to each transformer layer to predict probability $p$:

$$p_{i}=C_{i}(h_{i})=C_{i}(L_{i}(h_{i-1})).$$

(2)

We denote the similarity score between the prediction results of layer $i-1$ and $i$ as $s(p_{i-1},p_{i})$ ($s(p_{i-1},p_{i})\in\mathbf{R}$). The smaller the value of $s(p_{i-1},p_{i})$, the prediction distributions are more consistent with each other. The premise of the model’s early exit is that the comparison scores between successive layers are relatively small; The similarity threshold $thre$ is a hyper-parameter. We use $pat_{i}$ to store the times that the cross-layer comparison scores are consecutively less than the threshold $thre$ when the model reaches current layer $i$:

$$pat_{i}=\left\{\begin{array}[]{rcl}pat_{i-1}+1&&s(p_{i-1},p_{i})<thre\\ 0&&s(p_{i-1},p_{i})>=thre\end{array}\right\}$$

(3)

If $s(p_{i-1},p_{i})$ is less than the similarity score threshold $thre$, then increase the patience counter by 1. Otherwise, reset the patience counter to 0. This process is repeated until $pat$ reaches the pre-defined patience value $P_{0}$. The model dynamically stops inference and exits early. However, if this condition is never met, the model uses the final classifier layer to make predictions. This way, the model can stop inference early without going through all layers.

Under the framework of F-PABEE, we can adopt different similarity measures for predicted probability distributions. This work uses the knowledge distillation objectives as the similarity measures [34]. When the model reaches the current layer $l$, for SLC tasks, we compare a series of similarity measures of F-PABEE, denoted as:

In addition, for MLC tasks, we transform them into multiple binary classification problems and sum the similarity scores of all categories, and the formulas are denoted as:

The formulations of F-PABEE-SymKD and F-PABEE-JSKD for MLC tasks are similar to those of SLC tasks.

F-PABEE is trained on SLC and MLC tasks, while the activation and loss functions are different. For SLC tasks, we use the softmax activation function and cross-entropy function according to the tasks. In contrast, we use the sigmoid activation function and binary cross-entropy function for MLC tasks.

After that, we optimize the model parameters by minimizing the overall loss function $L$, which is the weighted average of the loss terms from all classifiers:

$$L=\sum_{j=1}^{n}jL_{j}/\sum_{j=1}^{n}j$$

(10)

We evaluate F-PABEE on GLUE benchmark [35] for SLC tasks and four datasets for MLC tasks: MixSNLPS [36], MixATS [37], AAPD [38], and Stackoverflow [39]. we compare F-PABEE with three groups of baselines: (1) BERT-base; (2) Static exiting; (3) Dynamic exiting methods, including BrachcyNet [40], Shallow-Deep [28], BERxiT [30], and PABEE. Considering the flops of inferencing one with the whole BERT as the base, the speed-up ratio is defined as the average ratio of reduced flops due to early exiting.

In training process, we perform grid search over the batch size of {16, 32, 128}, and learning rate of {1e-5, 2e-5, 3e-5, 5e-5} with an AdamW optimizer [41] . The batch size in the inference process is 1. We implement F-PABEE on the bases of HuggingFace Transformers [42]. All experiments are conducted on two Nvidia TITAN X 24GB GPUs.

In Table 1, we compare F-PABEE with other early exiting strategies. We adjust the hyper-parameters of F-PABEE and other baselines to ensure similar speedups with PABEE. It shows that F-PABEE balances speedup and performance better than baselines, especially for a large speedup ratio. Moreover, we draw the score-speedup curves for BERxiT, PABEE, and F-PABEE. It shows that F-PABEE outperforms the baseline models on both SLC (Fig 2) and MLC tasks(Fig 3). Furthermore, the distribution of executed layers (Fig 4) indicates that F-PABEE can choose the faster off-ramp and achieve a better trade-off between accuracy and efficiency by flexibly adjusting similarity score thresholds and patience parameters.

Ablation on different PLMs  F-PABEE is flexible and can work well with other pre-trained models, such as ALBERT. Therefore, to show the acceleration ability of F-PABEE with different backbones, we compare F-PABEE to other early exiting strategies with ALBERT base as the backbone. The results in Fig 5 show that F-PABEE outperforms other early exiting strategies under different backbones by large margins on both SLC and MLC tasks, indicating that F-PABEE can accelerate the inference process for numerous PLMs.

Comparisons between different similarity measures  We consider F-PABEE with different similarity measures, denoted as F-PABEE-KD, F-PABEE-ReKD, F-PABEE-SymKD, and F-PABEE-JSKD, and the results are presented in Fig 6. F-PABEE-JSKD performs the best on both SLC and MLC tasks. We suppose that F-PABEE-JSKD is symmetric, and the similarity discrimination is more accurate than asymmetric measures. Therefore, it is better at determining which samples should exit at shallow layers and which should go through deeper layers.