[논문 리뷰] Improving Knowledge Distillation via Regularizing Feature Norm and Direction

Yuzhu Wang, Lechao Cheng외 4명|arXiv (Cornell University)|2023. 05. 26.

COVID-19 diagnosis using AI인용 수 5

한 줄 요약

이 논문은 특성 노름과 방향을 함께 정규화함으로써 학생 모델 성능을 햖스르는 새로운 지식 증류 방법을 제안한다: 이는 학생 모델이 교사의 클래스 평균 특성과 일치하는 고노름 특성을 생성하도록 유도한다. 핵심 기여는 특성 크기와 정렬을 동시에 향상시키는 단순하면서도 효과적인 ND 손실이며, 기존의 KD 방법과 결합할 경우 ImageNet, CIFAR-100, COCO 벤치마크에서 최신 기술 수준(SOTA)의 정확도를 달성한다.

ABSTRACT

Knowledge distillation (KD) exploits a large well-trained model (i.e., teacher) to train a small student model on the same dataset for the same task. Treating teacher features as knowledge, prevailing methods of knowledge distillation train student by aligning its features with the teacher's, e.g., by minimizing the KL-divergence between their logits or L2 distance between their intermediate features. While it is natural to believe that better alignment of student features to the teacher better distills teacher knowledge, simply forcing this alignment does not directly contribute to the student's performance, e.g., classification accuracy. In this work, we propose to align student features with class-mean of teacher features, where class-mean naturally serves as a strong classifier. To this end, we explore baseline techniques such as adopting the cosine distance based loss to encourage the similarity between student features and their corresponding class-means of the teacher. Moreover, we train the student to produce large-norm features, inspired by other lines of work (e.g., model pruning and domain adaptation), which find the large-norm features to be more significant. Finally, we propose a rather simple loss term (dubbed ND loss) to simultaneously (1) encourage student to produce large-\emph{norm} features, and (2) align the \emph{direction} of student features and teacher class-means. Experiments on standard benchmarks demonstrate that our explored techniques help existing KD methods achieve better performance, i.e., higher classification accuracy on ImageNet and CIFAR100 datasets, and higher detection precision on COCO dataset. Importantly, our proposed ND loss helps the most, leading to the state-of-the-art performance on these benchmarks. The source code is available at \url{https://github.com/WangYZ1608/Knowledge-Distillation-via-ND}.

연구 동기 및 목표

• 기존 지식 증류 방법이 학생 특성과 교사 특성을 직접 정렬하지만 특성 크기나 분류기 구조를 고려하지 않는 한계를 해결하기 위해.
• 교사의 클래스 평균 특성을 강력한 내재 분류기로 활용하여 학생의 특성 학습을 유도함으로써 학생 성능을 향상시키기 위해.
• 모델 프루닝과 도메인 어댑테이션 연구에서 밝혀진 대규모 노름 특성의 역할을 탐구하기 위해.
• 특성 노름과 방향을 동시에 정규화하는 단순하고 즉각 적용 가능한 손실(ND 손실)을 개발하기 위해.

제안 방법

• 교사 모델의 사전 구축된 특성에서 클래스 평균 특성을 생성하고, 이를 학생 특성 학습을 유도하는 강력한 분류기로 활용한다.
• 새로운 ND 손실을 도입하여 두 목표를 동시에 최적화한다: (1) 학생 특성의 L2 노름을 최대화하고, (2) 학생 특성과 해당 교사 클래스 평균 특성 간의 코사인 거리를 최소화한다.
• ND 손실은 특성이 로그릿을 계산하는 데 사용되는 임베딩 레이어(마지막에서 두 번째 레이어)에 적용되며, 분류 성능와 직접적으로 관련된다.
• 로짓 증류 및 특성 증류 프레임워크 양쪽 모두와 호환되며, 기존의 KD 파ipeline에 원활하게 통합될 수 있다.
• ImageNet, CIFAR-100, COCO 데이터셋에서 실험적으로 검증되었으며, 아블레이션 스터디를 통해 각 구성 요소의 효과가 확인되었다.
• 구현 코드는 KD++로 공개되었으며, 최소한의 계산 오버헤드로도 뚜렷한 정확도 향상을 보였다.

Figure 1 : Our main contribution is a simple loss, termed $\mathcal{L}_{nd}$ , that regularizes the n orm and d irection of the student features. $\mathcal{L}_{nd}$ is applicable to different KD methods which are categorized into two types in the context of classification: (left) logit distillation

실험 결과

연구 질문

• RQ1 학생 특성을 더 큰 노름과 교사 클래스 평균과의 일치를 유도하도록 정규화하면 지식 증류 성능이 향상되는가?
• RQ2 단순하고 통합된 손실 함수가 지식 증류에서 노름 최대화와 방향 정렬을 효과적으로 동시에 구현할 수 있는가?
• RQ3 기본 정규화 기법(예: 코사인 손실, 노름 정규화)과 비교해 ND 손실은 성능 향상에서 어떤가?
• RQ4 교사 모델의 용량이 증가함에 따라 ND 손실의 이점이 비례 증가하는가?
• RQ5 아키텍처의 차이가 있음에도 불구하고, Transformer에서 CNN으로 증류할 때 ND 손실이 성능 향상에 기여하는가?

주요 결과

• 표준 KD에 ND 손실을 적용한 KD++는 ResNet-18 학생 모델과 ResNet-50 교사 모델을 사용해 ImageNet에서 72.53%의 Top-1 정확도를 달성했으며, 원래의 KD(71.35%)와 비교해 뛰어난 성능을 보였다. 이는 ReviewKD(71.10%) 및 DKD(71.87%)와도 비교해 우월하다.
• ResNet-18 학생 모델을 ImageNet에서 훈련할 때 KD에 비해 1.45%p의 절대 정확도 향상을 기록했으며, 이는 ND 손실의 뛰어난 효과를 입증한다.
• 더 큰 교사 모델(ResNet-34에서 ResNet-152)을 사용할 경우 KD++는 일관된 성능 향상을 보였으며(예: 71.98% → 72.54%), 반면 기존의 KD 방법은 성능 저하 또는 변동을 보였다.
• COCO 객체 검출에서 ND 손실을 사용하면 더 높은 검출 정밀도를 기록했으며, 이는 이미지 분류를 넘어서도 유의미한 이점이 있음을 확인한다.
• 시각화 결과, KD++는 더 높은 노름 특성과 더 나은 클래스 간 분리도를 생성했으며, 기존의 표준 KD 및 단순 학생 모델이 놓친 클래스들(예:_PURPLE 클래스)을 성공적으로 복구했다.
• ND 손실은 교차 엔트로피나 KL 손실과 함께 단독으로 사용되어도 정확도를 약 1%p 향상시키며, 계산 비용은 거의 증가하지 않아 효과적이다.

(a) teacher (Res56)