[논문읽기] Knowledge Distillation via Route Constrained Optimization

이 글은 2019 ICCV Oral 논문, Knowledge Distillation via Route Constrained Optimization를 참고하여 작성하였습니다.

이번 서울 코엑스에서 열린 ICCV 학회에 참가하여 직접 발표를 듣고 감명받아 읽어보게 된 논문입니다. 일단 굉장히 단순한 기법으로 Knowledge distillation 성능을 높였다는 점이 흥미로웠던 것 같습니다.

시작하기에 앞서, 기본적인 knowledge distillation 개념 관련 설명은 Lunit Tech Blog를 참고하면 좋을 것 같습니다!

그럼, 가볍게 route constrained optimization(RCO) 논문 리뷰해보도록 하겠습니다:)

Background

일반적으로 teacher와 student model을 갖고 knowledge distillation을 학습할 때는 아래 그림 (a), (b)와 같이 이미 수렴된 teacher model을 사용합니다.

하지만, 이 논문에서는 수렴된 teacher 모델을 사용하는 것에 대한 문제점을 지적하고 있습니다. 이미 수렴된 teacher 와 student 사이의 gap이 굉장히 크기 때문에 knowledge를 전달하는데 한계가 있다고 주장합니다.

수렴된 teacher 모델보다는 early training stage에 있는 teacher 모델이 student 모델과의 gap이 더 작다고 합니다. curriculum learning의 관점에서는 teacher의 knowledge 가 training process를 거칠수록 점차적으로 어려워진다라고 이해할 수 있습니다.

따라서, teacher 모델을 점진적으로 학습시키자가 이 논문에서 제시한 아이디어라고 볼 수 있겠습니다.

Gradually mimicking the route sequence of teacher.

Route Constrained Optimization

위 그림은 논문에서 제시한 RCO의 전반적인 프레임워크입니다. Teacher model의 step1 부터 stepN 까지를 이용하여 점진적으로 student를 학습하는 것을 볼 수 있습니다. 이때, \(T_{step1}\)~\(T_{stepN}\) 은 모두 pre-trained 모델들 입니다.

아래 알고리즘 그림을 보시면 더욱 이해하는데 도움이 될 것 같습니다.

KD Loss는 다음과 같습니다.

\[L_{KD} = H(y, P_s) + \lambda KL(P_t^\tau, P_s^\tau)\]

여기서 KL은 Kullback-Leibler divergence이고, \(\tau\)는 relaxation hyperparameter (Temperature), \(P_t\)는 \(softmax(z_t)\) 입니다.

RCO 알고리즘을 보다보면, teacher model의 anchor point를 어떻게 지정하는지에 대해서도 의문점이 생깁니다. 이 논문에서는 greedy search를 사용하여 최적의 anchor point를 찾는 과정을 설명하고 있습니다. 하지만 직관적으로는 learning rate를 tuning 하는 시점이 anchor point로 지정하기에 좋다고 합니다.

Results

여러 실험결과들이 논문에 나와있지만, ImageNet으로 실험해본 결과만 가져왔습니다. 이전 accuracy보다 더 좋은 성능을 보이고 있는 것을 볼 수 있습니다. 더 많은 결과들은 논문에서 볼 수 있습니다.

---

Reference:
[1] RCO 논문