Teacher Intervention: Improving Convergence of Quantization Aware Training for Ultra-Low Precision Transformers (2023)

Abstract: Pre-trained Transformer models such as BERT have shown great success in a wide range of applications, but at the cost of substantial increases in model complexity. Quantization-aware training (QAT) is a promising method to lower the implementation cost and energy consumption. However, aggressive quantization below 2-bit causes considerable accuracy degradation due to unstable convergence, especially when the downstream dataset is not abundant. This work proposes a proactive knowledge distillation method called Teacher Intervention (TI) for fast converging QAT of ultra-low precision pre-trained Transformers. TI intervenes layer-wise signal propagation with the intact signal from the teacher to remove the interference of propagated quantization errors, smoothing loss surface of QAT and expediting the convergence. Furthermore, we propose a gradual intervention mechanism to stabilize the recovery of subsections of Transformer layers from quantization. The proposed schemes enable fast convergence of QAT and improve the model accuracy regardless of the diverse characteristics of downstream fine-tuning tasks. We demonstrate that TI consistently achieves superior accuracy with significantly lower fine-tuning iterations on well-known Transformers of natural language processing as well as computer vision compared to the state-of-the-art QAT methods.

## introduction the transformer-based pre-trained neural networks have significantly improved the performance of various applications of artificial intelligence, including natural language processing (nlp) @xcite @xcite and computer vision (cv) @xcite @xcite . the self-attention mechanism represents these models * corresponding author figure 1 : comparison of fine-tuning time and accuracy of ternary weight quantized transformer models between teacher intervention (ti) and ternarybert (bert-base and tinybert-6l for cola task). ti achieves higher accuracy within 12.5× shorter finetuning time. @xcite , which links different symbols within a sequence to obtain a relational representation. thanks to the exceptional performance of the pre-trained transformer models, there have been increasing needs for their efficient deployment. however, the gigantic size of the pre-trained transformer models hinders straightforward implementation. even relatively small models like bert-base @xcite contain a few hundred million parameters, incurring profound memory and computation overhead for resourceconstrained devices with limited memory and computing fabric. therefore, seminal research efforts attempted to reduce this burden via model compression. @xcite and @xcite pruned unimportant weights to reduce the number of parameters, while @xcite further employed low-rank matrix factorization. in addition, knowledge distillation (kd) @xcite was employed in @xcite @xcite to transfer knowledge of the original model (teacher) to the compressed one (student) by mimicking the teacher's behavior. among many model compression techniques, quantization-aware training (qat) stands out for its recent success in reducing computational complexity and memory requirements of transformer models @xcite @xcite . qat reflects quantization errors during the forward pass computation of stochastic gradient descent to train a more accurate quantized model. however, quantizing weight parameters of transformers to a precision lower than 2-bits degrades the accuracy, especially when the dataset size for the target downstream tasks is not large enough @xcite . although few-sample fine-tuning of transformer models has been reported to be highly unstable @xcite @xcite @xcite , efforts on understanding and improving qat on small dataset tasks are limited. xtc @xcite recently revealed that a simple expansion in fine-tuning iterations could heal the qat accuracy. still, it caused an order-of-magnitude increase in qat time, hindering the broad deployment of quantized transformers. this work proposes a proactive kd method called teacher intervention (ti) for fast converging qat of ultra-low precision pre-trained transformers. we reveal that the difficulty of quantization on few-sample fine-tuning originates from disruption of loss surface due to quantization error propagation. to mitigate this undesirable phenomenon, we propose ti to intervene layer-wise signal propagation with the intact signal from the teacher. ti removes the interference of propagated quantization errors to smooth out the loss surface and expedite the convergence. we further discover that subsections of transformer layers exhibit different susceptibility to quantization error for diverse downstream tasks. thus, we propose a gradual intervention mechanism that first intervenes at the attention output for stable tuning of the feed-forward network, followed by self-attention map intervention for its recovery from quantization. the proposed gradual intervention along the subsections of transformer layers enables fast convergence of qat and improves the model accuracy regardless of the diverse characteristics of downstream fine-tuning tasks. we perform an extensive evaluation on various fine-tuned transformers (bert-base/large @xcite , tinybert-4l/6l @xcite , and skipbert-6l @xcite for nlp, and vit (dosovitskiy et al., 2020) for cv) and demonstrate that ti consistently achieves superior accuracy with lower fine-tuning iterations compared to the state-of-the-art qat methods. in particular, ti outperforms ternarybert @xcite ) on glue tasks with 12.5× savings in fine-tuning hours, as shown in fig. 1 . ## experiments in this section, we evaluate ti for qat of finetuned bert and vision transformer with ternary weight quantization. we demonstrate that ti significantly boosts the convergence of qat to achieve higher model accuracy within shorter fine-tuning time compared to the state-of-the-art qat methods (ternarybert @xcite and xtc ( wu et al., 2022 ## conclusion in this work, we proposed a proactive knowledge distillation method for improving the convergence of qat for ultra-low precision transformers called teacher intervention. the proposed method intervenes in the propagation of quantization error to suppress accuracy degradation and improve qat's convergence speed. we demonstrate that the proposed method outperforms the state-of-the-art qat methods in achieving higher accuracy on various fine-tuned transformers with smaller fine-tuning budgets. our code is available at https://github.com/marsjacobs/ti-kd-qat .