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

Abstract: Knowledge distillation (KD) has been a ubiquitous 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 Transformer encoders like BERT to improve the accuracy 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 recovering 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 results 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.

## introduction knowledge distillation (kd) @xcite is a transfer learning framework to pass on knowledge of a large model (the teacher) to a lightweight model (the student). numerous kd methods have been developed regarding the source of knowledge and distillation objectives. in many cases, kullback-leibler divergence (kl-div) is used as a default distillation objective to match the soft labels of the teacher and the student @xcite . but further studies on kd suggest that internal representations also convey the intermediate knowledge of the teacher @xcite . thus minimizing the distance (e.g., mean squared error, mse) of hidden state knowledge (hsk) of layers between the teacher and the student has also been proposed @xcite . kd becomes an essential model compression technique for efficiently deploying large-scale transformerbased language models. for example, a popular transformer encoder model, bert, contains hundreds of millions of parameters, incurring profound memory and computation overhead @xcite . 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 models @xcite @xcite @xcite , but their focus is limited to achieving fewer parameters. quantization-aware training (qat) stands out for its recent success in reducing not only the memory requirements but also the computational complexity of transformer models @xcite @xcite . although qat reflects quantization errors during the forward pass computation of stochastic gradient descent to train a model more robust to quantization errors, quantizing weight parameters of transformers to a precision lower than 2-bits degrades the accuracy. therefore, many recent qat techniques employed the kd framework to distill the capability of the full-precision teacher to the student model with reduced-precision parameters @xcite @xcite . however, little is understood about which of the various kd approaches best fits the qat of large transformers. without careful justification, most prior works adopted the layer-wise distillation of the attention score and the transformer output with the mse loss in addition to the basic kl-div loss on the model output. therefore, it is unclear if such kd setting is the most helpful for qat on large-scale transformer encoders like bert-large. in this work, we provide an in-depth analysis of kd on attention recovery for qat of transformers in terms of the knowledge sources and the objectives. we first reveal that all-layer kd of the intermediate transformer layer is essential for qat, in contrast to the kd-based model compression. in the case of bert-base, we further discover that the kl-div-based kd on attention-map (called attention-map loss) outperforms the prior kd technique that takes mse loss on the attention score. however, the attention-map loss is insufficient for the large transformer encoders since weight quantization disrupts attention propagation for specific nlp tasks when there are many layers. therefore, we devise an insightful kd, mse loss on attention output (called attention-output loss), and help preserve attention recovery along with many layers. the proposed attention-map and output losses and their combination are evaluated on various transformer encoder models (bert-base/large and a bert-like korean language model (ulm). the experimental results demonstrate that the proposed kd methods significantly boost the model accuracy surpassing the state-of-the-art for qat with aggressive sub-2-bit weight quantization. we summarize our contributions as follows: ## prior kd techniques for qat in this section, we investigate prior kd techniques for qat evaluated on bert-base. as discussed earlier, kd techniques commonly used for qat include 1) all-layer distillation and 2) distillation on sa-gen. first, we provide justification and improvement on these techniques. then we further showcase the limitation when they are applied to large-scale transformer encoders. ## kd for qat on large transformers we extend the investigation of kd techniques to qat on large transformer models. in this section, we first reveal the limitation of the attention-map loss due to the task-dependent characteristics. then we propose a new kd loss, the attention-output loss, to address this challenge. lastly, we propose a combination of the two losses to handle taskdependent characteristics. ## conclusion in this work, we investigate the mechanism of knowledge distillation (kd) for qat of large transformers. we propose two kd methods, attention-map, and attention-output losses, to improve the recovery of the self-attention information. the experimental results on various transformer encoder models demonstrate that the proposed kd methods and their combination achieve state-of-the-art accuracy for qat with sub-2-bit weight quantization. our code is available at https://github.com/marsjacobs/kd-qat-large-enc .