[論文レビュー] QKD: Quantization-aware Knowledge Distillation

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

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

ひとこと要約

本論文は Quantization-aware Knowledge Distillation (QKD) を提案する。これは three-phase framework を通じて、量子化と知識蒸馏を協調的に学習し、非常に低ビットの量子化ネットワークの性能を向上させる。

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を安定化・向上させるための3段階フレームワーク（Self-studying、Co-studying、Tutoring）を提案する。
• QKD が CIFAR と ImageNet で 2-bit、3-bit、4-bit 量子化の状態最先端の精度を達成することを示す。
• 通常量子化が難しい depthwise-separable convolution（MobileNetV2、EfficientNet）での有効性を示す。

提案手法

• 重みと活性化のために層ごとの区間値（I_W、I_X）を持つ学習可能な一様量子化方式を採用する。
• 3段階のトレーニングプロトコルを実装する： Phase 1 自習（タスク損失のみで低ビットの学生を訓練）、Phase 2 共習（KD で量子化に適した教師と学生をオンライン訓練）、Phase 3 指導（教師を固定し、学生を KD で微調整）。
• 教師と学生の間で温度 T=2 のKL発散ベースのオンラインKDを使用し、両ネットワークのクロスエントロピー損失を組み込む。
• 微分不可能な量子化器を逆伝播させるためにストレートスルー推定機（STE）を用い、重みとともに区間値を訓練する。
• 量子化器をすべての Conv および Linear 層に適用する；ハードウェア互換性を確保するため、最初と最後の層は 8 ビットに量子化する。
• 指導フェーズは、共習のみと比較して学生の性能を同等またはそれを上回ることができることを示す。

実験結果

リサーチクエスチョン

• RQ1 量子化とKDを異なるトレーニング段階で協調させることは、 ultra-low-bit ネットワークの精度を向上させるか。
• RQ2 共習中に適応される学習可能な教師は、固定された事前訓練教師よりも量子化KD設定で有効か。
• RQ3 3段階の QKD フレームワークは MobileNetV2 や EfficientNet のような depthwise separable convolution に対して有効か。
• RQ4 QKD を用いた 2-bit、3-bit、4-bit 量子化で CIFAR-10/100 および ImageNet の精度向上はどの程度か。

主な発見

• QKD は既存の最先端手法を上回る（例：ResNet-18 の W4A4 で 1.3% 改善、MobileNetV2 の W4A4 で 2.6%）。
• QKD は ResNet の W3A3、MobileNetV2 の W6A6 でそれぞれ全精度に回復させることができる。
• Self-studying は低ビット量子化ネットワークにおけるKD正則化を緩和するための良い初期化を提供する。
• Co-studying は固定教師より量子化に優しく強力な教師を生み出し、KDの伝達を改善する。
• Tutoring フェーズは co-studying の性能をさらに改善または同等にしつつ、教師を固定することで訓練コストを削減する。