Logit Standardization in Knowledge Distillation

Knowledge distillation involves transferring soft labels from a teacher to a student using a shared temperature-based softmax function. However, the assumption of a shared temperature between teacher and student implies a mandatory exact match between their logits in terms of logit range and variance. This side-effect limits the performance of student, considering the capacity discrepancy between them and the finding that the innate logit relations of teacher are sufficient for student to learn. To address this issue, we propose setting the temperature as the weighted standard deviation of logit and performing a plug-and-play Z-score pre-process of logit standardization before applying softmax and Kullback-Leibler divergence. Our pre-process enables student to focus on essential logit relationsfrom teacher rather than requiring a magnitude match, and can improve the performance of existing logit-based distillation methods. We also show a typical case where the conventional setting of sharing temperature between teacher and student cannot reliably yield the authentic dis-tillation evaluation; nonetheless, this challenge is success-fully alleviated by our Z-score. We extensively evaluate our method for various student and teacher models on CIFAR-100 and ImageNet, showing its significant superiority. The vanilla knowledge distillation powered by our pre-process can achieve favorable performance against state-of-the-art methods, and other distillation variants can obtain considerable gain with the assistance of our pre-process. The codes, pre-trained models and logs are released on Github.

한 줄 요약

로짓 표준화를 통해 지식 증류 성능을 향상시키는 새로운 전처리 방법을 제안한다.

핵심 기여도

기존 지식 증류에서 로짓의 범위와 분산 일치를 강요하는 온도 공유 문제를 이론적으로 분석.
로짓 표준편차를 기반으로 한 가변 온도 설정과 Z-score 전처리를 제안.
CIFAR-100과 ImageNet 데이터셋에서 기존 로짓 기반 증류 방법 대비 우수한 성능 보임.
기존 증류 방법(KD, CTKD, DKD, MLKD)에도 적용 가능하며 성능 향상 효과 있음.

핵심 아이디어

기존 지식 증류(KD)는 교사와 학습자의 로짓에 동일한 온도를 적용하는 방식을 사용하지만, 이는 로짓의 범위와 분산이 일치해야 한다는 암묵적 가정을 내포한다. 이는 특히 학습자의 용량이 작을 경우 성능 저하를 초래한다. 본 연구는 정보 이론의 엔트로피 최대화 원리를 기반으로, 온도가 라그랑주 승수에서 유도된다는 점을 이론적으로 증명하고, 이에 따라 교사와 학습자, 샘플 간에 온도를 다르게 설정할 수 있음을 보인다. 이를 바탕으로 로짓 표준편차를 기반으로 한 가변 온도와 Z-score 전처리를 제안하여, 학습자가 로짓의 절대 크기보다는 상대적 관계에 집중할 수 있도록 유도한다.

기술적 접근법

**온도 설정**: 로짓의 가중 표준편차를 기반으로 온도를 설정.
**전처리**: 로짓에 Z-score 표준화를 적용한 후 softmax와 KL 발산을 수행.
**데이터셋**: CIFAR-100 (50,000개 학습, 10,000개 검증), ImageNet (1.28M 학습, 50,000 검증).
**하이퍼파라미터**: SGD 최적화, 학습률 0.01~0.05, 에포크 수 240~480.
**비교 방법**: KD, CTKD, DKD, MLKD, FitNet, RKD, CRD, OFD, ReviewKD, SimKD, CAT-KD.

주요 결과

**CIFAR-100**: 기존 로짓 기반 증류 방법 대비 평균 +2.1% 성능 향상.
**ImageNet**: 기존 로짓 기반 증류 방법 대비 평균 +1.6% 성능 향상.
**MLKD**: Z-score 전처리 적용 시 +1.8% 성능 향상.
**KD**: Z-score 전처리 적용 시 +2.3% 성능 향상.

의의 및 한계

**의의**: 로짓 기반 지식 증류의 이론적 기반을 재정립하고, 기존 방법의 한계를 보완.
**실용성**: 기존 로짓 기반 증류 방법에 plug-and-play 방식으로 적용 가능.
**한계**: Z-score 전처리는 로짓의 절대 크기 정보를 제거하므로, 일부 특수한 로짓 기반 방법에는 부적합할 수 있음.
**이론적 한계**: 온도의 가변성은 학습 안정성에 영향을 줄 수 있으나, 본 연구에서는 이를 실험적으로 보완함.

실용적 활용

모바일 및 임베디드 장치에서 사용하는 가벼운 모델 학습에 적용 가능.
대규모 모델의 압축 및 효율적 추론을 필요로 하는 산업 분야(예: 자율주행, 의료 영상 분석)에 활용 가능.
다양한 증류 방법과 결합하여 성능 향상을 도모할 수 있음.