[Paper Review / Transformer]DeiT 리뷰

Training data-efficient image transformers & distillation through attention

 

초록 및 서론

기존의 ViT 연구에서는 “do not generalize well when trained on insufficient amounts of data”라고 결론지었다. 실제로도 ViT는 inductive bias가 거의 없으므로, model capacity가 높을 것이다.

이 논문에서는 이 문제를 해결하기 위해서 (1) Distillation architecture를 정의하고 (2) token-based distillation을 활용해서 적은 epoch으로 더 좋은 성능의 Vision Transformer를 학습한다. 이를 Data-efficient image Transformers (DeiT)라고 명명한다.

Contribution

1. ResNet-50, ResNet-18 보다도 파라미터가 더 적은데 좋은 성능
2. distillation token을 class token으로 사용하면서, teacher를 따라한다.
3. image transformer가 다른 transformer보다 CNN에서 더 잘 distillation 효과를 본다. 
4. CNN에서 distillation을 수행해서 CNN의 inductive bias를 Transformer로 soft하게 부여한다.

Transformer

논문에서는 ViT의 기본적인 설명을 한다. 이에 대해서는 추후 ViT 리뷰로 다시 살펴보고자 한다. 논문에서 제안하는 DeiT는 모델 구조가 이전의 ViT 그대로이며, Distillation Token이 Class Token 위에 추가된 것이 끝이다.

Distillation

soft distillation : teacher model의 logit을 그대로 활용
hard distillation : teacher model의 logit을 threshold한 결과를 활용

논문에서는 hard distillation이 더 성능이 좋았다고 말한다. 핵심적으로 제안하는 것은 distillation token을 ViT에 활용하고, Multi-head Attention을 통해서 token mixing을 하는 것.

Experiments

실제로 CNN이 더 좋은 Teacher 성능을 냈다.

 

논문에서 제안한 distillation token method가 더 성능이 좋음을 보였다. 논문에서는 이러한 성능의 결과가 CNN의 inductive bias와 연관이 크다고 주장한다.

As to be expected, the classifier associated with the distillation embedding is closer to the convnet that the one associated with the class embedding, and conversely the one associated with the class embedding is more similar to DeiT learned without distillation.

즉, distillation embedding은 CNN의 inductive bias에 영향을 받고, class embedding은 ViT의 inductive bias를 받게되니, 이 둘이 잘 앙상블 되어서 DeiT가 되고, 더 데이터 효율적인 학습이 가능하다는 의미이다.

심지어 더 적은 epoch로 수렴도 잘 된다고 한다.

But while for the latter the performance saturates with longer schedules, our distilled network clearly benefits from a longer training time.

또한, 당시 기준으로는 성능이 CNN을 못따라갔는데, 본 연구를 통해 성능이 CNN의 성능을 넘어서기 시작했다. 또한 연산량도 많이 줄일 수 있게 되었다.

마지막으로, CIFAR-10과 같은 작은 데이터셋에 대한 transfer learning의 성능도 비교하였고, ImageNet pretraining이 효과가 있음을 보였다.