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

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3033–3044 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 3033 GAML-BERT: Improving BERT Early Exiting by Gradient Aligned Mutual Learning Wei Zhu1,2, Xiaoling Wang1∗, Yuan Ni2, Guotong Xie2,3,4 Zhen Guo5, Xiaoming Wu5 1 East China Normal University, Shanghai, China 2 Ping An Health Tech, Beijing/Shanghai, China 3 Ping An Healthcare and Tech, Shanghai, China 4 Ping An International Smart City, Shenzhen, China 5 Qingdao Eye Hospital of Shandong First Medical University, Qingdao, China Abstract In this work, we propose a novel framework, Gradient Aligned Mutual Learning BERT (GAML-BERT), for improving the early ex- iting of BERT. GAML-BERT’s contributions are two-fold. We conduct a set of pilot ex- periments, which shows that mutual knowl- edge distillation between a shallow exit and a deep exit leads to better performances for both. From this observation, we use mutu- al 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 nov- el 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 outperfor- m the state-of-the-art (SOTA) BERT early ex- iting methods. 1 Introduction Since BERT (Devlin et al., 2018), 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 (Radford et al., 2019), XLNet (Yang et al., 2019), and ALBERT (Lan et al., 2020), and so forth. These BERT-style models achieved consid- erable improvements in many Natural Language Processing (NLP) tasks by pre-training on the un- labeled corpus and fine-tuning on labeled tasks, such as text classification natural language infer- ence (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 (Fan et al., 2020; Michel et al., 2019; Zhou et al., 2020) find that ∗Corresponding author. Email: xlwang@cs.ecnu.edu.cn. large PLMs with dozens of Transformer layers are over-parameterized and suffer from the “overthink- ing” problem (Kaya et al., 2019). That is, for many input samples, their shallow representations at a shallow layer are enough to make a correct classifi- cation. Moreover, the final layer’s representations may be too overfitting to generalize.1 The over- thinking problem leads to not only poor generaliza- tion but also wasted computation. Addressing the above two issues, a branch of lit- erature focuses 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; Zhu et al., 2021; Zhu, 2021c,a), knowl- edge distillation (Sun et al., 2019b; Sanh et al., 2019; Jiao et al., 2020), weight quantization (Zhang et al., 2020; Bai et al., 2020; Kim et al., 2021) and adaptive …

Content truncated to 500 words