Revisiting Knowledge Distillation: An Inheritance and Exploration Framework (CVPR 2021) 논문 리뷰
카테고리 없음 2021. 8. 26. 14:54이번에는 CVPR 2021에 poster paper로 accept 된 Revisiting Knowledge Distillation: An Inheritance and Exploration Framework을 리뷰하려고 합니다. 자세한 내용은 원문을 참고해주세요. 코드의 경우, 비공식 코드로 구현해두었습니다.
Introduction
Class distribution사이에서 similarity와 consistency의 정보를 전달함으로써 knowledge distillation 분야가 성공할 수 있었다고 저자는 말하고 있습니다.(Hinton의 KD, Tf-KD etc.) 또한, layer 중간의 output에서 feature representation을 transfer하는 방법도 성공을 이끌었다고 말하고 있습니다.
하지만, 직접적으로 teacher를 따라하게 student에게 knowledge를 전달하는 것은 student가 알려지지 않은 feature를 탐색할 확률을 줄일 것이라 생각했고, 이를 확인한 사진이 아래의 사진입니다.
사진을 보시면 knowledge distillation을 진행했을 때, teacher 가 바라보는 feature부분을 그대로 보는 것을 알 수 있습니다. 실제 답은 치타임에도 꼬리부분을 악어라고 인식하게 되는 것이라고 할 수 있습니다.
또 이러한 사실을 더 명확히 확인하기 위해서 small network에서 large network로의 knowledge distillation을 진행했을 때, 아래의 사진처럼 loss가 오히려 증가하는 현상을 발견했습니다. 즉 teacher의 network가 작더라도 compact하게 지식을 갖고 있다는 것입니다. (ind와 비교할 때 kd가 train loss 학습은 더 빠르지만, evaluation에서 보면 점점 더 높은 loss로 saturation됨)
이러한 점을 들어 단순히 KD에서 teacher 정보를 받아들이기만 해서는 안된다고 저자는 말하고 있습니다.
Related Work
Knowledge Distillation(KD)에는 전달하는 지식, 전달하는 시점에 따라서 분류가 갈리는데 여기서는 크게 3가지만 소개하도록 하겠습니다.
- Offline-KD
출력(logits)으로 부터 학습하는 Knowledge Distillation
중간 Feature로 부터 학습하는 KD: 1) Factor Transfer, FitNet, Attention Transfer, etc.
- Online-KD
Student의 ensemble이 학습하는 방법: DML(Deep Mutual Learning)
Sub-net의 classifier ensemble이 학습하는 방법: CL(Collaborative Learning)
- Teacher free Regularization
Student가 teacher 없이 referred target에 의해서 knowledge 가 transfer되는 방식(Tf-KD, Snapshot Distillation, etc.)
Motivation of Proposed Method
앞서 저자가 teacher에 의해 transfer되는 knowledge에 의해 같은 위치의 feature만 관찰한다고 말했습니다. 이를 극복하는 방법을 저자는 진화이론에서 떠올렸다고 합니다.
진화 이론은 inheritance와 variation of traits에서 발생한다고 하는데, 이 중 variation of traits를 달성할 방법에 대해서 고민하게 됩니다. 저자는 강화학습의 Q learning에서의 epsilon greedy에 의한 exploration과 과거 논문 중 image를 학습하기 어렵게 만드는 universal adversarial network의 loss term에서 영감을 받았다고 합니다. 즉, exploration을 통해 다른 feature도 관찰하고 이는 더 좋은 위치를 보게 만들어 줄 수도 있다라는 생각을 갖게 해준 것입니다.
아래 사진을 보시면 TPAMI 2020에 게재된 논문에서 $L_{dstc}(x)$ 를 사용함으로써 original의 이미지에 loss에 의해 발생하는 gradient를 전달하여 동일하게 보이지만 network에 의해서 다른 곳을 관찰하게 유도하는 이미지를 생성합니다.
이런 것을 활용하여 다른 feature도 관찰할 수 있게 유도한 논문을 통해 다른 feature를 관찰할 수 있는 exploration method를 선택하면 단순히 teacher를 따라하지 않아도 된다! 라고 생각한 것이라고 할 수 있겠습니다.
Proposed Method
이러한 inheritance(KD)와 exploration을 같이한 framework를 제시한 것의 구조는 다음과 같습니다.
제시한 loss term은 다은과 같습니다. $L_{goal}$은 classification의 CE Loss나 object detection에서 일반적으로 사용하는 loss term입니다.
$$L=L_{goal}+\lambda_{inh} L_{inh}+\lambda_{exp} L_{exp}$$
즉 추가한 두 loss term을 통해서 이를 조절한다는 것입니다.
$L_{inh}$: teacher의 feature를 따라하는 loss(inheritance) , $L_{inh}=||{F_{inh}}/{||F_{inh}||_2}-(F_T)/{||F_{T}||_2}||$
$L_{exp}$: teacher의 feature대신 다른 feature를 배우는 loss(exploration) $L_{exp}=||{F_{exp}}/{||F_{exp}||_2}-(F_T)/{||F_{T}||_2}||$
위 두 loss는 encoding 된 latent vector를 비교함으로써 사용되고
여기서 Encoder는 teacher의 경우, pretrained model을 사용하므로 auto-encoder를 통해 미리 학습한 후 사용하며, student는 학습과 동시에 사용합니다.
Auto-encoder
- Auto-encoder는 이 논문에서 teacher와 student간의 shape이 다른 것을 해결하기 위해서 제시하였습니다. (Tecaher만 사용하고 Student는 학습과 동시에 진행합니다) 동일 latent space로 옮겨서 확인하려고 만들었습니다. 논문에서는 아래의 구조로 사용했다고 합니다. Loss term의 경우 reconstruction loss인 L2 norm을 사용하였으며 식은 $L_{rec}=||f_T-R(f_T)||^2$ 을 사용하였습니다.
이러한 구조를 사용하여 학습하는 방법은 다은과 같습니다.
1. 먼저 reconstruction loss를 사용하여 teacher's factor를 학습하는 auto-encoder를 만듭니다.
2. Student의 encoder와 backbone network를 동시에 학습합니다.
IE-KD in Deep Mutual Learning(IE-DML)
- Online learning에서는 IE framework를 적용하기엔 pretrained_model이 존재해야 teacher의 faeture를 학습하는 auto encoder 먼저 학습하므로 다른 방식이 필요합니다. DML의 경우 AutoEncoder까지 동시에 학습시킵니다.
$$L_{\theta_1}=L_{goal1}+\lambda_{inh1} L_{inh1}+\lambda_{exp1} L_{exp1}+\lambda_{exp} L_{exp1}$$
Experiments
IE-KD의 경우 feature 기반 Knowledge transfer에 모두 적용할 수 있습니다. 그래서 본 논문에서는 SOTA였던 모든 method에 대해 적용하여 성능을 비교하였습니다.
- Attention Transfer(ICLR 2017) -> IE-AT
- Factor Transfer(NIPS 2018) -> IE-FT
- Overhaul of Feature Distillation(ICCV 2019) -> IE-OD
Ablation Study
저는 개인적으로 진짜 feature들이 보는 곳이 보기만 하는 것이 맞는지 궁금했습니다. Ablation Study에서 이를 보여주고 있는데 기존 baseline에 대비해서 Centered Kernel Alignment (CKA)를 측정한 결과, 기존보다 다소 낮은 것을 확인할 수 있었습니다. 그리고 baseline에 비해서 더 많은 neuron들이 activate됨을 통해 효율적으로 사용하고 있음 또한 확인할 수 있었습니다.
그리고 앞서 언급한 epsilone greedy와 같이 exploration과 exploitation 비율이 어느정도일 때 적절한지 확인한 실험도 존재하였습니다. 적절하게 balanced 되어있을 때, 더 좋은 성능을 보여줌을 확인할 수 있었습니다.
Conclusion
Knowledge Distillation이라는 분야에서 단순히 지식을 넘겨주는 측면 뿐만 아니라 새로운 것 또한 배울 수 있어야 발전한다고 제시한 논문입니다. 그 크기에 비해서 성능 개선폭이 낮은 것은 아쉽지만, 관전의 변환을 준다는 점에서 흥미로운 논문이라고 생각합니다.
자세한 내용은 논문을 참조하시거나 댓글로 남겨주시기 바랍니다.
코드는 official code가 없어 제가 구현해둔 코드에서 성능을 확인해보셔도 좋을 것 같습니다.
