NLP 논문리뷰 - Distilling the Knowledge in a Neural Network

논문 제목 : Distilling the Knowledge in a Neural Network

논문 링크 : arxiv.org/abs/1503.02531

---

Background

NN의 오버피팅을 피하기 위해 앙상블 기법이 사용되었다. 하지만 앙상블은 여러 모델을 사용하여 계산시간이 많이 걸린다는 단점이 있어, 앙상블만큼의 성능과 '적은 파라미터 수'를 가진 nn모델이 필요하여 이 Knowledge Distillation모델이 나오게 되었다.

Knowledge Distillation이란?

Knowledge는 지식, Distillation은 증류이다. 화학에서 액체를 가열하여 생긴 기체를 냉각하여 다시 액체로 만드는 것을 증류라고 부르는데, 이러한 개념을 NN에서 사용한 것이다. 즉 위 그림과 같이 NN에서 지식 증류는 큰 모델(teacher network)로부터 증류한 지식을 작은 모델(student netwrok)로 transfer하는 과정이다.

---

How to Knowledge Distillation

Soft Label

label을 예측한다고 가정했을 때, 기존 one hot encoding은 [0, 1, 0, 0]과 같이 0아니면 1이라는 숫자로 출력하게 되어 값이 아닌 다른 값은 완전히 무시하게 된다. 하지만 이 다른 값도 어느 정도 영향을 받아야 된다고 생각하여 soft label이라는 개념이 나왔고, [0, 0, 1, 0] 대신 [0.1, 0.1, 0.7, 0.1]로 출력하여 정답이 아닌 다른 값에 작은 영향을 받는다. 이 논문에서는 T라는 파라미터를 추가했는데 아래에서 설명하겠다.

T를 넣은 이유는 불필요한 데이터를 걸러주고, 벡터의 분포를 더 부드럽게 만들기 위해서 넣습니다. T를 넣어 T가 낮으면 logit값이 큰 값에 대해서는 확률값이 1이 되고, 나머지 값에 대해서는 0의 값을 갖게 된다. 높으면 확률같이 다 비슷해지게 됨. 보통 T는 2~5 사용

 

distillation loss

distillation knowledge모델에서 loss를 구하는 방법이 가장 중요하다. loss를 구하는 방법은 큰 모델을 학습 시킨 후, 작은 모델을 위 그림과 같은 손실함수를 통해서 학습한다고 보면 된다.

soft label : lager model의 softmax output([0.1, 0.1, 0.7, 0.1])

hard label : 기존 one hot encoding방식 output ([0, 0, 1, 0])

정리하자면 Teacher Network 학습 -> Student Network 학습 -> Student Network soft Prediction + Teacher Network soft label을 통해 distillation loss 구성 -> Student Network hard prediction + Original hard label을 통해 studnet loss구성 -> Teacher 모델의 손실값과 Student 모델의 Cross entropy값을 더해 갱신해나가는 것이다.

 

단계를 정리하자면

1) Teacher Network : training set(x, hard label)을 사용해 large model 학습

2) large model 학습 뒤, large model의 output(soft label)을 target으로 하는 transfer set(x, soft label)을 생성. 이 때 soft label의 T는 1이 아닌 높은 값 사용.

3) Student Network :

transfer set을 사용해 small model 학습. T는 soft label을 생성할 때와 같은 값 사용. -> soft predictions

transfer set을 사용해 small model 학습. T는 1로 고정. -> hard predictions

4) loss 생성

distillation loss : soft label과 soft predictions의 차이를 Kullback-Leiber Divergence를 통해 구함.

student loss : hard predictions과 hard label을 Cross-entropy를 통해 구함.

※ Cross-entropy와 Kullback-Leiber Divergence의 차이 : 

Cross-entorpy는 실제 분포인 q를 모르고, 모델링을 하여 q분포를 예측하고자 하는 것이다. 훈련 데이터를 사용한 예측 모형에서 실제 값과 예측값의 차이를 계산할 때 사용되는데 아래 예시를 통해 자세히 설명하겠다.

예를 들어 가방에 실제 0.8/0.1/0.1의 비율로 빨간/녹색/노랑 공이 들어가 있다고 하고, 가방을 모르는 사람이 예측했을 때는 0.2/0.2/0.6의 비율로 들어가 있을 것이라고 생각했다. H(q)는 -(0.8log0.8 + 0.1log0.1 + 0.1log0.1) = 0.63이 되고, Hp(q)는 -(0.8log0.2 + 0.1log0.2 + 0.1log0.6) = 1.5가 된다.

KL Divergence라고 주로 부르는 이 개념을 서로 다른 두 분포의 차이를 측정하는 데 쓰인다. 이를 entropy 개념에 대입하면 두 entropy의 차이로 계산된다. 수식을 보면 Hp(q)는 실제 entropy인 H(q)보다 항상 크기 때문에 0보다 큰 값을 갖게 되고, 예측 분포인 p를 실제 분포 q에 가깝게 하는 것이 목표이다. p가 q에 가까워질수록 KL Divergence가 0에 가까워지고, 이 값을 최소화시키는 것이 목적이다.

5) 두 loss를 더해서 최종 loss를 구함.

 

---

Experiment

Mnist 데이터 셋 기준으로

large model : unit=1200, 2개의 hidden layer, FC model, drop out 포함, 60,000개의 trainsing set을 사용

small model : unit=800, 2개의 hidden layer, drop out 포함 x,

error의 개수는 large는 67개 smll(student)는 146개, distillied(transfer set 적용)는 74개가 나왔다.

최적의 파라미터 : unit=300, T>=8 / unit=30, 2.5<=T<=4

model의 크기가 작아질수록 T역시 함께 작아져야 좋은 성능을 보장한다는 것을 알 수 있다. 직관적으로 생각해보면 model의 크기가 크다는 것은 parameters가 많다는 것이고, 결국 smooth하지 않은 function이라는 것이다. 따라서 generalize가 아닌 overfitting될 가능성이 높다는 것인데, T는 output을 더 soft하게 만들어주기 때문에 overfitting을 방지하는 효과라고 이해할 수 있다. 따라서 model의 크기가 클수록 overfitting이 발생할 가능성이 높은 것이므로, 높은 T를 줘 강하게 soft 효과를 주는 것이 좋은 성능을 낼 것이라는 것을 유추할 수 있다.

 

 

speech recognition 기준으로

기준 모델은 58.9, 앙상블은 61.1, distilled는 60.8퍼센트가 나왔다.

 

soft targets as regularizers

soft target은 regularization 효과를 낼 수 있어, overfitting을 방지하는 효과를 가져다준다. 3%의 트레이닝 셋으로 soft targets을 적용했을 때 100% traing set의 baseline과 비슷한 성능을 보인 것을 볼 수 있다.