이번 글에서는 Knowledge Distillation (이하 KD)에 대해 알아보자.
KD (Knowledge Distillation, 지식증류)의 목적
KD는 이미 학습이 된 모델이 갖고 있는 지식을 다른 모델에 전이시키는 방법이다. 그러면 이렇게 물을 수 있다. 이미 학습된 모델을 사용하면 되지 왜 굳이 지식을 전이 시키냐고? 이에 대해서 답변을 해보면 딥러닝 모델이 활약하게 될 환경을 생각해야 겠다. 모델의 크기를 키우고 이런저런 layer를 추가하다보면 모델의 성능이 높아질 확률이 높다. 그러나 이런 큰 모델을 어떤 기기에 장착시켜 실행시킨다고 하면 어려울 수 있다. 아주 큰 모델을 스마트폰에 넣는다고 하면 스마트폰의 하드웨어 스펙이 아주 큰 딥러닝 모델을 돌릴만큼 스펙이 좋지 않은 경우가 태반이기 때문이다. 이를 해결하기 위해서 스마트폰에 들어갈만한 아주 작은 모델을 만들고 이 작은 모델을 학습시킬 때 크고 성능이 좋은 모델의 지식을 전이시킨다. 이러한 방식을 KD (Knowledge DIstillation) 이라고 부른다.
참고영상: [딥러닝] knowledge distillation(지식전이), teacher student learning
KD (Knowledge Distillation, 지식증류) 학습 방법
이제는 KD의 학습방법에 대해 알아보도록 하자.
teacher student 구조
KD를 위해서는 teacher model(지식이 방대한) 모델을 student model이 필요하다. teacher student 모델은 크기가 큰 teacher model f_{\theta_t}(x) 이것보다 크기가 작은 student model f_{\theta_s}(x) 두 모델로 구성되어있다. 여기서 x는 입력데이터고 \theta_t 는 teacher 모델의 parameter \theta_s 는 student 모델의 parameter를 의미한다.
teacher student 학습방법
teacher student 모델을 이용한 학습 순서는 아래와 같다.
- teacher 모델 f_{\theta_t}(x) 를 미리 학습해둔다.
- student 모델을 학습할 때 task에 맞는 loss \mathcal{L}_{task} 와 지식증류를 위한 KL-divergence \mathcal{L}_{KL} 를 사용한다. KL-divergence는 아래와 같이 생겼다.
\mathcal{L}_{KL} = KL ( \mathbf{p}/T, \mathbf{q} )
위의 식에서 \mathbf{p}는 teacher model 의 logit 값이 되겠고 \mathbf{q}는 student model 의 logit 이 된다. 그리고 T 는 teacher를 믿을 만한 정도를 뜻하는 hyper parameter를 의미한다. KL divergence는 두 분포간의 차이(discrepancy)를 보는 measure이다. 이것을 최소화 한다는 것은 teacher model 만들어내는 분포와 student model이 만들어내는 분포를 비슷하게 만든다라는 것을 의미한다. 그렇게 함으로써 student model이 teacher model과 유사한 성능을 발휘하도록 유도한다. - 2에서 제시한 task loss와 kl-divergence를 이용해 최종적으로 loss를 아래와 같이 세운 후 학습을 진행한다. 여기서 \rho 는 0과 1사이의 숫자로 hyperparameter이다. \mathcal{L} = \rho \mathcal{L}_{task} + (1-\rho) \mathcal{L}_{KL}
KD 응용
KD는 teacher 모델의 결과값을 이용해서 student model의 결과값과 비교하는 식으로 지식을 전이했는데, 사실 그럴필요는 없다. 중간 layer의 결과값이 가진 지식을 KL divergence를 통해 전이할 수도 있다. 그리고 teacher 모델의 아키텍처와 student model 의 아키텍처를 달리해서 모델이 갖고 있는 구조가 지닐 수 있는 지식을 전이할 수도 있다.