KDLSQ-BERT: A Quantized Bert Combining Knowledge Distillation with Learned Step Size Quantization

TL;DR

KDLSQ-BERT introduces a distillation-aware learned step size quantization framework to compress BERT by training a quantized student with LSQ while distilling knowledge from a full-precision teacher. It combines multiple distillation signals (hidden states, attention, and logits) with optional ground-truth guidance to stabilize training, and it employs a novel scale-factor initialization to accelerate convergence. Empirical results on GLUE and SQuAD show significant improvements over existing BERT quantization methods, including near-full-precision accuracy at 2-bit weights and up to 14.9x compression. This approach enables efficient transformer inference on resource-constrained devices and demonstrates robust performance across low- to high-bit quantization settings.

Abstract

Recently, transformer-based language models such as BERT have shown tremendous performance improvement for a range of natural language processing tasks. However, these language models usually are computation expensive and memory intensive during inference. As a result, it is difficult to deploy them on resource-restricted devices. To improve the inference performance, as well as reduce the model size while maintaining the model accuracy, we propose a novel quantization method named KDLSQ-BERT that combines knowledge distillation (KD) with learned step size quantization (LSQ) for language model quantization. The main idea of our method is that the KD technique is leveraged to transfer the knowledge from a "teacher" model to a "student" model when exploiting LSQ to quantize that "student" model during the quantization training process. Extensive experiment results on GLUE benchmark and SQuAD demonstrate that our proposed KDLSQ-BERT not only performs effectively when doing different bit (e.g. 2-bit $\sim$ 8-bit) quantization, but also outperforms the existing BERT quantization methods, and even achieves comparable performance as the full-precision base-line model while obtaining 14.9x compression ratio. Our code will be public available.

Figures (10)

• Figure 1: The training framework of our proposed KDLSQ-BERT. In fact, the student BERT is inserted into LSQ quantization operations already after the training framework is built. The distillation loss $Loss_{kd}$ related to the knowledge of teacher BERT includes the hidden-states-based distillation loss $Loss_{hidden}$, the attention-based distillation loss $Loss_{att}$, as well as the prediction-layer-based distillation loss $Loss_{pre}$.
• Figure 2: The statistical histogram of a weight tensor from a well-trained TinyBERT. Compared with $s_1$ and $s_2$, $s_0$ is a much better initialized scale-factor (or truncated threshold) for quantization training since the main information of the weight tensor can be retained.
• Figure 3: The impact on total training loss, distillation loss, ground truth loss and accuracy when implementing our proposed KDLSQ-BERT by using different scale-factor initialization. The quantization bit "W-E-A (#bit)" is set to 2-2-8. The experimental results are tested by adopting "mnli" and "TinyBERT".
• Figure 4: The impact on total training loss, distillation loss, ground truth loss and accuracy when implementing our proposed KDLSQ-BERT by using different scale-factor initialization. The quantization bit "W-E-A (#bit)" is set to 4-4-8. The experimental results are tested by adopting "mnli" and "TinyBERT".
• Figure 5: The impact on total training loss, distillation loss, ground truth loss and accuracy when implementing our proposed KDLSQ-BERT by using different scale-factor initialization.The quantization bit "W-E-A (#bit)" is set to 6-6-8. The experimental results are tested by adopting "mnli" and "TinyBERT".