본문 바로가기

Papers/ML

[리뷰] Distilling the Knowledge in a Neural Network (NIPS Workshop 2014)

by Hinton et al.

Paper Link : https://arxiv.org/pdf/1503.02531.pdf

Contents

1. Introduction
2. Distillation
3. Experiments (MNIST, Speech Recognition)
4. Training Ensembles on Very Big Datasets (JFT)
5. Soft Targets as Regularizers
6. Summary

 


Introduction


다수의 사용자가 사용하는 서비스를 배포하고자 할 때, latency와 computational resources는 중요한 문제입니다. 예를 들어, ensemble 모델은 모바일 프로그램으로 배포하기에 부적합할 수 있습니다. 이러한 경우 도입해볼 수 있는 방법론이 "Distillation"입니다. 

Distillation은 ensemble과 같은 cumbersome 모델의 지식을 small 모델로 전달(knowledge transer)하는 기법입니다. 위 맥락에서 cumbersome 모델을 teacher 모델, small 모델을 student 모델이라고 하겠습니다.

모델의 일반화 능력(generalization ability)을 전달하는 방법 중 하나는 cumbersome의 output인 class probability를 "soft target"으로 사용하는 것입니다.

여기서 soft target/label은 hard target/label에 반대되는 개념으로 다음과 같이 설명할 수 있습니다.

Soft target : class probabilities produced by the cumbersome model
- Teacher 모델의 output에 softmax를 취하여 얻은 각 class에 대한 확률 값. 예를 들어, (개, 고양이, 사람)에 대한 확률이 (0.5, 0.2, 0.3)으로 얻어진 경우.
Hard target : binary class from class probabilties (e.g. one-hot encoding)
- Class probability에서 가장 높은 값을 찾아(argmax) 해당 example이 특정 클래스에 속하는지에 대한 이진 분류. 예를 들어, (개, 고양이, 사람)에 대한 결과가 (1, 0, 0)으로 얻어진 경우.

 

* cumbersome: large or heavy and therefore difficult to carry or use (google dicionary)

 

Distillation


Converting Logit to Probability

$$q_i = \frac{exp(z_i/T)}{\sum_{j} exp(z_i/T)}$$

$z$를 logit이라고 하면 temperature $T$로 나눈 logit에 softmax를 취함으로써 cumbersome 모델의 class probability $q$를 얻을 수 있습니다.

$T$는 보통 1로 설정되며, T가 높아질수록 class probability의 distribution이 smoothing되는 효과를 얻을 수 있습니다.

 

Two Objective functions

$$\mathcal{L}(x;W) = (1-\alpha) * \mathcal{L}_{CE}(y, \sigma(z_s; T=1)) + \alpha * \mathcal{L}_{CE}(\sigma(z_t; T=\tau), \sigma(z_s, T=\tau))$$

  1. 첫 번째 항: Cross entropy with correct labels (T=1)
  2. 두 번째 항: Cross entropy with soft targets (high T)

$\alpha$를 높여 첫 번째 항에 적은 weight를 주는 것이 좋은 결과를 보였다고 합니다.

 

Matching logits

Cross-entropy loss의 logit에 대한 derivative $dC/dz_i$를 계산해봅시다. $z_i$는 student (distilled) 모델의 logit, $v_i$는 teacher (cumbersome) 모델의 logit입니다. $T$가 logit의 magnitude에 비해 클 때, gradient는 다음과 같이 근사됩니다.

 

$${\partial C \over \partial z_i} = {1 \over T} (q_i - p_i) \approx {1 \over {NT^2}} (z_i - v_i)$$

 

결국 distillation은 $1/2(z_i - v_i)^2$를 최소화하는 것이 되고, teacher와 student의 logit을 같게 만드는 것($z_i = v_i$)과 같다고 할 수 있습니다.

 

출처: https://arxiv.org/abs/1611.01144

높은 temperature(e.g. $\tau=10.0$)를 사용하면, distilled 모델이 large negative logit을 고려하고 logit은 쓸모 있는 정보를 전달할 것입니다.

낮은 temperature(e.g. $\tau=0.1$)를 사용하면, large negative logit의 영향이 적어져 cumbersome 모델의 logit이 noisy할 때 사용할 수 있습니다.

증류된(distilled) 모델이 너무 작아서 지식을 담기 어려울 땐, 중간 정도의 temperature를 사용하면 좋은 성능을 보였습니다. 이는 large negative logit을 무시하는 게 성능을 향상하는 데 도움이 되기도 한다는 것을 보여줍니다.

 

 

Training Ensembles on Very Big Datasets


Large neural networks와 large dataset을 사용한 ensemble의 경우, 학습 시 computation 양이 엄청날 것입니다. 여기서 class의 부분 집합만 사용해 학습한 specialist 모델을 이용한 ensemble이 전체 training set을 이용해 학습한 generalist 모델을 이용했을 때에 비해 어떻게 개선될 수 있는지 보여줍니다.

이 문제를 보여주기 위해 구글의 비공개 데이터셋인 JFT(100M images with 15k labels)를 실험에 이용했습니다.

출처: https://www.cs.ubc.ca/~lsigal/532S_2018W2/4b.pdf

이러한 specialist 모델은 쉽게 과적합(overfitting)할 수 있습니다. 이를 해결하기 위해 두 가지 방법을 제안합니다.

  1. 각각의 specialist 모델을 generalist 모델의 weight으로 초기화하여, specialist를 lower level feature dector로 사용합니다.
  2. Soft target을 사용합니다.

Table 3의 실험 결과에 따르면, specialist 모델을 이용한 ensemble은 학습 시간을 단축하여 distillation할 수 있고, test accuracy도 1.1%가량 향상되었습니다. Table 4에 따르면, specialist 모델이 더 많은 class를 가진 subset에서 학습될 때 성능이 더욱 향상되었습니다. Specialist 모델은 각각 독립적이기 때문에 병렬화하여(parallelize) 학습하기 쉽습니다.

 

Soft Targets as Regularizers


Hard target 대신 soft target을 사용하면 모델이 generalization을 잘하도록 만들어 줍니다. 또한, soft target을 사용하면 overfitting을 방지해주고 early stopping을 필요로 하지 않는다는 장점이 있습니다.

 

Summary


  • 모바일 기기 등으로의 모델 배포에서 어려움을 겪을 수 있는 ensemble, 또는 매우 큰 regularized 모델을 보다 작은 모델로 증류(distillation)하여 knowledge를 전이할 수 있음을 보였습니다.
  • Distillation에서 사용하는 transfer set은 일반적인 supervised learning에서 쓰는 training set과 달리 unlabeled data를 포함할 수 있습니다.
  • 굉장히 큰 neural network의 경우엔 full ensemble을 학습하는 게 어렵지만, small specialist 모델 여러 개를 ensemble함으로써 성능 향상을 얻을 수 있습니다.

 

※ 위 내용은 스터디 및 개인 학습 용으로 작성된 글로, 일부 정확하지 않은 내용을 포함할 수도 있습니다.

(피드백 환영합니다!)


Q&A

Q. Is the transfer set same with the training set?

A. (2 Distillation) 윗 문단을 보면 transfer set은 완전히 unlabeled data 또는 original training set을 쓸 수 있다. supervised loss와 unsupervised loss("soft target")를 결합해서 사용하면 성능이 잘 나온다.

 

Q. Why does setting the temperature high help distilled models to learn the cumbersome model? (Section 2.1)

A. High temperature를 사용할 경우 distribution이 smoothing되고 large negative logit을 고려하게 되면서 더 많은 정보를 student model에 전달할 수 있다.


Q. Why should we multiply $T^2$ to the gradient? (Section 2.1)

A. 식 4에서 cross-entropy loss에 대한 gradient가 $1/T^2$로 스케일된 것을 볼 수 있다. 매 실험마다 다른 temperature를 쓴다고 해도, $T^2$를 곱해 스케일을 복구함으로써 soft target의 상대적인 영향을 비슷한 수준으로 만들 수 있다.

 

Q. What is the difference between generalist vs. specialist model? (Section 5.2)

- Generalist: 전체 training data에 대해 학습된 모델
- Specialist: 일부 class에 대해 학습되어 작은 문제에 특화된 모델들

(Ref) UPC Presentation https://www.cs.ubc.ca/~lsigal/532S_2018W2/4b.pdf

반응형