GAML - BERT : Improving BERT Early Exiting by Gradient Aligned Mutual Learning (2021)

Abstract: In this work, we propose a novel framework, Gradient Aligned Mutual Learning BERT (GAML-BERT), for improving the early exiting of BERT. GAML-BERT’s contributions are two-fold. We conduct a set of pilot experiments, which shows that mutual knowledge distillation between a shallow exit and a deep exit leads to better performances for both. From this observation, we use mutual learning to improve BERT’s early exiting performances, that is, we ask each exit of a multi-exit BERT to distill knowledge from each other. Second, we propose GA, a novel training method that aligns the gradients from knowledge distillation to cross-entropy losses. Extensive experiments are conducted on the GLUE benchmark, which shows that our GAML-BERT can significantly outperform the state-of-the-art (SOTA) BERT early exiting methods.

## introduction since bert @xcite , the pre-trained language models (plms) are dominating the field of natural language processing (nlp). the recent years have witnessed the rise of many plms, such as gpt @xcite , xlnet @xcite , and albert @xcite , 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 classification natural language inference (nli), sequence labeling. however, plms are notorious for being gigantic and slow in both training and inference. their significant inference latencies pose great challenges to deployment in real-time applications, such as chat-bots and search engines. in addition, previous literature @xcite @xcite find that large plms with dozens of transformer layers are over-parameterized and suffer from the "overthinking" problem @xcite . that is, for many input samples, their shallow representations at a shallow layer are enough to make a correct classification. moreover, the final layer's representations may be too overfitting to generalize. 1 the overthinking problem leads to not only poor generalization but also wasted computation. addressing the above two issues, a branch of literature focuses on making plms' inference more efficient via network pruning @xcite @xcite @xcite , knowledge distillation @xcite @xcite , weight quantization @xcite @xcite and adaptive inference @xcite @xcite . the adaptive inference has drawn much attention. the adaptive inference is designated to process simple examples with only shallow layers of bert and predict 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 controlled with certain hyper-parameters to handle drastic changes in request traffic. what is more, it can address the over-thinking problem and improve the model's generalization ability. early exiting is one of the most important adaptive inference methods @xcite . as depicted in figure 1 , 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 is designed to decide whether to exit at each layer given the currently obtained predictions (from previous and current layers) @xcite @xcite . in this mode, different samples can exit at different depths. knowledge distillation (kd) @xcite is of essential importance for improving early exiting performances. the traditional belief of kd is that the teacher model (usually a stronger model) teaches a lower-capacity student model "dark knowledge" through providing soft targets. as an application of this traditional belief, recent studies by phuong and lampert (2019); @xcite improve the training procedure by incorporate kd losses, which encourages the early exits to mimic the output distributions of the last exit. @xcite challenge the common belief of kd by revealing that knowledge distillation is actually a learned label smoothing (ls) regularization @xcite , and a weaker teacher can also enhance a stronger student's performance via knowledge distillation. @xcite shows that asking all the exits to learn from one another (mutual learning) is beneficial for early exiting. @xcite introduce dense pairwise knowledge matching operations at certain intermediate layers of deep convolutional networks during training, which are demonstrated to be beneficial for the whole network's generalization. however, @xcite propose this mutual learning framework intuitively, and does not fully explore the underlying motivations. @xcite focuses on improving the whole network, and does not investigate how these mutual learning (or knowledge matching) affects each layer's exiting performances. in this work, we first conduct a series of exploratory experiments called pairwise mutual learning (pml). pml selects two exits of bert and consider the finetuning of these exits with or without adding different knowledge distillation settings. our experiments show the following three approaches can imrpove the shallower exit's performances: (1) adding supervisions to deeper exits; (2) kd from deeper exits; (3) further asking the deeper exits to learn from this shallow exits. in addition, via the above approaches, the deeper exit's performance also improves. our experimental findings is consistent with and complement the conclusions of @xcite and @xcite . and thus we adopt the mutual learning (ml) framework to enhance the training of early exits. ml asks all the exits to learn from one another (depicted by figure 1 ), thus fully releasing knowledge transfer and regularization capabilities of kd. further, by analyzing the directions of gradients from cross-entropy loss and distillation loss (denoted as g ce and g kd ), we find that g kd is often in a conflicting direction with g ce . we hypothesize that g kd⊥ce , g kd 's part that is orthogonal with g ce , will extract the model from moving toward optimum. thus we propose a novel optimization mechanism called gradient alignment (ga). as depicted by 2, ga will project g kd onto the direction of g ce . we propose two versions of ga, ga-soft, and ga-hard, which only differ in how to deal with g kd when the angle of the two gradients is larger than 90 • . we will call our framework gradient aligned mutual learning for bert (gaml-bert). extensive experiments are conducted on the glue benchmark @xcite and show that gaml-bert outperforms existing sota bert early exiting methods, sometimes by a large margin. deeper analysis and ablation studies result in the following main takeaways: (a) knowledge distillation among the exits can improve their performances, especially for the shallow ones; (b) our gradient alignment method can improve the training procedure and thus improve the model's generalization performances. our contributions are summarized as follows: ## preliminaries 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), @xmath0, @xmath1, @xmath2, 2, ..., n }, e.g., sentences, and the number of classes is k. ## mutual learning in light of the above analysis, and following zhu (2021b) and @xcite , we adopt the mutual learning framework (figure 1 ) to explore the potentials of early exits. that is, all the exiting classifiers learn from one another. the loss terms from this fully mutual learning framework are added to the cross-entropy losses in eq. 1, and the loss objective becomes ## gradient alignment the ml training mechanism can stabilize the training process and lead to better optimization by implementing rich regularizations over all the exits. however, during experiments, we still observe that the weight of the kd loss α has a large impact on the model performances, and the better performances are achieved by setting α to be relatively small (e.g., 0.1 or 0.2). our observations are consistent with the experimental observations of @xcite ; @xcite . intuitively, it seems that the kd objective is conflicting with the ce loss to a certain degree. to visualize the interactions between crossentropy loss and kd loss, we separately compute the gradients derived from the two loss objectives, g @xmath3. their an-gle is given by during the bert model fine-tuning on sst-2 with the mutual learning training objective, we calculate the angles of the gradients g kd and g ce on each training step. the distribution of the gradient angles are plotted in figure 3 . we can see that about half of the optimization steps γ is larger than 90 • , meaning that they have conflicting directions for optimization. this observation motivates us to think that we may obtain better convergences if we can somehow align the two gradients. thus, the above observation naturally leads to the following hypothesis: hypothesis 1 (h1): dropping off the part of kd's gradient that conflicts with ce's gradient and only keeps the part aligned with the latter can improve the trained model's performances. thus, we propose a novel optimization method, gradient alignment (ga), to align kd's gradient g kd with ce's gradient g ce . ga is depicted in figure 2 . when we project g kd on g ce , the projected vector is given by denote the final modified gradient as g ga . we propose two versions of ga, as follows. ga-soft when the angle θ is larger than 90 • , g p kd is also added to g ce . thus g ga is given by in ga-soft, when the angle θ is larger than 90 • , g p kd is in the opposite direction with g ce , thus might slow down or reverse this gradient descent direction. ga-hard when the angle θ is larger than 90 • , g p kd is not added to g ce . thus g ga is given by ## experiments 6.1 datasets 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 @xcite . ## conclusion in this work, we propose gaml-bert, a novel framework for improving plms' early exiting. our contributions are three-fold. firstly, we conduct a series of exploratory experiments, which shows that mutual knowledge distillation between a pair of exits can boost the performances of both. following this observation, it is natural to apply mutual learning (ml), that is, asking all the exits to learn from each other, to enhance the performances of bert early exiting. second, we propose ga, a novel training method that aligns the gradients from knowledge distillation to cross-entropy losses. experiments on the glue benchmark datasets show that our framework can improve plms' early exiting performances, especially under high latency requirements. in addition, our framework is off-theshelf and can be adapted to various early exiting strategies. the 11th international joint conference on natural language processing: student research workshop, pages 33-43, online. association for computational linguistics. wei zhu. 2021b. leebert: learned early exit for bert with cross-level optimization. in proceedings of the 59th annual meeting of the association for computational linguistics and the 11th international joint conference on natural language processing (volume 1: long papers), pages 2968-2980, online. association for computational linguistics. wei zhu. 2021c. mvp-bert: multi-vocab pretraining for chinese bert. in proceedings of the 59th annual meeting of the association for computational linguistics and the 11th international joint conference on natural language processing: student research workshop, pages 260-269, online. association for computational linguistics. wei zhu, yuan ni, xiaoling wang, and guotong xie. 2021. discovering better model architectures for medical query understanding. in proceedings of the 2021 conference of the north american chapter of the association for computational linguistics: human language technologies: industry papers, pages 230-237, online. association for computational linguistics. in figure 5 , we demonstrate the over-thinking problem of bert on sst-2 and cola. to obtain the performance of bert-base's layer i, we insert a classifier at layer i and fine-tune bert-base on the train set. we can see that the last layer does not obtain the best performances.