[AI 논문] AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

Transformer는 본래 NLP 쪽에서 주로 활용되었다면, 이를 이미지에도 적용시키고자 나온 개념이 바로 ViT(Vision Transformer)이다. 
오늘은 ViT를 다루는 논문, 'AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE'에 대해 공부해보겠다.
 
출처

- Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2021). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv.
- 김성법[교수/산업경영공학부] Youtube: https://www.youtube.com/watch?v=bgsYOGhpxDc&t=1201s

 

Introduction

NLP 분야에서는 Transfomers 아키텍처가 대세가 되었는데, 주로 large text corpus에 pre-train 시키고 이를 smaller task-specific dataset에 fine-tuning하는 방식으로 이루어졌다.
 
그러나 CV 쪽에서는 그 당시에는 ResNet 같은 CNN이 여전히 대세였다.
그래서 이 연구팀은 CV 쪽에서도 Transformer 구조를 조금 변형시켜서 적용시킨다면 활용할 수 있지 않을까라는 발상으로 실험을 시작했다고 한다.
 
NLP에서 문장을 단어로 쪼개는 것처럼, 이미지도 patch로 잘개 자르면 된다. 본 논문에서는 supervised 방식으로 image classification task를 수행시킨다.
 
한 가지 주의점은, 만약 데이터셋이 적다면 기존 CNN의 성능을 뛰어넘지 못하기 때문에 ImageNet 같은 대량의 dataset으로 pre-trained시켜야 한다!
 
*여기서 inductive bias라는 용어가 나오는데, Method에 들어가기 전에 짚고 넘어가자.
Inductive bias란 우리가 특정 데이터셋에 대해 더 좋은 성능을 얻고자 의도적으로 강제해주는 가정이다.
2가지 예를 들어서 쉽게 설명해보겠다.
CNN 같은 경우에는 인접하는 픽셀들간의 locality가 존재하기 때문에, 인접 픽셀간의 정보를 추출하기 위한 목적으로 설계된다.
마찬가지로 RNN 같은 경우에는 순차적인 정보를 잘 처리하도록 설계된다.
하지만 MLP(fully-connected layer)나 Transformer는 이러한 inductive bias가 약하다. 
 
Bias랑 Variance는 trade-off 관계이기 때문에, inductive bias가 낮을수록 성능을 높이기 위해선 대량의 데이터셋을 필요로 한다! 
 
 

Method

 

 
1. Vision Transformer (ViT)
 
제목에서도 알수 있듯이, 이 논문은 48x48의 이미지를 16x16의 패치 9개로 자른다. 
그리고 각 패치를 1차원의 벡터로 바꿔준다. 이 때 벡터의 크기는 16x16x3(RGB 채널 개수)=768이다.
이렇게 만들어진 9개의 벡터를 trainable linear projection을 통해 patch embedding으로 표현해준다. 
 
이러한 patch embedding이 바로 Transformer Encoder로 들어가기 전에, learnable class과 poositional embedding이라는 두 가지 개념이 등장한다.
 
(Learnable class embedding, Patch embedding) + Positional embedding = Transformer encoder
 
이때, Learnable class embedding이란 BERT의 class token과 유사하게 classification을 위해 사용되는 token이고,
Position embedding은 이전 논문에서 보았듯이 각 patch의 위치 정보를 나타낸다. 
(이 둘의 차원은 마찬가지로 768이고 학습을 통해 결정되는 파라미터이다.)
 
내가 아직 BERT 논문을 보지 않아서 class token 개념을 이해하지 못했다. 
인터넷에 찾아보니 sequence data를 single vector로 변환하는 방법은 여러 가지 있는데, 그 중 fisrt token을 class token으로 설정하여 이를 sequence를 대표하는 벡터로 활용한다고 한다. 
 
'ViT' Transformer encoder의 구조는 'Vanilla' Trasnformer encoder와 구조가 다르다. 
Layer normalization의 위치가 Multi-Head Attention보다 앞에 온 것을 알 수 있다. 
 

 
Multi-Head Attention에서는 Q, K, V 벡터가 생성되어 self-attention을 12번 수행한다. 
그리고 이것을 다시 연결하여 64차원의 벡터를 768차원으로 다시 복원시킨다.
 

 
최종적으로 Multi-Head Attention으로부터 나온 값을 MLP(1-hidden layer)을 통과시킨다.
그 후 class token에 해당하는 벡터(맨 왼쪽)를 normalization -> MLP해서 최종 class를 분류한다! 
 
*Hybrid Architecture
지금까지는 raw image patches를 사용했다면, CNN으로부터 추출한 raw image의 feature map을 활용할 수도 있다고 제시한다.
 
 
2. Fine-Tuning and Higher Resolution
 

이 파트에서는 ViT의 fine-tuning을 어떻게 하는지 설명한다.
먼저 pre-trained된 prediction head를 제거하고, zero-initialized된 D(768)XK(1000) 피드포워드 레이어를 연결한다. 
 

Experiment

생략
 

Conclusion

CV에 self-attention 기법을 적용시키려는 이전 방식과 ViT가 다른 점은, image-specific한 inductive bias를 introduce하지 않고 NLP에서와 같이 image를 sequence of patches로 해석했다는 것이다.
 
Large dataset로 pre-train만 된다면 ViT는 기존의 image classification SOTA 모델을 능가한다. 
 
그러나, 여전히 몇 가지 해결해야할 과제가 있다.
1. ViT를 detection, segmentation에 적용하기
2. self-supervised learning 방식 pre-training 더 연구하기 (for low cost)
3. ViT 성능 향상 위해서 모델 확장하기 
 

---

 
블로그를 쓰기 전까지는 'ViT? 그냥 이미지를 패치로 자르고 학습시키는 거겠지'라고만 생각했는데,
Inductive bias, learnable class, pre-training 등 새로운 개념에 대해 새롭게 알게 되어서 신기했다.
 
다음은 multi-modal 모델인 CLIP 논문을 공부하겠다!