[논문 리뷰] Rethinking Knowledge Distillation via Cross-Entropy

Zhendong Yang, Zhe Li외 5명|arXiv (Cornell University)|2022. 08. 22.

Advanced Neural Network Applications인용 수 13

한 줄 요약

이 논문은 고전적 KD 손실을 교차 엔트로피(CE) 손실과 확률 척도가 맞지 않는 데 기인한 최적화 왜곡을 유발하는 추가 손실으로 분해함으로써 지식 증류를 재고한다. 이를 해결하기 위해 비대상 분포 지식을 위한 분산 손실과 교사의 대상 출력을 부드러운 레이블로 사용하는 소프트 손실을 제안하여, CIFAR-100과 ImageNet에서 최신 기준 성능을 달성하는 새로운 손실(NKD)을 도출한다. 교사 없는 변종(tf-NKD)은 학생 예측을 스무딩하여 성능을 더욱 향상시키며, ResNet-18을 사용할 경우 ImageNet에서 Top-1 정확도를 +0.86% 향상시킨다.

ABSTRACT

Knowledge Distillation (KD) has developed extensively and boosted various tasks. The classical KD method adds the KD loss to the original cross-entropy (CE) loss. We try to decompose the KD loss to explore its relation with the CE loss. Surprisingly, we find it can be regarded as a combination of the CE loss and an extra loss which has the identical form as the CE loss. However, we notice the extra loss forces the student's relative probability to learn the teacher's absolute probability. Moreover, the sum of the two probabilities is different, making it hard to optimize. To address this issue, we revise the formulation and propose a distributed loss. In addition, we utilize teachers' target output as the soft target, proposing the soft loss. Combining the soft loss and the distributed loss, we propose a new KD loss (NKD). Furthermore, we smooth students' target output to treat it as the soft target for training without teachers and propose a teacher-free new KD loss (tf-NKD). Our method achieves state-of-the-art performance on CIFAR-100 and ImageNet. For example, with ResNet-34 as the teacher, we boost the ImageNet Top-1 accuracy of ResNet18 from 69.90% to 71.96%. In training without teachers, MobileNet, ResNet-18 and SwinTransformer-Tiny achieve 70.04%, 70.76%, and 81.48%, which are 0.83%, 0.86%, and 0.30% higher than the baseline, respectively. The code is available at https://github.com/yzd-v/cls_KD.

연구 동기 및 목표

• 고전적 지식 증류(KD) 손실과 원래 교차 엔트로피(CE) 손실 간의 기본적인 관계 분석
• 학생의 상대적 확률과 교사의 절대적 확률 간 척도 불일치로 인해 발생하는 KD 최적화 과제 규명
• 비대상 분포 지식 전달 최적화를 분리하고 향상시키기 위한 새로운 손실 형식 개발
• 교사의 대상 출력을 부드러운 레이블로 사용하여 증류 성능 향상
• 사전 학습된 교사 없이도 학생 예측을 스무딩하여 안정적인 소프트 타겟을 생성함으로써 교사 없는 지식 증류 가능화

제안 방법

• 고전적 KD 손실을 원래 CE 손실과 학생이 비대상 클래스에 대해 교사의 절대 확률을 일치시키도록 유도하는 추가 손실로 분해
• 확률 척도를 정규화함으로써 비대상 분포 지식 전달을 더 효과적으로 수행하기 위한 분산 손실 제안
• 교사의 대상 출력을 소프트 레이블로 사용하여 하드 레이블보다 더 부드러운 지도를 제공하는 소프트 손실 도입
• 분산 손실과 소프트 손실을 원래 CE 손실과 조합하여 성능 향상을 위한 새로운 지식 증류(NKD) 손실 구성
• 학습 동안 학생 자신의 타겟 출력을 스무딩하여 안정적인 소프트 타겟을 생성함으로써 교사 없는 변종(tf-NKD) 개발
• tf-NKD에서 학생 예측을 안정화하고 소프트 타겟 품질을 향상시키기 위해 특별한 스무딩 함수 $ S_t + V_t - \text{mean}(S_t) $ 적용

실험 결과

연구 질문

• RQ1 고전적 KD 손실이 최적화 역학 측면에서 원래 교차 엔트로피 손실과 어떤 관계가 있는가?
• RQ2 표준 KD 손실이 최적화가 비최적화되는 이유는 무엇이며, 학생과 교사의 확률 척도 불일치는 어떤 원인으로 발생하는가?
• RQ3 KD 손실을 재구성함으로써 비대상 분포 지식을 더 효과적으로 전달할 수 있는가?
• RQ4 교사의 대상 출력을 소프트 레이블로 사용하여 증류 성능을 향상시킬 수 있는가?
• RQ5 자신의 스무딩된 예측을 소프트 타겟으로 사용함으로써 교사 없이도 학생 모델을 효과적으로 훈련시킬 수 있는가?

주요 결과

• ResNet-34를 교사로 사용할 경우, 제안된 NKD 손실은 ResNet-18의 ImageNet Top-1 정확도를 69.90%에서 71.96%로 향상시킨다.
• 교사 없는 tf-NKD 방법은 ResNet-18을 사용할 경우 ImageNet에서 70.76%의 Top-1 정확도를 달성하여 기준 모델 대비 0.86% 향상된다.
• CIFAR-100에서 NKD 손실은 ResNet-18과 ResNet-34를 교사로 사용할 경우 각각 80.76%의 Top-1 정확도를 기록하며 이전 방법들을 능가한다.
• 스무딩 함수 $ S_t + V_t - \text{mean}(S_t) $ 를 사용한 tf-NKD 변종은 ImageNet에서 0.86%의 최대 향상률을 기록했으며, 교사의 출력을 소프트 타겟으로 사용하는 것조차 초월한다.
• 시각화 결과는 스무딩된 학생 타겟이 학습 에포크 전반에 걸쳐 진정한 클래스 소속 관계를 더 일관되게 반영함을 확인한다. 이는 레이블 안정성을 향상시킨다.
• 모든 가중치가 훈련 중 학생 예측에서 유도되므로, 기준 모델과 동일한 추가 훈련 시간이 필요로 한다.