[논문 리뷰] Comparing Kullback-Leibler Divergence and Mean Squared Error Loss in Knowledge Distillation

Taehyeon Kim, Jaehoon Oh외 3명|arXiv (Cornell University)|2021. 05. 19.

Advanced Neural Network Applications참고 문헌 21인용 수 13

한 줄 요약

이 논문은 지식 정련에서 표준 Kullback-Leibler (KL) 발산 손실을 교사 및 학생 로짓 벡터 간의 평균 제곱 오차 (MSE) 손실로 대체하는 것을 제안하며, 여러 벤치마크에서 직접 로짓 매칭을 통한 MSE가 KL 발산보다 우수한 성능을 보임을 입증한다. 주요 기여는 MSE가 더 정확한 로짓 매칭을 가능하게 하여, 특히 학생 모델이 교사보다 작을 경우 일반화 성능을 향상시킨다는 점이다.

ABSTRACT

Knowledge distillation (KD), transferring knowledge from a cumbersome teacher model to a lightweight student model, has been investigated to design efficient neural architectures. Generally, the objective function of KD is the Kullback-Leibler (KL) divergence loss between the softened probability distributions of the teacher model and the student model with the temperature scaling hyperparameter tau. Despite its widespread use, few studies have discussed the influence of such softening on generalization. Here, we theoretically show that the KL divergence loss focuses on the logit matching when tau increases and the label matching when tau goes to 0 and empirically show that the logit matching is positively correlated to performance improvement in general. From this observation, we consider an intuitive KD loss function, the mean squared error (MSE) between the logit vectors, so that the student model can directly learn the logit of the teacher model. The MSE loss outperforms the KL divergence loss, explained by the difference in the penultimate layer representations between the two losses. Furthermore, we show that sequential distillation can improve performance and that KD, particularly when using the KL divergence loss with small tau, mitigates the label noise. The code to reproduce the experiments is publicly available online at https://github.com/jhoon-oh/kd_data/.

연구 동기 및 목표

• KL 발산 손실에서의 온도 스케일링 하이퍼파rameter τ가 지식 정련 성능에 미치는 영향을 조사하는 것.
• 부드러운 라벨이 지식 정련에서 일반화를 향상시키는 이유에 대한 이론적 및 실증적 이해의 부족을 해결하는 것.
• KL 발산의 대안으로 직접 로짓 매칭을 위한 MSE 손실을 사용하는 방법을 제안하여 학생 모델 성능을 향상시키는 것.
• 다양한 손실 함수가 노이즈가 있는 라벨 설정에서 얼마나 견고한지 평가하는 것.
• 교사와 학생 모델 간의 용량 격차가 클 경우 성능 향상을 위해 순차적 정련 전략을 탐색하는 것.

제안 방법

• 이론적 분석을 통해 τ → ∞일 때 KL 발산 손실은 로짓 매칭을 근사하고, τ → 0일 때는 라벨 매칭을 근사함을 보여줌.
• 학생 및 교사 로짓 벡터 간의 MSE 손실을 사용하는 직접 로짓 매칭 프레임워크를 제안하며, 부드러운 확률 분포를 거치지 않고 매칭을 수행함.
• Müller 등(2019)의 시각화 기법을 사용하여 KL 손실과 MSE 손실로 훈련된 모델의 최전단 레이어 표현을 비교함.
• 먼저 KL 손실(τ = 3)로 훈련한 후 MSE 손실로 미세조정하는 순차적 정련을 통해 용량 격차가 클 경우 성능 향상을 도모함.
• CIFAR-100에서 라벨을 무작위로 뒤집어 노이즈 라벨 설정을 만들고, 다양한 손실 함수 하에서 테스트 정확도를 측정하여 노이즈에 대한 견고성 평가.
• GitHub에서 공개된 코드를 사용하여 실험을 재현함으로써 재현 가능성을 확보함.

실험 결과

연구 질문

• RQ1 KL 발산 손실에서 온도 스케일링 하이퍼파rameter τ는 지식 정련의 일반화 행동에 어떻게 영향을 미치는가?
• RQ2 MSE 손실을 통한 직접 로짓 매칭이 지식 정련에서 KL 발산 손실을 초월할 수 있는가?
• RQ3 KL 발산과 MSE 손실로 훈련된 모델 간의 최전단 레이어 표현 학습 방식의 차이는 무엇인가?
• RQ4 순차적 정련(먼저 KL 손실, 그 후 MSE)은 교사와 학생 모델 간 용량 격차가 클 경우 성능 향상에 기여하는가?
• RQ5 노이즈가 있는 라벨 설정에서 다양한 손실 함수는 어떻게 작동하는가, 특히 잘못된 라벨에 대한 과적합을 억제하는 데서 성능은 어떠한가?

주요 결과

• 모든 평가된 데이터셋에서 MSE 손실이 KL 발산 손실을 능가하며, CIFAR-100에서 75.78%의 테스트 정확도를 기록한 반면 KL 손실(τ=20)은 75.47%를 기록함.
• 동일한 교사 및 학생 아키텍처를 사용할 경우, MSE 손실은 τ=20를 포함한 모든 KL 발산 설정보다 높은 성능을 달성함.
• KL 발산 손실(τ=20)은 75.47%의 정확도를 기록한 반면, MSE 손실은 75.78%를 기록하여 직접 로짓 매칭이 부드러운 라벨 매칭보다 더 효과적임을 시사함.
• 순차적 정련(KL 이후 MSE)은 용량 격차가 클 경우 성능 향상에 기여하며, 단독으로 어떤 손실 함수를 사용하는 것보다 우수한 성능을 보임.
• 40%의 대칭 노이즈가 있는 라벨 설정에서 τ=0.5인 KL 손실이 MSE 및 고온도 KL보다 우수한 일반화 성능을 보였으며, 이는 암묵적인 라벨 스무딩 효과 때문임.
• MSE 손실로 훈련된 모델은 KL 손실로 훈련된 모델보다 교사의 최전단 레이어 표현을 더 정확하게 유지함. 반면 KL 손실은 표현을 연장시켜 왜곡시킴.