[논문 리뷰] Efficient Knowledge Distillation from Model Checkpoints

Chaofei Wang, Qisen Yang외 3명|arXiv (Cornell University)|2022. 10. 12.

Machine Learning and Data Classification인용 수 7

한 줄 요약

이 논문은 지식 정복에서 최종 수렴 모델 대신 중간 단계의 모델 체크포인트를 사용하여 더 나은 교사 모델로 활용하는 것을 제안한다. 정보 버블 이론을 바탕으로 중간 단계의 모델은 더 높은 입력-특징 상호정보를 유지함으로써 더 많은 '다크 지식'을 보존함을 밝혀내었으며, 이는 학생 모델의 성능 향상으로 이어진다. 실험 결과, 최적의 중간 단계 교사 모델은 최종 모델과 강력한 앙상블 모두를 능가하며, 제안된 선택 알고리즘이 CIFAR-100에서 정확도를 최대 0.8% 향상시킨다.

ABSTRACT

Knowledge distillation is an effective approach to learn compact models (students) with the supervision of large and strong models (teachers). As empirically there exists a strong correlation between the performance of teacher and student models, it is commonly believed that a high performing teacher is preferred. Consequently, practitioners tend to use a well trained network or an ensemble of them as the teacher. In this paper, we make an intriguing observation that an intermediate model, i.e., a checkpoint in the middle of the training procedure, often serves as a better teacher compared to the fully converged model, although the former has much lower accuracy. More surprisingly, a weak snapshot ensemble of several intermediate models from a same training trajectory can outperform a strong ensemble of independently trained and fully converged models, when they are used as teachers. We show that this phenomenon can be partially explained by the information bottleneck principle: the feature representations of intermediate models can have higher mutual information regarding the input, and thus contain more "dark knowledge" for effective distillation. We further propose an optimal intermediate teacher selection algorithm based on maximizing the total task-related mutual information. Experiments verify its effectiveness and applicability.

연구 동기 및 목표

• 지식 정복에서 높은 성능을 보이는 완전히 훈련된 교사 모델이 학생 모델을 최고로 만든다는 기존의 가정을 도전하기 위해.
• 낮은 정확도를 보임에도 불구하고 중간 단계 모델이 완전히 수렴한 모델보다 더 나은 교사로 기능하는 이유를 조사하기 위해.
• 정보 이론에 기반한 이론적이고 효율적인 방법으로 최적의 중간 단계 교사 모델을 선택하는 데에.
• 단일 훈련 경로에서의 약한 스냅샷 앙상블이 독립적으로 훈련된 강력한 앙상블보다 더 나은 성능을 내는지 입증하기 위해.
• 전체 훈련 비용을 줄이며 전략적 중간 단계 모델 선택을 통해 지식 정복 성능을 향상시키기 위해.

제안 방법

• 비모수적 엔트로피 추정기를 사용하여 각 훈련 체크포인트에서 입력과 특징 간 상호정보 I(X;F)와 타깃과 특징 간 상호정보 I(Y;F)를 추정한다.
• I(X;F)와 I(Y;F)의 합을 최대화함으로써 지식의 풍부함과 과제의 관련성 간 균형을 이루는 최적의 중간 단계 교사 선택 기준을 수립한다.
• 다중 독립적 교사 훈련 런이 필요 없이 단일 훈련 경로 상의 최적 체크포인트를 식별하는 알고리즘을 설계한다.
• 표준 KD 손실 함수와 호환되며, 어떤 학생-교사 아키텍처에도 적용 가능하다.
• 체크포인트 간 비교를 위해 정규화된 상호정보 점수를 사용하여 총 과제 관련 정보를 최대화하는 체크포인트를 선택한다.
• 비용이 많이 드는 완전 훈련 모델이나 앙상블 대신 최적의 시간에 선택된 단일 중간 단계 체크포인트로 교체함으로써 효율적인 정복을 가능하게 한다.

실험 결과

연구 질문

• RQ1 낮은 정확도를 보임에도 불구하고 중간 단계 모델 체크포인트가 완전히 수렴한 모델보다 지식 정복에서 더 뛰어난 성능을 보이는 이유는 무엇인가?
• RQ2 단일 훈련 경로에서의 약한 스냅샷 앙상블이 독립적으로 훈련된 강력한 앙상블보다 더 나은 성능을 낼 수 있는가?
• RQ3 정보 버블 원리가 중간 단계 모델의 뛰어난 지식 전이 능력을 어떻게 설명하는가?
• RQ4 지식 정복 성능을 최대화하기 위해 최적의 중간 단계 교사 모델을 선택하는 전략은 무엇인가?
• RQ5 비용이 많이 드는 앙상블을 대체할 수 있는 단일 최적 시간 체크포인트가 학생 정확도를 유지하거나 향상시킬 수 있는가?

주요 결과

• 상호정보 최대화를 통해 선별된 최적의 중간 단계 교사 모델은 WRN-40-1 학생 모델을 사용해 CIFAR-100에서 73.26%의 KD 정확도를 달성하였으며, 완전 수렴 모델의 72.68%를 0.58% 뛰어넘었다.
• ResNet-110의 경우 최적의 중간 단계 교사 모델이 72.63%의 KD 정확도를 기록하여 전체 모델의 72.48%와 0.7x 체크포인트의 72.46%를 모두 초월했다.
• 반정도 체크포인트(0.5x)는 평가된 모든 아키텍처에서 완전 수렴 모델을 일관되게 능가했으며, CIFAR-100에서 평균 0.31% 향상된 성능을 보였다.
• 단일 훈련 런에서의 약한 스냅샷 앙상블이 독립적으로 훈련된 강력한 앙상블보다 더 뛰어난 성능을 보여, 제안된 선택 전략의 유효성을 입증했다.
• 다중 완전 훈련 모델이나 앙상블을 훈련할 필요를 제거함으로써 훈련 비용을 줄였으며, 동시에 뛰어난 지식 정복 성능을 달성했다.
• 연구는 과도한 훈련이 클래스 상관관계 정보를 억압함으로써 정복에 사용 가능한 '다크 지식'의 품질을 떨어뜨린다는 점을 확인했다.