KDST

이번에 소개해 드릴 논문은 데이터 없이 트레이닝된 모델만 있을 때, 우리가 knowledge distilling을 통해서 새로운 모델을 얻을 수 있을까하는 문제를 다루고 있습니다. 여러가지 문제 - 보안 레벨이라든지, 개인 정보 보호라든지 - 때문에 원 모델을 학습시킨 데이타가 접근 가능 하지 않을 경우를 가정해본다면, 충분히 생길 수 있는 문제일 것 같습니다.

일단, 원 모델을 학습시킨 데이타가 접근 가능하지 않기 때문에, 다음과 같은 상황에 유사하다고 할 수 있겠습니다. 선생님과 학생이 있는데, 사용가능한 교재가 없는 것이죠. 다음과 같은 문제들을 해결해야, 학생을 가르칠 수 있을 겁니다.

1. 교재를 어떻게 준비할까요? (우리가 데이타 포인트들을 만들 수 있을까요?)
2. 교재의 양이 충분치 않을 경우, 어떻게 잘 가르칠 수 있을까요?

이와 같은 문제를 해결하기 위해, 저자들은 다음과 같은 구조의 뉴럴 네트워크를 제안합니다.

M은 학습된 모델입니다. (따라서 고정되어있습니다) 학습되는 부분은 G와 D인데, 다음과 같이 동작합니다.

• D는 z hat 을 x hat으로 부터 생성하는 역할입니다. z hat은 x hat의 dimensional representation이라고 볼 수 있습니다. z bar 와 z hat, y bar 와 y hat의 차로부터 네트워크를 업데이트 할 loss function을 구합니다.
• 위 네트워트 (저자들은 KegNet 이라고 이름 붙였습니다)가 학습이 끝나면, 이젠 학생 네트워크를 학습시킬 수 있습니다. G는 학생 네트워크를 학습시킬 데이타 포인트들을 만들어네는 네트워크 역할을 하게 됩니다.
• 학생 네트워크는 M에서 tensor decomposition을 통해서 만들어냅니다. 저자들은 많은 방법 중에서 Tucker decomposition을 선택했구요. Tensor decomposition을 통해서 만들어낸 경량화된 모델은 fine tuning이 필요한데, fine tuning을 위한 데이터를 KegNet을 통해서 만들겠다는 것입니다.
• 만들어지는 데이타 포인트들이 너무 유사해 지는 것을 방지하기 위해, loss function에 diversity를 고려한 항이 추가되어 있습니다.
• KegNet을 y와 z를 만들기 위한 auto-encoder의 형태를 하고 있다고 이해할수 있다고 합니다.
• 또한, 학습시키기위한 데이타 포인트들을 만들기 위한 Generative Adversarial Network (GAN)의 형태로 볼 수도 있다고 합니다.

생성된 이미지들을 보면 저자들이 의도한대로 구현된 것을 보실 수 있습니다. 물론, 사용한 데이터셋이 다른 데이터셋들에 비해 상대적으로 작은 편을 감안하면, 발전시킬 여지가 많은 것이 아닐까 추측해봅니다.

정량적으로 성능 분석한 결과 일부를 발췌해 보았습니다. Tucker decomposition후에 생긴 성능 손실을 상당한 수준으로 복구하신 것을 볼 수 있습니다. 괄호 안의 숫자는 Tucker decomposition후에 얼마나 모델 사이즈가 줄었나하는 것을 표시하기 위한 scaler입니다.

더 자세한 부분을 알고 싶은 분들을 위해 원문 링크를 첨부합니다.

본 포스팅은 과거 제가 포스팅하였던 모델 압축 기술의 한 가지인 Pruning에서 화제를 전환하여 Knowledge Transfer의 일부인 Distillation 에 대한 연구 결과를 리뷰하도록 하겠습니다. 포스팅에 앞서, 주제와 관련된 모든 연구 내용은 Relational Knowledge Distillation과 Author의 YouTube를 참조했음을 먼저 밝힙니다.

Preliminary

Knowledge Transfer는 크게 Knowledge Distillation과 Transfer Learning으로 구분할 수 있습니다. 차이점을 설명하기 위해, Transfer Learning은 A라는 도메인에 대해 학습된 A'라는 모델이 준비가 된 경우, 그리고 A가 매우 크고 복잡한 데이터라고 가정합시다. 이 때, A'라는 모델을 보다 단순한 B라는 도메인에 대해 학습시키고자 할 때, A'으로부터 미리 일반화된 파라미터를 전이하는 과정을 말합니다. 한편, Knowledge Distillation은 A라는 도메인에 대해 학습된 A'라는 모델이 존재하는 경우, 이 A'라는 모델이 매우 잘 일반화되었다고 가정합시다. 이 때, 같은 A라는 도메인에 대해 B'으로 정의된 모델을 학습하고자 할 때, A'에서 학습된 정보를 B'에 전이하는 과정을 말합니다. Knowledge Distillation은 같은 domain에 대하여 일반화가 잘 된 모델과 유사하도록 다양한 모델의 성능을 확보할 수 있는 기회가 있다는 장점이 있는데요. 이를 통해, 최근 연구를 살펴보면, 경량화의 한 축으로 Knowledge distillation을 활용하는 사례도 다수 등장하였습니다. 또한, Knowledge Distillation은 일반화가 잘 된 모델 (Teacher model)로부터 보다 작은 사이즈의 모델 (Student model)에 학습 내용을 전이할 때 Student model의 성능을 잘 이끌어낼 수 있다는 장점도 있지만, Teacher model과 같은 규모의 모델로 전이하는 경우에는 기존 모델보다 더 우수한 성능이 나타나는 가능성도 확인이 되었습니다.

History

Knowledge Distillation의 전략은 몇 가지 Methodology로 구분할 수 있습니다. 먼저, Teacher model과 Student model이 규모 차이가 있는 경우, Hinton et al., NIPS, 2015에서는 아래의 그림.1처럼 Class Probability를 Transfer하여 Student model을 학습하였습니다. Softmax의 결과를 Temperature로 더욱 Smoothing하여 Feature를 추출할 물리적인 node 개수가 적은 Student model의 classifcation 성능을 일반화가 잘되도록 Support 요소를 적용한 부분이 인상적입니다.

한편, Romero et al., ICLR, 2015처럼 Hidden Activation을 Transfer하는 방식도 소개되었습니다. 규모가 큰 Teacher model의 hidden activation (C')과 보다 규모가 작은 Student model의 hidden activation (C)이 아래의 그림처럼 주어질 때, 본 방식의 목적은 C'과 C가 같아지도록 만드는 것을 목표로 합니다. 그러나, C'과 C는 기본적으로 서로 다른 채널 개수를 갖기 때문에 L2-distance를 구할 수가 없으므로, C'을 C와 같은 벡터 공간으로 transform하기 위해 적절한 transform layer를 선택해야하는 조건이 있습니다.

또 하나의 방법으로는 Zagoruyko et al., ICLR, 2017처럼 Attention을 Transfer하는 방법이 있습니다. 아래의 그림처럼 Convolution Output Tensor가 존재할 때, 이를 Channel wise하게 평균을 취하여 Special attention Matrix를 생성하고 이들 간의 Transfer를 하는 방식입니다.

이와 다르게, 규모가 유사한 모델 간의 Knowledge Distillation에 대한 방법들이 있습니다. (본 포스팅에서는 첨부하지 않겠습니다.) Furlanello etl al., ICML, 2018, Bagherinezhad et al., arXiv, 2018.

Summary

본 논문은 이와 같은 유형의 Distillation기법 즉, Teacher model의 Output을 Student model의 Output에 전이하는 방식을 일반화 하여 Individual Knowledge Distillation (IKD)라 정의하였으며, 이러한 방식과는 다르게 아래의 그림.4처럼 Relational Knowledge Distillation (RKD)을 통한 Knowledge Distillation을 제안하였습니다.

이는 그림.5에서 의미하 듯, Teacher model과 Student model 간의 output vector를 individual하게 mapping하는 것이 아니라 이들의 관계 즉, distance와 angle 두 가지 측면을 고려한 potentiality를 전이하도록 하여 보다 student model이 output의 상관관계는 이해하지만 mapping 결과는 주도적으로 만들도록 하였습니다. 본 논문의 전이 방식에서 제안하는 2가지 loss function (distance, angle)은 해당 논문을 확인해주세요.

기술한 내용 외에 본 논문에 대한 궁금한 점은 댓글로 남겨주세요.

감사합니다.

Mincheol Park