[논문 리뷰] Multi-head Knowledge Distillation for Model Compression

Huan Wang, Suhas Lohit외 2명|arXiv (Cornell University)|2020. 12. 05.

Advanced Neural Network Applications참고 문헌 24인용 수 5

한 줄 요약

이 논문은 다중 헤드 지식 정련(MHKD)을 제안한다. 이는 교사 및 학습자 네트워크의 중간 레이어에 보조 분류기 헤드를 추가하는 단순하면서도 효과적인 모델 압축 방법이다. 이러한 헤드의 출력 간 KL 발산을 최소화함으로써, MHKD는 다양한 아키텍처와 데이터셋에서 학습자 성능을 향상시키며, 이전의 지식 정련 방법들을 일관되게 능가하고, 대조 학습과 조합될 경우 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Several methods of knowledge distillation have been developed for neural network compression. While they all use the KL divergence loss to align the soft outputs of the student model more closely with that of the teacher, the various methods differ in how the intermediate features of the student are encouraged to match those of the teacher. In this paper, we propose a simple-to-implement method using auxiliary classifiers at intermediate layers for matching features, which we refer to as multi-head knowledge distillation (MHKD). We add loss terms for training the student that measure the dissimilarity between student and teacher outputs of the auxiliary classifiers. At the same time, the proposed method also provides a natural way to measure differences at the intermediate layers even though the dimensions of the internal teacher and student features may be different. Through several experiments in image classification on multiple datasets we show that the proposed method outperforms prior relevant approaches presented in the literature.

연구 동기 및 목표

• 중간층에서의 감독을 활용하여 모델 압축에서 지식 정련 성능을 향상시키는 것.
• 정련 과정에서 학습자 및 교사 네트워크 간 특징 차원 불일치 문제를 해결하는 것.
• 복잡한 지식 정의나 아키텍처를 도입하지 않으면서도 학습자 성능을 향상시키는 단순하고 즉시 적용 가능한 방법을 개발하는 것.
• 다양한 네트워크 아키텍처와 데이터셋, 특히 교차 아키텍처 쌍에서도 강건성을 확보하는 것.
• 기타 정련 기법(예: 메모리 백업을 활용한 대조 학습)과 호환되어 추가 성능 향상을 가능하게 하는 방법을 제공하는 것.

제안 방법

• 교사 및 학습자 네트워크의 다수의 중간 레이어에 보조 분류기 헤드를 도입한다.
• 각 중간 레이어에서 학습자 및 교사의 보조 헤드 출력을 KL 발산 손실을 통해 비교한다.
• 정규 교차 엔트로피 손실과 함께 KL 발산 손실을 학습자 훈련 중에 결합하여 특징 학습을 이끈다.
• 보조 헤드가 입력 차원과 관계없이 유사한 출력을 생성하도록 훈련되기 때문에, 이 방법은 특징 차원 불일치 문제를 자연스럽게 처리한다.
• 보조 헤드는 정련 과정 동안 훈련되지만 추론 시에는 제거되어 오직 학습자의 주 분류기만 남는다.
• 기타 정련 기법(예: 메모리 백업을 활용한 대조 학습)과 호환되어 추가 성능 향상을 가능하게 한다.

실험 결과

연구 질문

• RQ1 보조 분류기로 중간층에서의 감독이 모델 압축에서 지식 정련 성능을 향상시킬 수 있는가?
• RQ2 MHKD는 교차 아키텍처 쌍을 포함한 다양한 교사-학습자 아키텍처에서 어떻게 성능을 발휘하는가?
• RQ3 FitNet, RKD, CC와 같은 기존 정련 방법에 비해 MHKD는 정확도와 강건성 측면에서 뛰어나게 성능을 발휘하는가?
• RQ4 MHKD는 음성 쌍을 활용한 대조 학습과 같은 다른 정련 기법과 효과적으로 조합될 수 있는가?
• RQ5 보조 헤드를 위한 중간층 선택에 민감한가? 다수의 헤드를 사용할 경우 성능 향상이 이루어지는가?

주요 결과

• MHKD는 평가된 모든 데이터셋과 네트워크 쌍에서 표준 지식 정련(KD)을 일관되게 능가한다.
• CIFAR-100에서 ResNet32x4 교사와 ShuffleNetV2 학습자 조합으로 96.68%의 top-1 정확도를 기록했으며, KD보다 최대 1.2%p 높은 성능을 달성했다.
• VGG13/MobileNetV2 조합에서는 FitNet, RKD, CC에 비해 정확도를 1% 이상 향상시켰다.
• 제거 실험을 통해 다수의 보조 헤드(예: 헤드 1–3)를 사용할 경우 단일 헤드보다 더 우수한 성능을 내어 수동으로 헤드를 선택할 필요가 없음을 확인했다.
• CRD와 조합했을 때 MHKD는 최신 기술 수준의 성능을 달성했으며, 대조 학습과의 호환성과 상호보완성을 입증했다.
• 이 방법은 데이터셋 간 강건성을 보였으며, CIFAR-100과 SVHN에서 모두 모든 베이스라인보다 뛰어난 성능을 보였다. 다른 방법들의 성능 추세가 데이터셋 간으로 다를 수 있음에도 불구하고도 마찬가지로 성능을 확보했다.