Knowledge Distillation

Knowledge Distillation = 지식(Knowledge) + 증류(Distillation)

- 지식 : 어떤 대상에 대하여 배우거나 실천을 통하여 알게 된 명확한 인식이나 이해
- 증류 : 액체를 가열하여 생긴 기체를 냉각하여 다시 액체로 만드는 일. 여러 성분이 섞인 혼합 용액으로부터 끓는점의 차이를 이용하여 각 성분을 분리할 수 있다.

그렇다면 지식 증류란 무엇인가하면, 큰 모델(Teacher Network)로부터 증류한 지식을 작은 모델(Student Network)로 transfer하는 일련의 과정이라고 할 수 있다.

Knowledge Distillation은 도대체 왜 등장했을까?

지식 증류를 처음으로 소개한 논문은 모델 배포(model deployment) 측면에서 지식 증류의 필요성을 찾고 있다. 우리가 딥러닝을 활용해 인공지능 예측 서비스를 만든다고 가정해봅시다. 연구 및 개발을 통해 만들어진 딥러닝 모델은 다량의 데이터와 복잡한 모델을 구성하여 최고의 정확도를 내도록 설계되었을 것이다. 하지만 모델을 실제 서비스로 배포한다고 생각했을 때, 이 복잡한 모델은 사용자들에게 적합하지 않을 수 있다. 아래 슬라이드처럼 모델이 배포된 모바일 장치는 (복잡한 모델이 작동하는데 필요한) 강력한 하드웨어가 아니기 때문이다.

간단한 예를 들어보자. 다음의 두 모델이 있다면 어떤 모델을 사용하는 게 적합할까?

- 복잡한 모델 T : 예측 정확도 99% + 예측 소요 시간 3시간
- 단순한 모델 S : 예측 정확도 90% + 예측 소요 시간 3분

어떤 서비스냐에 따라 다를 수 있겠지만, 배포 관점에서는 단순한 모델 S가 조금 더 적합한 것으로 보인다. 더 빠른 prediction이 가능하기 때문에 더 빠른 서비스가 가능하기 때문이다.

여기서, 복잡한 모델과 단순한 모델을 잘 활용해서 더 나은 서비스를 제공하기 위해서 고안된 방법이 바로 Knowledge distillation이다. 바로 복잡한 모델이 학습한 generalization 능력을 단순한 모델 S에 전달(transfer)해주는 것이다.

위 슬라이드는 복잡한 모델 T와 단순한 모델 S가 가지는 특징에 관해 설명하고 있다.
이 개념을 처음으로 제시한 논문에서는 복잡한 모델을 cumbersome model로, 단순한 모델을 simple model로 나누어 설명하고 있다. 하지만 이후 등장하는 논문에서는 일반적으로 Teacher model과 Student model로 표현하고 있다. 먼저 배워서 나중에 지식을 전파해주는 과정이 선생님과 학생의 관계와 비슷하여 이렇게 표현한 듯 하다.

Knowledge Distillation은 언제 처음으로 등장했을까?

Knowledge Distillation은 NIPS 2014 workshop에서 발표한 논문 “Distilling the Knowledge in a Neural Network”에서 처음으로 등장한 개념이다. 논문을 요약하자면 다음과 같다.

- 앙상블과 같은 복잡한 모델을 다량의 유저에게 배포하는 것은 하드웨어적으로 엄청 힘듦
- 앙상블이 가진 지식을 단일 모델로 전달해주는 기존 기법이 있으나 그보다 더 일반적인 연구를 함
- MNIST 데이터를 사용해, 큰 모델로부터 증류된 지식이 작은 모델로 잘 전달되는지 확인함
- 앙상블 학습 시간을 단축하는 새로운 방법의 앙상블을 제시함

Knowledge Distillation은 어떻게 하는 걸까?

추상적인 개념에 대해 주야장천 설명은 들었는데, 구체적인 방법론은 어떻게 하면 될까? 어떻게 큰 모델로부터 작은 모델로 지식을 전달할 수 있는 걸까? 이는 신경망과 손실함수를 살펴보면 쉽게 이해할 수 있다. 다음은 Hinton이 제시한 Knowledge Distillation을 기반으로 살펴보도록 하자.

1) Soft Label

이를 레이블을 통해 구체적으로 살펴보도록 하자.

위의 경우, 기존의 개 사진의 레이블을 Original (Hard) Targets라고 생각할 수 있다. (Hard = discrete)

우리가 학습한 딥러닝 모델로 어떤 사진(실제는 개)을 넣는다면 클래스마다 확률값을 출력할 것이다. 그에 따라 가장 높은 출력값인 0.9의 클래스 “개”를 예측하게 되는 구조이다. 이때, Hinton은 예측한 클래스 이외의 값을 주의 깊게 보았다. 개를 제외한 고양이나 자동차 그리고 젖소의 확률을 살펴봤다는 뜻이다.

그렇다면 이 출력값들을 통해 무엇을 알 수 있을까? 결론부터 말하면, 이 출력값들이 모델의 지식이 될 수 있다고 말하고 있다. 우리가 예측하려고 하는 이 사진이 강아지인 것은 알겠는데 자동차나 젖소보다 고양이에 더 가까운 형태를 띠고 있다는 것을 알 수 있다는 것이다. 하지만, 이러한 값들은 softmax에 의해 너무 작아 모델에 반영하기 쉽지 않을 수 있다.

따라서 출력값의 분포를 좀 더 soft하게 만들면, 이 값들이 모델이 가진 지식이라고도 볼 수 있을 듯 하다. 이것이 바로 Knowledge Distillation의 시초(Hinton’s KD)이다. 그리고 해당 논문에서는 이러한 soft output을 dark knowledge라고 표현하고 있다.

soft하게 만들어주는 과정을 수식으로 표현하면, 위와 같습니다. 기존 softamx output과 다른 점은 T 라는 값이 분모로 들어갔다는 점이다. 본 논문에서는 온도(temperature)라고 표현하고 있고, 이 값이 높아지면 더 soft하게 낮아지면 hard하게 만드는 것을 살펴볼 수 있다. 결국, 이 T(온도)라는 설정 때문에 증류(Distillation)라는 과정이 나오게 된 것으로 볼 수 있다.

2) distillation loss

위에서 정의한 Hinton의 soft target은 결국 큰 모델(T)의 지식을 의미한다. 그렇다면 이 지식을 어떻게 작은 모델(S)에게 넘길 수 있을까? 먼저, 큰 모델(T)을 학습을 시킨 후 작은 모델(S)을 다음과 같은 손실함수를 통해 학습시킨다.

여기서 L은 손실함수, S는 Student model, T는 Teacher model을 의미한다. 또한 (x,y)는 하나의 이미지와 그 레이블, θ는 모델의 학습 파라미터, τ는 temperature를 의미한다.

위 슬라이드에서 정의하는 손실함수는 크게 기존 이미지 분류에서 사용하는 Cross Entropy Loss (LCE)와 Distillation Loss (LKD)로 구성되어 있다. LKD는 잘 학습된 Teacher model의 soft labels와 Student model의 soft predictions를 비교하여 손실함수를 구성한다. 이때, 온도(τ)는 동일하게 설정하고 Cross Entropy Loss를 사용한다.

3) Summary

- Teacher Network 학습
- Student Network 학습
- Student Network soft prediction + Teacher Network soft label → distillation loss 구성
- Student Network (hard) prediction + Original (hard) label → classification loss 구성