Knowledge distillation

Knowledge distillation은 NIPS 2014에서 제출된 Distilling the Knowledge in a Neural Network라는 논문에 제시된 개념입니다. 이 네트워크의 목적은 미리 잘 학습된 큰 네트워크(Teacher network)의 지식을 실제로 사용하고자 하는

작은 네트워크(Student network) 에게 전달하는 것입니다.

작은 네트워크로 큰 네트워크와 비슷한 성능을 낼 수 있도록, 학습과정에서 큰 네트워크의 지식을 작은 네트워크에게

전달하여 성능을 높이겠다는 것입니다.

 

Soft Label

- 일반적으로 이미지 클래스 분류와 같은 task는 마지막 layer의 softmax를 통해 각 클래스의 확률값을 output으로 합니다. 

곰, 고양이, 개 3가지 클래스를 구분하는 모델이 있을 때 위 그림의 왼쪽과 같은 결과는 Hard label, 오른쪽과 같다면 Soft label이라고 합니다. Soft label로 하면 고양이라는 걸 유추 할 수 있는데 고양이와 개가 함께 가지고 있는 특성들이 어느 정도 있기 떄문에 0.2라는 score가 나왔다고 생각 할 수 있는데 hard label로 표현하면 이러한 정보가 사라지는 거고 이러한 출력 값들이 모델의 지식이 될 수 있다고 하는 것입니다. 

 

출력 값의 분포를 soft label로 만들어주는 과정입니다. soft output을 dark knowledge라고 논문에서는 표현합니다.

기존의 softmax output과 다른 점은 T라는 값이 분모로 들어가는데 논문에서는 온도(temperature)라고 표현하고 이 값이 높아지면 더 soft하게 되는 거고 낮아지면 hard하게 만드는 과정입니다. 이 soft targe은 큰 모델(Teacher)의 지식을 의미합니다. 이 모델을 학습을 시킨 후 작은 모델(Student)을 다음과 같은 loss 함수를 통해 학습시킵니다.

Cross Entropy Loss (LCE)와 Distillation Loss (LKD)로 구성되어 있습니다. LKD는 잘 학습된 Teacher model의 soft labels와 Student model의 soft predictions를 비교하여 손실함수를 구성합니다. 이때, 온도(τ)는 동일하게 설정하고 Cross Entropy Loss를 사용합니다.

 

쉽게 말해 Student network의 분류 성능에 대한 loss로 ground truth와 Student의 분류 결과 차이를 Cross entropy loss로 계산하고 하나는 Student network와 Teacher network의 분류 결과의 차이를 loss에 포함시키는 것입니다.

 

결론적으로 Knowledge distillation 은 미리 학습시킨 Teacher network 의 출력을 내가 실제로 사용하고자 하는 작은 모델인 Student network 가 모방하여 학습함으로써, 상대적으로 적은 Parameter 를 가지고 있더라도 모델의 성능을 높이는 방법론입니다.

 

 

참고자료 : https://light-tree.tistory.com/196 ,

https://baeseongsu.github.io/posts/knowledge-distillation/