Triplet Knowledge Distillation

https://arxiv.org/abs/2305.15975

Triplet Knowledge Distillation

Xijun Wang et al. 

 

 

KD에서, teacher가 student보다 크기 때문에, teacher의 solution을 student가 학습하기 어렵다.

이를 해결하기 위해, Triplet Knowledge distillation인 TriKD를 제안했다.

* Teacher, Student, Anchor

노란점선:subspace, 빨간별:anchor, 초록:teacher, 노랑:student

 

distillation을 하기 전에, pretrained anchor model은 target problem에 대한 full solution space의 subspace를 구분한다.

이 subspace의 solution은 student가 잘 학습할 것이라고 기대한다.

 

online 방식으로 distillation이 시작되면, teacher모델은 위의 subspace안에서만 탐색이 가능하다.

위의 hint를 통해서, student는 좋은 성능을 낼 수 있고, 잘 학습된 student는 다음 step에서 anchor로 사용되면서, 새로운 학습 전략을 구성할 수 있다.

 

다양한 모델과 face recognition 실험에서 효과를 입증했고, overfitting문제에도 도움이 됐다.

또한 이론적인 분석으로 triple distillation의 필요성을 입증했다.

 

 

 

Introduction

기존의 KD에서는 student가 teacher의 지식을 잘 받는 것에만 집중하고, teacher의 학습에는 거의 투자하지 않았다.

그러나 teacher를 건드리지 않으면, teacher model의 generalize는 잘 되지만, distillation이 어려운 solution이 나올 위험이 높아진다.

 

위의 주장과 비슷하게, "KD의 제일 큰 장벽은 optimization의 어려움이다" 라고 주장하는 (Does Knowledge Distillation Really Work?) 논문과 정확도가 높은 teacher일수록 오히려 성능이 나쁜 student를 만든다는 최근 3개의 논문도 있었다. (On the efficacy of knowledge distillation, Student customized knowledge distillation, Improved knowledge distillation via teacher assistant)

 

input image -> target output의 function space를 고려할때, teacher의 용량이 크기 때문에, teacher의 sub-space $F_{T}$가 student의 sub-space $F_{S}$보다 크다.

teacher의 solution이 student $F_{S}$가 얻을 수 있는 범주를 벗어난 결과가 나왔을때, student는 그 solution을 배울 수 없다.

 

논문의 TriKD는 online kd를 기반으로 하며, "teacher를 성능도 좋고, 따라하기도 쉽게 만드는 방법은 없을까?"에서 영감을 받았다. 

online teacher와  student에게 anchor를 제공하며, anchor는 small-model에 친숙한 접근법으로 target task를 풀도록 제한해준다.

pretrained anchor는 student와 같은 capacity를 갖기 때문에, anchor로 표현할 수 있는 $f_{A}$는 $F_{S}$에 있어서 student가 더욱 쉽게 따라할 수 있게 된다.

 

이러한 function distance 규제를 student와 teacher에게 걸면서, student와 teacher의 search space를 모두 $f_{A}$로 당겨준다.

위 규제를 걸었을때, teacher의 search space가 $f_{A}$ 근처로 당겨지고, 이는 $F_{S}$와도 가까워지는 효과가 있으므로, 더욱 따라하기 쉬운 teacher의 search space를 갖을 수 있다.

 

근데 이때, small search space를 갖도록 강제해도 여전히 높은 accuracy를 갖을 수 있다. (capacity가 높기 떄문이라고 한다.)

이렇게 높은 정확도지만 student가 따라하기 쉬워서 distillation이 잘 되는 모델이 만들어진다.

 

적합한 pretrained anchor를 찾기 위해 커리큘럼 전략을 사용했다.

커리큘럼 전략이란, TriKD로 학습한 student를 다음 generation의 anchor로 사용하는 것이다.

• Q. next generation이란? iteration? epoch? training schdule?
  • A. Convergence
• Q. 제일 처음 anchor는 무엇을 사용하나?
  • A. Online KD (teacher <-> student), Anchor 없음

논문의 contribution을 요약하면 다음과 같다.

1. We propose a novel triplet knowledge distillation mechanism named TriKD.

2. To find a proper anchor model for TriKD. (curriculum strategy)

3. TriKD achieves SOTA performance on KD.

4. Therorical analysis in statistical perspective is given to analyze the rationality of triplet distillation.

 

Larget Teacher, Worse Student.

직관적으로 생각할때, teacher의 accuracy가 높고 capacity가 클수록 student의 성능이 증가해야한다.

하지만, 2019년 ICCV On the efficacy of knowledge distillation에서 매우 큰 teacher는 student를 악화시킨다는 것을 보였다.

이러한 현상은 이후 2020년, 2021년 논문에서도 계속 관찰되었다.

 

ESKD (위 2019년 논문) 에서는 early-stopping strategy로 해결하는 모습을 보였고

SCKD (위 2021년 논문)에서는 teacher와 students의 distillation loss gradient의 similarity를 통해 distillation을 조절했다.

TAKD (위 2020년 논문)에서는 distillation process를 multi-stage로 나누고, 중간크기의 teacher assistant라는 모델을 하나 더 두어, capacity gap을 줄여서 해결했다.

근데, TAKD는 이 student가 학습하기 어려운 요소의 원인을 teacher model의 capacity가 너무 큰 것이라고 이야기한다. 

이러한 부분에 대해서 이 논문은 capacity를 건들일 것이 아니라, 이는 그대로 두고, express부분을 변경해야한다고 주장한다.

 

Method

위에서 말한 것 처럼, TriKD는 3가지의 모델을 사용한다.

• online teacher $T$
• student $S$
• anchor $A$

anchor는 teacher와 student를 supervise하고, teacher와 student는 mutually 각각 서로를 학습한다.

anchor는 pretrained를 사용하고, student와 teacher는 random initiialized를 사용한다.

학습할때는 anchor는 학습하지 않고, student와 teacher만 학습에 사용한다.

 

일단 anchor를 사용하는 목적은 teacher와 student모두 student에 친숙한 function express를 만들어주는 것이므로, anchor의 function express $f_{A}$를 student에 친숙한 형태로 만들어야한다. 

• anchor는 student와 동일한 architecture와 size를 사용한다.

이러한 anchor를 활용하여 teacher와 student의 search space를 $f_{A}$와 가깝도록 규제를 걸어준다.

 

규제는 KL Divergence로 건다.

softmax.

최종 Loss형태.

 

anchor KD를 통해, teacher-student는 점점 더 유사해질 수 있다.

이 과정에서 teacher의 capacity는 건들이지 않았기에, search space에 규제가 걸려서 Anchor와 비슷해졌더라도, 더욱 low-expected-risk는 잘 찾을 수 있고, 이를 통해서 student에게 더 좋은 지식을 전수해줄 수 있다.

 

규제가 강하게 걸리면, teacher model의 최대 성능이 anchor에 맞춰질 가능성이 있으므로 적절한 하이퍼파라미터를 선택해야한다.

 

위 모든 과정의 중심은, anchor인데 적절한 anchor를 어떻게 찾을 수 있을까?

• 논문에서는 curriculum strategy을 제안한다.

 

간단하게 설명하면, $g$번째 student가 $g+1$ 번째의 anchor가 된다.

• 제일 처음 anchor는 이전 세대의 student가 없으므로, 0세대는  teacher model과 online distillation으로 학습한다.
  • online distillation을 하지 않고, label로만 학습한 애들 anchor로 사용해보려고 했는데, 성능은 비슷한데 수렴속도가 느렸다.
• 몇번까지만 성능이 오르고, 그 후에는 성능이 수렴한다.

 

 

Q. S는 A에만 KD걸어도 되는거 아닌가? 굳이 teachr KL?

A.

classification baseline + Teacher KD + Anchor KD가 제일 높다.

Experiments

기존 KD방법들보다 성능이 좋음을 알 수 있다.

 

 

Classification loss + KD + Anchor KD 비교