Understanding and Improving Knowledge Distillation for Quantization Aware Training of Large Transformer Encoders

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 6713 - 6725 December 7-11, 2022 ©2022 Association for Computational Linguistics Understanding and Improving Knowledge Distillation for Quantization-Aware Training of Large Transformer Encoders Minsoo Kim1, Sihwa Lee2, Sukjin Hong3, Du-Seong Chang3, and Jungwook Choi1,2∗ 1Department of Electronic Engineering, Hanyang University 2Department of Artificial Intelligence, Hanyang University Seoul, Republic of Korea {minsoo2333, macto94, choij}@hanyang.ac.kr, 3KT, Seoul, Republic of Korea {sukjin.hong, dschang}@kt.com Abstract Knowledge distillation (KD) has been a ubiq- uitous method for model compression to strengthen the capability of a lightweight model with the transferred knowledge from the teacher. In particular, KD has been employed in quantization-aware training (QAT) of Trans- former encoders like BERT to improve the ac- curacy of the student model with the reduced- precision weight parameters. However, little is understood about which of the various KD approaches best fits the QAT of Transformers. In this work, we provide an in-depth analysis of the mechanism of KD on attention recovery of quantized large Transformers. In particular, we reveal that the previously adopted MSE loss on the attention score is insufficient for recov- ering the self-attention information. Therefore, we propose two KD methods; attention-map and attention-output losses. Furthermore, we explore the unification of both losses to address task-dependent preference between attention- map and output losses. The experimental re- sults on various Transformer encoder models demonstrate that the proposed KD methods achieve state-of-the-art accuracy for QAT with sub-2-bit weight quantization. 1 Introduction Knowledge distillation (KD) (Hinton et al., 2015) is a transfer learning framework to pass on knowl- edge of a large model (the teacher) to a lightweight model (the student). Numerous KD methods have been developed regarding the source of knowl- edge and distillation objectives. In many cases, Kullback-Leibler divergence (KL-Div) is used as a default distillation objective to match the soft la- bels of the teacher and the student (Hinton et al., 2015; Sanh et al., 2019). But further studies on KD suggest that internal representations also con- vey the intermediate knowledge of the teacher (Sun et al., 2019; Aguilar et al., 2020). Thus minimiz- ing the distance (e.g., mean squared error, MSE) ∗Corresponding Author of hidden state knowledge (HSK) of layers be- tween the teacher and the student has also been proposed (Sun et al., 2019; Liu et al., 2021a). KD becomes an essential model compression technique for efficiently deploying large-scale Transformer- based language models. For example, a popular Transformer encoder model, BERT, contains hun- dreds of millions of parameters, incurring profound memory and computation overhead (Devlin et al., 2019). These large-scale models require extreme compression to reduce the model footprint by 10 to 100 times. Therefore, extensive studies have been accomplished to distill efficient student mod- els (Sanh et al., 2019; Sun et al., 2019; Jiao et al., 2020; Wang et al., 2020, 2021), but their focus is limited to achieving fewer parameters. Quantization-aware training (QAT) stands out for its recent success in reducing not only the mem- ory requirements but also the …

Content truncated to 500 words