[논문 리뷰] An Embarrassingly Simple Approach for Knowledge Distillation

Mengya Gao, Yujun Shen외 6명|arXiv (Cornell University)|2018. 12. 05.

Advanced Neural Network Applications참고 문헌 40인용 수 12

한 줄 요약

이 논문은 단계별 지식 증류(Stage-by-Stage Knowledge Distillation, SSKD)를 제안한다. 이는 지식 증류를 두 단계로 분리하는 단순하면서도 효과적인 방법으로, 첫 번째로 레이블을 사용하지 않고 교사에서 학생의 백본으로 중간 특징 표현을 증류하고, 두 번째로는 진짜 레이블을 사용하여 태스크에 특화된 헤드만 미세조정하는 방식이다. SSKD는 손실 가중치 초모수 조정이 필요 없으며, ImageNet, CIFAR-100, COCO 객체 검출, IJB-A 얼굴 인식 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성한다. COCO에서 최대 2.8%p의 AP 향상과 ImageNet에서 최대 1.8%p의 정확도 향상을 기록하였다.

ABSTRACT

Knowledge Distillation (KD) aims at improving the performance of a low-capacity student model by inheriting knowledge from a high-capacity teacher model. Previous KD methods typically train a student by minimizing a task-related loss and the KD loss simultaneously, using a pre-defined loss weight to balance these two terms. In this work, we propose to first transfer the backbone knowledge from a teacher to the student, and then only learn the task-head of the student network. Such a decomposition of the training process circumvents the need of choosing an appropriate loss weight, which is often difficult in practice, and thus makes it easier to apply to different datasets and tasks. Importantly, the decomposition permits the core of our method, Stage-by-Stage Knowledge Distillation (SSKD), which facilitates progressive feature mimicking from teacher to student. Extensive experiments on CIFAR-100 and ImageNet suggest that SSKD significantly narrows down the performance gap between student and teacher, outperforming state-of-the-art approaches. We also demonstrate the generalization ability of SSKD on other challenging benchmarks, including face recognition on IJB-A dataset as well as object detection on COCO dataset.

연구 동기 및 목표

• 지식 증류에서 초모수 민감도 문제, 특히 태스크 손실과 KD 손실을 균형 잡는 데 어려움을 해결하기 위해.
• 훈련 과정을 두 개의 별도 단계로 분리함으로써 다양한 데이터셋과 태스크에서 지식 증류의 일반화성과 안정성을 향상시키기 위해.
• 중간 특징의 점진적, 단계별 증류가 단일 단계 증류보다 더 나은 지식 전달을 이끌 수 있음을 보여주기 위해.
• 이전 방법이 직접 적용되지 않았던 비분류 태스크, 예를 들어 객체 검출 및 얼굴 인식에 지식 증류를 확장하기 위해.
• COCO에서 FPN과 RetinaNet에 SSKD를 적용하여 객체 검출 분야의 지식 증류 기준 성능을 새롭게 설정하기 위해.

제안 방법

• 학생 훈련을 두 단계로 분해: (1) 진짜 레이블을 사용하지 않고 교사에서 학생의 백본으로 중간 특징 표현을 증류하는 단계, (2) 백본을 고정한 채로 레이블이 있는 데이터를 사용하여 태스크에 특화된 헤드만 훈련하는 단계.
• 단계별 지식 증류(SSKD)를 구현하여, 네트워크의 각 단계에서 순차적으로 증류를 수행하고, 한 서브넷씩 특징을 모방한다.
• 첫 번째 단계 동안 교사 및 학생 네트워크의 대응하는 레이어 간에 특징 수준의 증류 손실(예: MSE 또는 KL 발산)을 사용한다.
• 증류 후 학생의 백본을 고정하고, 두 번째 단계에서 교차 엔트로피 또는 기타 태스크에 특화된 손실을 사용하여 헤드만 미세조정한다.
• ResNet 기반 다양한 아키텍처와 태스크에 SSKD를 적용하며, ImageNet 및 CIFAR-100에서의 ResNets, COCO 객체 검출에서의 FPN 및 RetinaNet에 적용한다.
• 각 단계에서 대부분의 네트워크를 동결함으로써 훈련 효율성을 확보하여 훈련 시간이 크게 증가하지 않도록 한다.

실험 결과

연구 질문

• RQ1 지식 증류를 백본 특징 모방과 태스크 헤드 미세조정으로 분리하면 손실 가중치 초모수 조정이 필요 없어지는가?
• RQ2 중간 특징의 단계적, 점진적 증류가 고정된 손실 가중치를 가진 종단 간 공동 훈련보다 더 높은 성능을 낼 수 있는가?
• RQ3 제안된 방법이 이미지 분류, 객체 검출, 얼굴 인식 등 다양한 태스크에 일반화되는가?
• RQ4 SSKD는 ImageNet 및 COCO와 같은 표준 벤치마크에서 최신 기술 수준의 지식 증류 방법과 비교해 어떻게 성능을 낼 수 있는가?
• RQ5 SSKD는 KD가 덜 탐색된 복잡한 검출 아키텍처인 FPN 및 RetinaNet에 효과적으로 적용될 수 있는가?

주요 결과

• ResNet-18을 학생, ResNet-34를 교사로 사용한 ImageNet에서 SSKD는 두 번째로 좋은 방법보다 상위 1% 정확도를 1.8%p 향상시켜 80.5%의 정확도를 달성하였다.
• FPN에서 COCO 객체 검출을 수행한 결과, ResNet-152에서 증류한 ResNet-50의 AP는 37.7%에서 40.5%로 상승하여 2.8%p 향상되었다.
• RetinaNet에서 SSKD는 ResNet-152에서 증류한 ResNet-18의 AP를 32.4%에서 35.6%로 끌어올려 3.2%p 향상되었다.
• IJB-A 얼굴 인식에서 SSKD는 강력한 일반화 성능을 보이며 기준 방법보다 일관된 향상을 보였다.
• COCO에서 모든 객체 크기 범주(AP_S, AP_M, AP_L)에서 SSKD는 일관된 성능 향상을 보이며 尺도 변화에 대한 강건성을 입증하였다.
• 제거 분석 결과 단계별 증류가 단일 단계 증류보다 우수했으며, 공동 최적화보다 두 단계 훈련 방식이 더 안정적이고 효과적이었다.