DeiT : Training Data-efficient Image Transformers & Distillation through Attention 논문 리뷰

 

Abstract & Introduction

- ViT모델을 유지한 채 학습 방법을 조금 달리해서, 오로지 ImageNet 데이터셋과 single 8-GPU만으로 2~3일만에 학습하고 그에 준하는 Performance를 얻었습니다.

정확도  &  속도 간  trade-off  그래프

- 우측 상단에 있을 수록 좋은 모델입니다.

- 여기서 ViT-B는 SOTA ViT모델이 아니라 ImageNet 데이터만으로 학습했을 때의 ViT 모델인데 이 모델과 DeiT-B와는 꽤 많은 성능차이가 나는 것을 볼 수 있습니다

 

- DeiT-B 옆에있는 요상한 증류기 기호는 Distillation 학습방법을 사용한 모델에 붙는 기호인데 이에 대해서는 아래에서 설명하겠습니다.

---

Related work

- Image classification은 computer vision의 core task 입니다.

- Alexnet의 등장 이후, convolution network는 해당 분야의 standard였습니다.

- ViT 등장 전에도 transformer를 이용한 classification 시도가 있었지만, 성능이 부족했습니다.

ex) Squeeze-and-Excitation Networks, 2017,  우승모델

squeeze and excitation network란, 채널 별로 attention을 따로 계산해서 일종의 self attention 처럼 동작하는 모델을 말합니다.

 

- ViT closed gap, This performance is remarkable since convnet methods have benefited from years

(= ViT는 CNN과의 격차를 확연히 줄였습니다. CNN이 2012년부터해서 약 10년동안 발전해온 것을 생각하면, ViT는 이제 시작입니다!!)

- JFT dataset 사용, pre-training 시 굉장히 많은 양의 데이터가 필요하다는 단점이 존재합니다.

= ImageNet 데이터만으론 학습 시 성능이 감소합니다!

그러나, JFT 데이터셋은 오픈되지 않은 구글만의 데이터셋입니다..

 

이 논문에서는 ViT모델을 유지한채 학습방법을 조금 달리해서 오로지 ImageNet 데이터셋과 single 8-GPU만으로 2~3일만에 학습하고 그에 준하는 Performance를 얻을 수 있다고 이야기 합니다!

---

Architecture - ViT

- ViT 모델을 그대로 사용

---

Attention and Inductive bias

- Attention 모델들은 CNN이나 dense layer들 보다 Inductive bias가 적습니다.

=> 그래서, 더 많은 데이터들이 필요합니다.

1. cnn은 지역적으로 주변에 있는 정보들과 연결되어 있습니다, 멀리 떨어진 x4와는 연결 안되어있구요! , 또한 모든 입력에 대해 같은 가중치를 사용합니다.

 

2. dense layer는 fully connected되어 있습니다. 그리고 역시 모든 입력에 대해 같은 가중치를 사용합니다

 

3. attention layer는 dense처럼 fully connected 되어 있지만, 각각의 입력 X ( 엄밀히 말하면 x으로부터 나온 value v) 와 [q2와 key 1의 내적의 softmax 값]과 곱해집니다. (예시에서는 2번째 입력이므로 q2라 표현)

 

결국 1,2,3 모두 입력들과 가중치를 곱해서 weighted sum을 하는 것은 동일한데, 3번의 경우 이 과정에서 사용되는 weight가 input에 따라서 다릅니다! => 즉, 3번은 가장 제약이 없고, 가장 inductive bias가 적다는 것을 뜻합니다.

 

자유도가 가장 높아 그만큼 성능을 끌어올릴 수 있는 반면, inductive bias가 거의 없어 그만큼의 data가 필요한, 일장일단의 모습이라고 할 수 있습니다.

(inductive bias란, 모델이 학습과정에서 본 적이 없는 분포의 데이터를 입력 받았을 때, 해당 데이터에 대한 판단을 내리기 위해 가지고 있는, 학습과정에서 습득된 Bias(편향)이라고 말할 수 있을 것 같습니다.)

---

Knowledge Distillation

- teacher model의 inductive bias를 student에게 transfer 하는 방법입니다.

-student 모델은 teacher 모델보다 파라미터 수가 훨씬 적습니다. 크기가 더 적은 모델로도 크기가 큰 teacher 모델만큼의 성능을 낼 수 있게끔 하는 것이 목표입니다.

- student 모델이 teacher 모델의 softmax값을 전달받아, student 모델의 성능을 높이려는 방법입니다.

- teacher는 hard하게 학습, student는 soft하게 학습합니다

논문에서는 soft distillation 방법을 이용해 학습을 하였고, 논문에서는 hard distillation 방법도 진행해서 비교해봤다고 합니다.

논문에서 사용한 방법

- KL divergence란, 두 확률분포의 차이를 계산하는데 사용하는 함수입니다.

- zt는 teacher model output, zs는 student model output을 의미

 

1. soft distillation은 teacher의 softmax값과 student의 softmax값의 KL divergence를 최소화하는 식입니다.

식의 왼쪽 항은 student model의 output과 true label과의 cross entropy를 구하고,

오른쪽 항은 student와 teacher의 softmax값을 smoothing해서 KL divergence를 구합니다

 

λ : 왼쪽 항(cross entropy loss)과 오른쪽 항(KL divergence)의 중요도를 가르는 가중치

τ : softmax값을 soft하게 바꿔주는 일종의 smoothing 역할 (ex) 0.98, 0.01, 0.01 -> 0.80, 0.10, 0.10)

Ψ : softmax 기호

KL : Kullback-Leibler divergence는 두 확률분포의 차이를 수치적으로 나타내는 값을 의미하며 정답과 예측에 대한 두개의 확률 분포의 Cross Entropy - true분포에 대한 Entropy(정보량)로 구합니다. 즉 true 분포와 비교해서 상대적으로 얼마나 정보량의 차이가 나는지를 의미하기 때문에 relative entropy라고도 합니다

 

KL divergence의 특징 

1. 0보다 크다. 

2. KL(P||Q) != KL(Q||P)이다. 

 

2. hard distillation의 왼쪽 항은 student model output과 true label간의 cross entropy를 구하고 오른쪽 항은 student model output과 yt값 간의 cross entropy를 구하는데,

(yt란, teacher model output의 argmax하여 얻은 hard한 값)

soft distillation의 λ와 τ가 없어졌기 때문에 하이퍼 파라미터로부터 자유롭고 이 hard distiilation 또한 ε을 통해서 label smoothing을 시킬 수 있습니다.

---

Distillation Token

기존 ViT 모델의 입력에 distillation Token을 하나 추가합니다.

class token이 classification을 위한 토큰이 되는 것처럼, distillation token은 마지막 레이어에서 teacher model output과 cross entropy를 통해 학습됩니다.

즉, 이 토큰으로 생성되는 distillation embeding은 teacher 모델이 예측하는 결과를 갖고 있습니다. 이 결과를 student model이 활용하여 학습하는 것 입니다.

( inductive bias 때문에 teacher model을 transformer보다 CNN을 사용하는게 더 좋은 성능을 낸다고 합니다!)

 

이때, 마지막 layer에서 class token과 distillation token의 코사인 유사도(= 두 벡터가 가리키는 방향의 유사도) 를 측정했을 때, 0.93으로, 1보단 낮지만 엄청 유사한 것을 확인할 수 있습니다.

왜 1이 안되냐하면, augmentation을 할 때, 주로 crop and resize를 사용하는데 , 이 때 아래 그림처럼 이미지와 label이 잘못 매치될 수 있습니다. 그래서 코사인 유사도가 1이 안되고, 이러한 경우에는 teacher가 student에게 도움을 줄 수 있다고 합니다.

---

Experiment

- teacher를 cnn으로 사용했을 때, 정확도가 더 높음을 확인할 수 있습니다.

- 다른 task에서도 충분한 성능을 가짐을 확인할 수 있습니다.

---

Training details and ablation

찐하게 표시된 부분이 달라진 부분입니다.

 

Batch Augmentation

배치안에 있는 모두 다른 이미지들을 한번만 Augmentation하지 말고 현재 배치사이즈에 있는 이미지에 대해 여러번 Augmentation 해서 사용하면 모델을 더 잘 일반화 시킬 수 있다고 주장하는 방법입니다.

예를들어 배치사이즈가 16일때 이미지를 4번씩 Augmentation을 하게되면 이때 배치사이즈는 16 x 4 = 64 가 됩니다.

 

Repeated Augmentation

 Batch Augmentation거의 동일하지만, Batch Augmentation은 그만큼 배치사이즈가 늘어난다면 Repeated Augmentation은 기존 배치사이즈가 64일때 independent한 이미지 64장을 16장으로 줄이고 그것을 장당 4번씩 Augmentation 하게된다.

이렇게 되면 sample들이 not independent해지기 때문에 오히려 성능이 떨어지는데 배치사이즈 커지게 되면 오히려 배치내의 이미지가 independent할 때 보다 성능이 좋아진다고 합니다.