[논문 리뷰] QKD: Quantization-aware Knowledge Distillation

Jangho Kim, Yash Bhalgat외 3명|arXiv (Cornell University)|2019. 11. 28.

Advanced Neural Network Applications참고 문헌 52인용 수 46

한 줄 요약

본 논문은 Quantization-aware Knowledge Distillation (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를 안정화하고 향상시키기 위한 세 단계 프레임워크(Self-studying, Co-studying, Tutoring)를 제안한다.
• QKD가 CIFAR 및 ImageNet에서 2-, 3-, 4-bit 양자화로 최첨단 정확도에 도달함을 보여준다.
• 일반적으로 양자화가 어려운 depthwise-separable convolution(MobileNetV2, EfficientNet)에 대한 효과성을 제시한다.

제안 방법

• 가중치와 활성화에 대해 계층별 간격 값(I_W, I_X)을 갖는 학습 가능한 균일 양자화 스킴을 채택한다.
• 세 단계 학습 프로토콜을 구현한다: Phase 1 self-studying(작업 손실만으로 저비트 학생 모델을 학습), Phase 2 co-studying(KD를 활용한 양자화 친화적 교사와 학생의 온라인 학습), Phase 3 tutoring(교사를 고정하고 KD로 학생을 미세 조정).
• 온도 T=2인 KL 발산 기반 온라인 KD를 교사와 학생 간에 사용하고, 두 네트워크의 교차 엔트로피 손실도 포함한다.
• 미분 불가능한 양자화기들을 역전파하기 위해 Straight-Through Estimator (STE)를 활용하고 가중치와 함께 간격 값을 학습한다.
• 모든 Conv 및 Linear 계층에 양자화기를 적용하되, 하드웨어 호환성을 위해 첫 번째 및 마지막 계층은 8비트로 양자화한다.
• Tutoring 단계가 학생 성능 측면에서 Co-studying만으로도 달성한 것과 같거나 더 뛰어남을 보인다.

실험 결과

연구 질문

• RQ1 서로 다른 학습 단계를 가로지르는 양자화와 KD의 조정이 초저비트 네트워크의 정확도를 향상시킬 수 있는가?
• RQ2 학습 가능한 교사(코-studying 중에 적응된)가 양자화된 KD 설정에서 고정된 사전 학습 교사보다 더 나은 성능을 내는가?
• RQ3 세 단계의 QKD 프레임워크가 MobileNetV2, EfficientNet 등의 depthwise-separable 합성곱에 대해 효과적인가?
• RQ4 QKD를 사용하여 CIFAR-10/100 및 ImageNet에서 2-, 3-, 4-bit 양자화를 적용했을 때 어떤 정확도 향상을 얻을 수 있는가?

주요 결과

• QKD는 기존의 최첨단 방법들보다 성능이 뛰어나며(예: ResNet-18에서 W4A4로 1.3% 향상, MobileNetV2에서 W4A4로 2.6% 향상).
• QKD는 ResNet에서 W3A3에서, MobileNetV2에서 W6A6에서 풀 정밀도 정확도를 회복할 수 있다.
• Self-studying은 저비트 양자화된 네트워크에서 KD 정규화를 완화하는 데 좋은 초기화를 제공한다.
• Co-studying은 고정된 교사보다 양자화에 더 친화적이고 강력한 교사를 제공하여 KD 전달을 개선한다.
• Tutoring 단계는 교사를 고정함으로써 학습 비용을 줄이면서 코-studying의 성능을 더 개선하거나 비슷하게 만든다.