Knowledge Distillation

Knowledge Distillation

Knowledge distillation은 2014 NIPS 2014 workshop에서 Geoffrey Hinton, Oriol Vinyals, Jeff Dean이 발표한 논문 "Distilling the Knowledge in a Neural Network"에서 소개된 개념이다.

 

Knowledge distillation은 Pre-trained 된 대용량의 Teacher Network의 정보를 비교적 작은 크기의 Student Network로 전이시키는 것이다.

 

Teacher-Student model

 

Teacher-Student Network

Student 모델은 Teacher Model을 더 적은 parameter를 가지고 모방하도록 학습된다.

 

Teacher Student model structure

 unlabeled data만을 사용하여 student 모델의 학습이 가능하기 때문에  unsupervised learning이다.

 

 

labeled data가 존재할 경우 이를 활용하여 모델을 다음과 같이 수정할 수 있다.

 

Supervised Teacher Student model structure

 

• Distillation Loss
  • KL Divergence 사용
  • Teacher model의 soft prediction과 Student model의 soft prediction 분포의 차이가 Loss
  • Teacher model의 지식을 모방하기 위한 Loss
• Student Loss
  • CrossEntropy 사용 (Hard-label인 G.T와 soft prediction)
  • Student model의 True label에 대한 Confidence의 음의 값이 Loss
  • True Label을 예측하기 위한 Loss
• Loss = Weighted sum of two losses

 

Softmax function with Temperature

 

Softmax에 Temperature를 사용하면 Prediction Probability가 높은 label은 상대적으로 Loss에 적은 영향을 미치고 Prediction Probability가 낮은 label의 정보가 Loss에 큰 영향을 미친다. 이는 0에 가까운 여러 prob의 분포 또한 효과적으로 학습하여 Teacher 모델을 학습하는데 도움이 된다.

 

 

Semi-supervised Learning

Teacher Student 모델 구조의 Knowledge distillation을 활용하여 Semi-supervised Learning이 가능하다.

Semi-supervised Learning은 labeling 되지 않은 수많은 데이터를 활용하기 위해 고안된 방법으로 Labeled 데이터와 Unlabeled 데이터를 모두 사용한다.

 

Pseudo-labeling Semi-supervised learning

 

Pseudo-labeling Semi-supervised learning은 Label data로 pre-trained 한 model을 사용하여 Unlabeled data의 label을 예측하고 Confidence score가 threshold를 넘어 의미가 있는 예측이라고 판단될 경우 예측값을 data의 label로 지정하고 이를 기존 labeled data에 추가하여 다시 model을 학습시키는 방법이다.

 

Noisy Student training

 

ImageNet Classification Top-Model

Noisy Student Training은 EfficientNet-B7을 넘어 2019년 ImageNet Classification 분야에서 SOTA를 기록한 모델이다.

Student, Teacher 모델은 기존의 SOTA 모델인 EfficientNet을 사용하였으며 Data Augmentation, Knowledge Distillation, Semi-Supervised learning (Pseudo label-based method)를 적용하여 큰 성능 향상을 보였다.

 

Self-Training with noisy student

1. Labeled data인 1M ImageNet data로 Teacher Model을 학습시킨다.
2. 300M unlabeled data에 pre-trained Teacher Model을 사용하여 Pseudo-label을 부여한다.
3. Labeled data에 Noise를 적용하고 Teacher model 이상의 크기로 생성한 Student Model을 학습시킨다.
4. 학습된 Student 모델이 새로운 Teacher Model이 되고 2~3의 단계를 반복한다.

논문에서 Noise를 부여하는 방법으로 Dropout, Stochastic depth, Data augmentation을 제시하였다.

Noise를 Student 모델을 학습시킬 때 적용하여 Teacher model을 뛰어넘는 성능을 가진 모델을 학습시킬 수 있었다.

Unlabeled data에 Noise를 주거나 Teacher Model에 Noise를 주는 것은 정확도에 악영향을 미칠 때가 있었다고 한다.

 

Stochastic Depth

Stochastic depth는 2016년에 Gao Huang et al. 의 논문 "Deep Networks with Stochastic Depth"에 등재된 개념으로 학습 도중 확률적으로 특정 layer에서 학습을 건너뛰는 identity mapping을 사용하는 method로 gradient vanishing / exploding 문제를 해결하기 위해 고안되었다.

 

Proposed Linear decay rule

위 식은 논문에서 제시한 stochastic depth를 적용하기 위한 수식 중 하나로 위의 수식을 사용하면 단순하게 depth에 따라 선형적으로 layer를 skip 할 확률이 변화한다. depth l이 0일 때 survival probability $p_{l}$은 1이 되고 depth l이 최대인 L일 때 survival probablity는 hyperparameter인 $p_{L}$이 된다.

 

Stochastic Depth Performance on CIFAR-10 Dataset

Reference

[1] Qizhe Xie et al. (2019). "Self-training with Noisy Student improves ImageNet classification". arXiv: 1911.04252 url

[2] Gao Huang. (2016). "Deep Networks with Stochastic Depth". arXiv:1603.09382 url

[3] Geoffrey Hinton, Oriol Vinyals, Jeff Dean. (2015). "Distilling the Knowledge in a Neural Network". arXiv:1503.02531 url

[4] Naver Boostcamp AI-Tech 오태현 교수님 "Annotation data efficient learning"