SmartBERT: A Promotion of Dynamic Early Exiting Mechanism for Accelerating BERT Inference

1 Introduction

In recent years, large-scale pre-trained language models(PLMs) such as BERT [Devlin et al., 2019], GPT [Radford et al., 2018], ALBERT [Lan et al., 2019], and XLNET [Yang et al., 2019] RoBERTa [Liu et al., 2019] have made significant progress in the field of natural language processing(NLP). However, these models usually require large computation resources. It is difficult to deploy these models in the case of limited resources. To reduce redundant computation, various approaches have been proposed, including network pruning [Li et al., 2016; He et al., 2017], weight quantization [Jacob et al., 2018], knowledge distillation [Hinton et al., 2015; Sanh et al., 2019; Jiao et al., 2019] and dynamic early exiting [Xin et al., 2020; Liu et al., 2020; Schwartz et al., 2020].

In this work, we mainly focus on dynamic early exiting methods. Such methods do not change the original network structure, and they only add some light plugins in the original network to decide whether early exit or not in each layer which can reduce plenty of computation while keeping comparable performance.

Although current dynamic early exiting technology has shown excellent characteristics, samples must go through all consecutive layers before early exiting, and complex samples have to go through nearly all layers. That is, complex samples have more redundant computation. In the field of computer vision, Highway Network [Srivastava et al., 2015] and SkipNet [Wang et al., 2018b] have proved that redundant computations exist in neural networks and some blocks (layers) can be skipped directly. Motivated by these works, we assume that some layers also can be skipped directly before early exiting for BERT to further reduce computation.

In this paper, firstly, we introduce a novel dynamic early exiting combined with layer skipping technique for BERT inference named SmartBERT, which plugs a skipping gate and an exiting operator into each layer of BERT and performs adaptive inference based on the principle of the higher priority of skipping than exiting. We use the widely used PLM BERT [Devlin et al., 2019] as the backbone, and our method can be extended to other PLMs.

Secondly, in order to address the inconsistent usage of skipping gates between the train and inference stage, we design a hard weight mechanism. For adopting the transformation from soft to hard, we proffer a special training way. Thirdly, we propose cross-layer contrastive learning and combine it into our training phases to boost the intermediate layers and classifiers to achieve more efficient computation. Moreover, We conducted experiments on eight classification datasets of the GLUE benchmark[Wang et al., 2018a], which shows that SmartBERT achieves 2-3× computation reduction with minimal accuracy drops compared with BERT and has a better performance than other dynamic early exiting models at the same computational cost. We also prove skipping mechanism is highly effective in some complex datasets like RTE and WNLI.

Our main contributions can be summarized as follows:

• We propose a novel dynamic early exiting method combined with layer skipping for BERT inference. And we design a hard weight mechanism and a special training way for consistent usage and better training respectively.
• We propose cross-layer contrastive learning into the training phase to boost the intermediate layers and classifiers, which is beneficial for early exiting.
• Our method can achieve 2-3x computation reduction with comparable performance and outperforms previous methods in both efficiency and accuracy. In some complex datasets like RTE and WNLI, our method can save the quantity of computation time (2-4x) compared to current early existing methods while achieving better performance.

2 RelatedWork

2.1 Model Compression

Many pre-trained language models (PLMs) have emerged in the last few years. Although these PLMs have achieved high scores in many NLP tasks, their inference time is slow, and the cost of calculation is expensive. One of the most representative models is BERT [Devlin et al., 2019], which has made remarkable improvements in many NLP tasks. However, the inference speed of BERT is criticized. Therefore, a series of methods of model compression have been proposed to solve the above problems. Knowledge distillation [Hinton et al., 2015] aims to transfer knowledge from the teacher model to the student model, which is applied in DistillBERT [Sanh et al., 2019] and TinyBERT [Jiao et al., 2019]. ALBERT [Lan et al., 2019] uses sharing parameters to greatly reduce the number of parameters and memory consumption. Q8BERT [Zafrir et al., 2019] use symmetric linear quantization [Jacob et al., 2018] to reduce the number of bits about the parameters of BERT. For Pruning, [Gordon et al., 2020] mainly remove the unimportant part based on gradients of weights. Although these methods can improve the inference time of BERT, they can not adaptively change the architecture of the model according to the complexity of each sample. For example, simple samples must go through all layers for these methods. However, these samples may only go through early layers in early exiting models.

2.2 Adaptive Inference

Adaptive inference can adaptively change the architecture of the model according to the complexity of samples. DeeBERT [Xin et al., 2020] chooses whether to early exit according to the entropy of the output distribution on each layer. FastBERT divides the early exiting classifiers into student classifiers and the teacher classifier, and then trains all student classifiers through self-distillation. FastBERT also uses the early exit based on entropy. FastBERT and DeeBERT can adaptively adjust the model size according to different samples. However, samples need to go through all layers before early exiting, and if they are complex enough, they need to go through the whole model. Hence, we introduce cross-layer contrastive learning to obtain more powerful classifiers and use layer skipping mechanism to reduce redundancy further.

2.3 Contrastive Learning

Recently, contrastive learning (CL) has made significant progress in various domains. Especially, unsupervised CL can exploit a bulk of unlabelled data to train a model with generalization which can even surpass the model trained under supervised training in some situations. The aim of CL is to maximize the agreement between positive views which are jointly sampled and disparting negative views in the representation space. Pioneering works were mainly proposed in the CV domain [He et al., 2020; Chen et al., 2020]. Then in the graph domain, there are plentiful follow-up works [You et al., 2020; Zhu et al., 2022] which borrow ideas from pioneering works. However, in the field of natural language processing, there are few inventive works [Gao et al., 2021] that have shown up in recent years. The main difficulty in NLP domain is that augmentation is hard to design. Besides, previous works mainly focus on the quality of the final representations which may neglect the intermediate layers. But in our work, we hope each layer can be competent for good classification to early exit. So, we propose cross-layer contrastive learning in our method for both the first stage and the second training stage. In the cross-layer contrastive learning, we do not rely on data augmentation and assume that each token should have similar semantics across the consecutive layers because their representations do not change drastically. The same tokens will be treated as positive pairs in the consecutive layers, otherwise they shall be negative pairs.

3 Methodology

In this section, we first introduce the architecture of SmartBERT, which combines the early exiting mechanism and layer skipping mechanism. Secondly, we will illustrate how to combine cross-layer contrastive learning into the training phase to obtain powerful classifiers. Lastly, we introduce a hard weight mechanism and a special training way for skipping gates, further keeping the consistent usage of skipping gates between the training and inference phases.

3.1 Preliminaries

Given an input sentence S = {w1,w2, . . . ,wN} ∈ RN×W , N and L stand for the number of tokens and encoder layers respectively. Xi ∈ RN×D represents the output of the ith encoder layer E i (·). X0 stands for the output of embedding layer. G i (·) and C i (·) represent the skipping gate and classifier, and the superscript means they are inserted in ith layer.

3.2 Model Design

Different to [Liu et al., 2020; Xin et al., 2020] whose each layer will introduce one exiting classifier, we additionally plug a skipping gate into each layer. In order to understand SmartBERT better, we will first introduce the skipping gates and then introduce exiting classifiers.

The skipping gate aims to decide whether to execute or bypass the current layer like Figure 1.

With the skipping gate G i (·), the output of ith encoder block can be defined as

ଵି௜܆

where G i (Xi-1 ) ⊆ [0, 1] is the output of the ith skipping gate which represents the probability of skipping the current layer. In the training phase, G i (Xi-1 ) is a continuous value in the range 0 to 1. However, in the inference phase, G i (Xi-1 ) is a discrete value that is 0 or 1, and the inconsistent usage of skipping gates between the training and inference phase will be discussed in Section 3.5.

Next, we will introduce early exiting classifiers. All early exiting classifiers use the early exiting mechanism based on the entropy, and we define the output of classifiers as

where zi ∈ R1×C represents the probability of each category.

3.3 Training

SmartBERT is composed of pre-trained BERT, skipping gates, and early exiting classifiers. For stability, we use two stages training strategy for the model, and different stages will train different parts of the model. In this section, we will illustrate two stages separately.

First Training Stage

In the first stage, we train skipping gates and fine-tune pretrained BERT in a joint way. Other components will be fixed except the last classifier, as shown in Figure 2a. The total loss of the first stage consists of three main parts: classification error, the sloth of skipping, and cross-layer contrastive loss. It can be formalized as

where the first term represents the classification error(e.g., cross-entropy loss) like Equation 2. The second term indicates the sloth of skipping. In order to encourage the gates to skip to some extent rather than only considering the performance, we introduce the second term as a regular term. To enlarge the output of skipping gates. And λ, η ∈ [0, 1] are scaling factors that are used to balance the influence of the regular terms on the loss. The third term represents a crosslayer contrastive loss, which mainly aims to obtain a more powerful BERT in this stage.

In cross-layer contrastive learning (CCL), we assume that each token should have similar semantics across the consecutive layers because their representations do not change drastically. The experiments in Appendix C can support our motivation. The same tokens will be considered as positive pairs across consecutive layers. Otherwise, they shall be deemed to be negative pairs. The CCL loss in the first training stage can be formulated as

where pim = f(xim), xim represents the mth word representation in the ith layer. f(·) is a projection function that maps representations to another latent space where the contrastive loss is calculated. s(x, y) represents a score function (e.g., cosine similarity), and τ is temperature parameter. In this training stage, we only boost the transformer layer by the crosslayer contrastive loss.

Second Training Stage

In the second stage, the trained parameters of the first stage will be frozen, and all classifiers except the last one will be trained as shown in Figure 2b. In order to obtain more powerful classifiers, we also adopt CCL into the second training phase. Specifically, the token representations obtained by frozen transformer layers will be fed into classifiers firstly. Secondly, we extract the hidden states that only go through the first layer (i.e., self-attention layer) of classifiers, and these hidden states will be fed into the projection head to map them into the space where contrastive loss is occupied. Lastly, the contrastive loss (i.e., InfoNCE loss) is employed. Figure 3 shows the overall pipeline. The CCL loss in the second training stage can be formalized as

where ˆhim = f(him), him = Self-Attention(Xi )[m], Self-Attention(·) represents a self-attention function and [m] means we only select m-th token. Other notations have the same meanings as in Equation 4. Enumerating all tokens in all layers, the overall loss is:

In this way, the exiting classifiers can help each other in the consecutive layers, which will assist our model early exiting. In detail, our classifier is similar to [Liu et al., 2020], which consists of one attention layer and one linear layer. Each token in the i-th layer will acquire its high-level representation Hi through the attention layer. Through Equation 5, the selfattention layer will be boosted so that we can empower the classifier to produce more confident results which are helpful for early exiting based on entropy.

ଵା௜ଵܐመ

ଵା௜ଵܐ

ଵା௜ଵܠ

For keeping consistent with the inference phase, SmartBERT uses the output of skipping gates to update the hidden states of the current layer like Equation 1. Next, we define the loss of each classifier as

In the second stage, the loss function is the sum of the crossentropy loss of classifiers with cross-layer contrastive loss:

3.4 Inference

After finishing the two-stage training, SmartBERT can take advantage of skipping gates and early exiting to accelerate inference considering the complexity of samples. Only when the output of skipping gates is greater than or equal to 0.5, the model will decide to skip. The specific formula follows:

When skipping gates choose not to skip, the current early exiting classifier will be used, and all classifiers use an early exiting mechanism based on the entropy:

where zi represents a categorical distribution of the ith classifier, and S is the entropy threshold for early exiting, which is set manually. The entire inference algorithm is in Appendix A.

3.5 HardWeight Mechanism

In the previous sections, we have introduced SmartBERT training and inference. However, we find the inconsistent usage of skipping gates between the training and inference phase. Specifically, in the training phase, the output of skipping gates is a continuous probability value between 0 and 1 like Equation 1, and in the inference phase, we will discretize the output of skipping gates like Equation 9 for acceleration. To reduce the impact of inconsistent usage of skipping gates, we introduce a hard-weight mechanism for skipping gates in the training phase. And the hard weight mechanism is differentiable, we use the trick:

where 1Gsoft≥0.5 is an indicator function which returns 1 if Gsoft≥0.5 otherwise 0, and ˜Gsoft is the result of truncating the gradient of Gsoft. With this trick, we achieve two things: (1) makes the output value exactly one-hot (since we add then subtract soft value), (2) makes the gradient equal to the soft gradient (since we strip all other gradients). In the training phase, we define the output of skipping gates as the soft weight Gsoft. To keep the consistent usage of skipping gates between training and inference phases, we turn the soft weight Gsoft into discrete values (0 or 1), and keep the gradient equal to Equation 12. In this way, the output of the skipping gate becomes a discrete skipping decision in the training phase, which is consistent with the inference phase:

where fp represents forward pass, Gsoft is the soft weight and Ghard is the hard weight.

3.6 A Special TrainingWay for Skipping Gates

Although the hard weight mechanism can reduce the impact of inconsistent usage of skipping gates, directly using the hard weight mechanism will lead to some layers can not be well fine-tuned because they may be skipped. However, these layers may be used in the inference phase, which will result in unsatisfactory performance. To ensure all layers can be trained and keep the consistent usage of skipping gates in the training and inference phase, we introduce a special training way from soft weight mechanism to hard weight mechanism for skipping gates.

We use the transition from the soft weight mechanism to the hard weight mechanism to address the above problems. First, we use the soft weight mechanism to warm up all skipping gates and transformer layers and then use the hard weight mechanism to train them. Moreover, the soft weight mechanism aims to fine-tune these layers, and the hard weight mechanism is used for keeping the consistent usage of skipping gates between training and inference phases.

4 Experiments

4.1 Baselines

We compare our method with three baselines:

• BERT. A large-scale pre-trained language model based on Transformer, which is used as the backbone of our methods. In experiments, we only use the BERT-base model pre-trained by Google. [Devlin et al., 2019]
• DistillBERT. It is a smaller transformer-based model by distilling the BERT. [Sanh et al., 2019]
• Early exiting model. The dynamic early exiting method for BERT is an effective method to accelerate BERT inference. We choose some classical methods like DeeBERT [Xin et al., 2020] and FastBERT [Liu et al., 2020] as our baselines, which use the early exiting mechanism based on entropy.

4.2 Datasets

To verify the effectiveness of our methods, We conduct experiments on eight classification datasets of the GLUE benchmark[Wang et al., 2018a], including SST-2 [Socher et al., 2013], CoLA [Warstadt et al., 2019], MRPC [Dolan and Brockett, 2005], MNLI [Williams et al., 2017], QQP, QNLI [Rajpurkar et al., 2016], RTE [Bentivogli et al., 2009], and WNLI [Levesque et al., 2012].

4.3 Experimental Setup

The experiments are done on an NVIDIA 2080Ti GPU. We adopt the same parameters for BERT, DeeBERT, FastBERT, and SmartBERT. In experiment, these models use the pretrained parameters(bert-base-uncased) released by the Hugging Face Transformer Library [Wolf et al., 2019]. In the pre-trained parameters, the number of transformer layers, the dimension of hidden states, and the max length of the input sentence are set to 12, 768, and 128. We use AdamW [Kingma and Ba, 2014] to train these models with a default batch size of 32. For each task, we select the best fine-tuning learning rate(among 1e-5,2e-5,5e-5).

In the first stage of SmartBERT, we train the model with 5 epochs and select one with the best accuracy for the second stage. In the second stage, we adopt cross-layer contrastive learning to train the classifiers and train these classifiers for 4 epochs. We slightly tune the hyper-parameters across the different tasks.

Following prior work, our batch size of inference is set to 1 in the inference phase.

In inference, We followed FastBERT and used FLOPs as an indicator to evaluate the computational cost. Generally speaking, the size of the FLOPs reflects the inference speed of the model, and the smaller the FLOPs of the model are, the shorter the inference time will be. Table 2 presents the computational cost of each operation within the SmartBERT, which shows that the computational cost of the Skipping Gate and Classifier is much lower than the Transformer Layer.

4.4 Main Results

We evaluate these models in eight classification datasets of the GLUE benchmark and select different entropy thresholds to test the performance and sample-averaged FLOPs for DeeBERT, FastBERT and SmartBERT. We set three different entropy thresholds for early exiting models and compare the results with other baselines in Table 1. The results of the experiment show that our model achieves 2-3× computation reduction with minimal accuracy drops and even has a better performance than BERT. In the results, due to DistillBert cannot adaptively change the model architecture according to the sample complexity, it has a fixed computational complexity. In other words, the computational cost of DistillBert is independent of sample difficulty. In fact, the difficulty of each dataset is different, and in more datasets, our method has a better performance than DistillBert.

In Figure 4, we set several entropy thresholds and compare different models' tradeoffs in accuracy and computational cost on some GLUE datasets. From Figure 4, we conclude that:

1. Compared with FastBERT and DeeBERT, SmartBERT have a better performance at the same computational cost, which proves our model is more effective than other approaches.
2. In all datasets, the scores(accuracy/F1 score/FLOPs) of these models start to decline when the computational costs reach certain values which verify the redundant computation in original BERT, and SmartBERT dominates the performance under different computation costs comparing with FastBERT and DeeBERT.

4.5 Effectiveness of SkippingMechanism

The early exiting mechanism is based on the entropy of the corresponding classifier output. However, previous methods hardly reduce computation cost when samples are complex and difficult enough. Because the entropy of classifier output will be extremely high that leads to most exiting classifiers do not exit early. In Table 1, we can see that even if the threshold S is set as 0.5, the DeeBERT and FastBERT still have an extremely high computational cost for RTE and WNLI dataset, which means early exiting mechanism is invalid on such situations. But SmartBERT achieves a relatively large computation reduction with minimal accuracy drop, which shows the skipping mechanism is effective.

To further prove the effectiveness of the layer skipping mechanism, we disable the early exiting mechanism and only use the layer skipping mechanism during the inference phase. The empirical results are shown in Table 3, which proves that the skipping mechanism is able to reduce the computational cost compared with BERT. Moreover, combined Table 1 and Table 3, FLOPs and Accuracy are hardly changed on the RTE and WNLI datasets, which further proves that the early exiting mechanism is invalid in difficult datasets and layer skipping mechanism plays an important role in our method on such situations.

4.6 Effectiveness of Cross-Layer Contrastive Learning

We have introduced the cross-layer contrastive learning, which takes the same tokens across consecutive layers as positive pairs and the different tokens across consecutive layers as negative pairs. To further evaluate the effectiveness of contrastive learning, we compare the two cases of using contrastive learning and not using contrastive learning. Empirical results are shown in Figure 5, which shows cross-layer contrastive learning is effective for further reducing the computation.

4.7 Effectiveness of Our Special TrainingWay

In the previous section, we introduce the soft weight mechanism and hard weight mechanism. Meanwhile, we propose a special training way from soft weight mechanism to hard weight mechanism. In the experiment, we compare three different training strategies: (1) soft weight mechanism: using the continuous probability value of skipping gates for training but using the discrete value of skipping gates for inference, (2) hard weight mechanism: using the discrete value of skipping gates for training and inference and (3) the special training way: firstly using continuous probability value for training, and then using the discrete value for training.

We compare these methods on the SST-2 and QNLI datasets, and the results are shown in Figure 6. The results prove that the special training way is effective for balancing the inconsistent usage of skipping gates between training and inference phases.

5 Conclusion

In this paper, we propose a novel dynamic early exiting combined with layer skipping for BERT inference. To address the inconsistent usage of skipping gates in the inference and training phase, we propose hard weight mechanism and a special training way. In addition, we propose cross-layer contrastive learning and combine it into our training phases to boost the intermediate layers and classifiers which can further reduce computation cost. We evaluated our model on eight classification datasets of the GLUE benchmark. Empirical results show that SmartBERT could achieve performance comparable to BERT while significantly reducing computational cost. Compared to other dynamic early exiting models, SmartBERT obtain better accuracy with lower computation. Moreover, we conduct a series of ablation studies to demonstrate that each component is beneficial. In the future, we will adopt our methods to other pre-trained language models. We also will further study the combination of other acceleration methods and layer skipping mechanism.

Acknowledgements

We would like to thank anonymous reviewers for their valuable comments and suggestions. This work has been supported in part by the NSFC (No. 62272411), the Zhejiang NSF (LR21F020004), Ant Group and Alibaba-Zhejiang University Joint Research Institute of Frontier Technologies.

Contribution Statement

Boren Hu, Yun Zhu designed the research and wrote the paper. Siliang Tang was an early advocate for the research, and advised Boren Hu and Yun Zhu. Jiacheng Li provided opinions for the writing of the paper. Moreover, Boren Hu and Yun Zhu have equal contributions for this paper, and Siliang Tang is the corresponding author.