[论文解读] QKD: Quantization-aware Knowledge Distillation

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

Advanced Neural Network Applications参考文献 52被引用 46

一句话总结

本文提出量化感知知识蒸馏（QKD），一个三阶段框架，协同学习量化与知识蒸馏以提升极低位量化网络。

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.

研究动机与目标

• 动机：在边缘设备效率下，联合优化量化与知识蒸馏的必要性。
• 提出一个三阶段框架（自我学习、共同学习、辅导）以在低位量化下稳定并提升KD。
• 证明QKD在CIFAR和ImageNet上在2、3、4位量化下达到最先进的准确率。
• 展示对深度可分离卷积（MobileNetV2、EfficientNet）等通常难以量化的网络的有效性。

提出的方法

• 采用权重和激活的可训练均匀量化方案，按层设定区间值（I_W、I_X）。
• 实现三阶段训练协议：阶段1 自我学习（仅用任务损失训练低位比特学生模型），阶段2 共同学习（对量化友好教师和学生进行在线训练并结合KD），阶段3 辅导（固定教师并用KD对学生进行微调）。
• 使用基于KL散度的在线KD，教师与学生之间温度T=2，同时结合两网络的交叉熵损失。
• 使用直通估计器（STE）在不可微的量化器上进行反向传播，并与权重一起训练区间值。
• 对所有卷积和线性层应用量化器；前后两层量化为8比特，以确保硬件兼容性。
• 表明辅导阶段在学生性能方面可达到甚至超过仅共同学习的效果。

实验结果

研究问题

• RQ1 在不同训练阶段协调量化与KD是否可以提高超低位网络的准确性？
• RQ2 在量化KD设置中，经过共同学习中适应的可学习教师是否优于固定的预训练教师？
• RQ3 三阶段的QKD框架对MobileNetV2、EfficientNet等深度可分离卷积是否有效？
• RQ4 使用QKD在CIFAR-10/100和ImageNet上进行2、3、4位量化时可以达到哪些准确度提升？

主要发现

• QKD 超越现有最先进的方法（例如在 ResNet-18 W4A4 上提升 1.3%，在 MobileNetV2 W4A4 上提升 2.6%）。
• QKD 可以在 ResNet 的 W3A3 和 MobileNetV2 的 W6A6 时恢复到全精度准确率。
• 自我学习提供了良好的初始化，以减轻低位量化网络上KD正则化的影响。
• 共同学习产生的教师比固定教师更具量化友好性与更强的 KD 迁移能力。
• 辅导阶段在冻结教师的情况下进一步提高或达到共同学习的表现，同时降低训练成本。