[Paper Review] QKD: Quantization-aware Knowledge Distillation

Jangho Kim, Yash Bhalgat+3 more|arXiv (Cornell University)|Nov 28, 2019

Advanced Neural Network Applications52 references46 citations

TL;DR

The paper introduces Quantization-aware Knowledge Distillation (QKD), a three-phase framework that coordinately learns quantization and knowledge distillation to improve very low-bit quantized networks.

ABSTRACT

Quantization and Knowledge distillation (KD) methods are widely used to reduce memory and power consumption of deep neural networks (DNNs), especially for resource-constrained edge devices. Although their combination is quite promising to meet these requirements, it may not work as desired. It is mainly because the regularization effect of KD further diminishes the already reduced representation power of a quantized model. To address this short-coming, we propose Quantization-aware Knowledge Distillation (QKD) wherein quantization and KD are care-fully coordinated in three phases. First, Self-studying (SS) phase fine-tunes a quantized low-precision student network without KD to obtain a good initialization. Second, Co-studying (CS) phase tries to train a teacher to make it more quantizaion-friendly and powerful than a fixed teacher. Finally, Tutoring (TU) phase transfers knowledge from the trained teacher to the student. We extensively evaluate our method on ImageNet and CIFAR-10/100 datasets and show an ablation study on networks with both standard and depthwise-separable convolutions. The proposed QKD outperformed existing state-of-the-art methods (e.g., 1.3% improvement on ResNet-18 with W4A4, 2.6% on MobileNetV2 with W4A4). Additionally, QKD could recover the full-precision accuracy at as low as W3A3 quantization on ResNet and W6A6 quantization on MobilenetV2.

Motivation & Objective

• Motivate the need to jointly optimize quantization and knowledge distillation for edge-device efficiency.
• Propose a three-phase framework (Self-studying, Co-studying, Tutoring) to stabilize and enhance KD with low-bit quantization.
• Demonstrate that QKD achieves state-of-the-art accuracy on CIFAR and ImageNet with 2-, 3-, and 4-bit quantization.
• Show effectiveness on depthwise-separable convolutions (MobileNetV2, EfficientNet) which are typically hard to quantize.

Proposed method

• Adopt a trainable uniform quantization scheme for weights and activations with per-layer interval values (I_W, I_X).
• Implement a three-phase training protocol: Phase 1 self-studying (train low-bit student with task loss only), Phase 2 co-studying (online training of a quantization-friendly teacher and student with KD), Phase 3 tutoring (freeze teacher and fine-tune student with KD).
• Use KL divergence-based online KD between teacher and student with temperature T=2, incorporating cross-entropy losses for both networks.
• Utilize straight-through estimator (STE) to backprop through non-differentiable quantizers and train interval values alongside weights.
• Quantizers applied to all Conv and Linear layers; first and last layers quantized to 8 bits to ensure hardware compatibility.
• Show that tutoring phase can match or surpass co-studying alone in terms of student performance.

Experimental results

Research questions

• RQ1 Can coordinating quantization and KD across distinct training phases improve accuracy for ultra-low-bit networks?
• RQ2 Does a learnable teacher (adapted during co-studying) outperform a fixed pre-trained teacher in quantized KD settings?
• RQ3 Is the three-phase QKD framework effective on depthwise separable convolutions like MobileNetV2 and EfficientNet?
• RQ4 What accuracy gains can be achieved on CIFAR-10/100 and ImageNet with 2-, 3-, and 4-bit quantization using QKD?

Key findings

• QKD outperforms existing state-of-the-art methods (e.g., 1.3% improvement on ResNet-18 with W4A4, 2.6% on MobileNetV2 with W4A4).
• QKD can recover full-precision accuracy at W3A3 for ResNet and at W6A6 for MobileNetV2.
• Self-studying provides a good initialization to mitigate KD regularization on low-bit quantized networks.
• Co-studying yields a more quantization-friendly and powerful teacher than a fixed teacher, improving KD transfer.
• Tutoring phase further improves or matches co-studying performance while reducing training cost by freezing the teacher.