[7] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

 

[Paper] https://arxiv.org/pdf/2010.11929.pdf

[Github] https://github.com/google-research/vision_transformer

GitHub - google-research/vision_transformer

 

이번 포스팅에서는 자연어처리 분야에서 눈부신 성능 발전을 가져온 Transformer 구조를 비전 분야에도 적용한 Vision Transformer(ViT) 에 대해 설명하겠습니다.

 

 

Transformer의 전체적인 구조와 작동과정은 아래 두 글을 참고하면 됩니다!

https://ga02-ailab.tistory.com/144

[딥러닝 기본지식] Self Attention과 Transformer (1)

https://ga02-ailab.tistory.com/145

[딥러닝 기본지식] Self Attention과 Transformer (2)

 

 

1. Introduction  

Transformer 구조는 자연어처리 분야에서 처음 등장했고 이는 굉장한 성능 발전을 이루는 결과를 가져왔습니다. 그럼에도 불구하고 비전 분야에서는 제한적으로 적용되어 왔고 최근까지 convolution 구조가 주로 사용되고 있습니다.  하지만 Transformer의 가장 핵심인 self-attention을 비전 분야에 적용하려는 시도가 아예 없었던 것은 아닙니다. CNN 구조를 self-attention과 결합하려는 시도가 있었으나 self-attention의 특수한 구조 때문에 현대의 hardware accelerators를 효과적으로 scaling하기가 어려웠습니다. 

 

따라서, 본 논문에서는 기본적인 transformer를 최소한으로 수정하여 비전 분야에 적용한 Vision Transformer를 소개합니다. NLP분야에서 문장을 단어 단위로 쪼갠 token이 입력으로 들어가는 것처럼 이미지를 일정크기의 patch로 쪼개어 입력으로 사용합니다. 

성능에 대한 결과를 살펴 보면 ImageNet처럼 약 130만 정도의 크기를 가지는 데이터 집합에서는 ResNet 보다 낮은 정확도를 기록했습니다. 이는 Transformer 구조 자체가 CNN보다 inductive bias(unseen data를 예측하기 위해 사용하는 가정, 모델의 일반화 능력)가 부족하기 때문입니다. 하지만 대형 데이터셋에서는 큰 scale의 학습이 inductive bias를 이겨 높은 성능을 달성했다고 합니다.

 

 

 

2. Model  

2.1 Vision Transformer (ViT)  

전체적인 작동과정을 살펴보겠습니다. 

이미지를 일정 크기로 여러 개의 patch로 쪼개어 linear projection 시킵니다. 이 값과 position embedding 값을 Transformer Encoder에 넣어줍니다. NLP 분야에서의 Transformer는 token emvedding을 1차원으로 하여 받습니다. 비전 분야에서도 이를 따르기 위해 2차원의 patch들을 flatten하게 펴서 1차원으로 만들어줍니다. 간단히 수식으로 정리해보겠습니다.

- 2차원의 이미지는 아래 patch 사이즈로 변환됩니다.

p: patch의 resolution

H,W :원본 이미지의 resolution

N: patch의 수(=HW/p^2)

그리고 모든 layer에서 동일하게 D 사이즈의 latent vector 를 갖습니다. flatten 시킨 각 patch를 학습 가능한 Linear projection을 사용해 D차원으로 매핑하여 patch embedding을 얻습니다.

 

BERT와 마찬가지로 [CLS] Token도 존재합니다. BERT의 [class] Token처럼 학습가능한 embedding patch를 추가해준 것입니다. 

 

classification head는 pre-training시에는 hidden MLP로 수행되고 fine-tuning시에는 1-linear layer로 구성됩니다. 

 

마찬가지로 position embedding도 사용되는데 이미지이지만 standard learnable 1D position embedding을 사용합니다.(2D를 사용했을시 보다 성능이 좋았다고 합니다.)

이러한 과정을 통해 계산되는 representation y는 아래와 같습니다.

 

- Inductive Bias

먼저 Inductive bias란 unseen data를 예측하기 위해 사용하는 가정, 모델의 일반화 능력을 말합니다.

ViT는 CNN보다 이미지별 Inductive Bias가 적습니다. CNN에서는 locality(근접한 픽셀들끼리만 종속성을 가짐), Two-Dimensional neighborhood structure , Translation equivariance 때문에 Inductive Bias가 존재합니다. ViT에는 이 세 가지가 없기 때문에 Inductive Bias가 훨씬 약하게 발생하는 것입니다. 때문에 ViT는 아래 두 가지 방법으로 Inductive Bias가 작동할 수 있도록 합니다. 

- MLP layer: mlp에서 locality하게 작용합니다.

- self attention: 이미지를 patch로 자르는 과정, position embedding(초기화시마다 모든 spatial relation을 처음부터 학습해야 함)

 

 

- Hybrid Architecture

Transformer의 입력으로 raw image patch가 아닌 CNN의 feature map patch를 넣어줄 수도 있습니다. 이때, patch의 크기를 1로 설정해줄 수 있는데 crop할 필요없이 바로 flatten하여  D차원으로 projection 하면 됩니다. classification input embedding과 position embedding은 기존과 같은 방법으로 추가하면 됩니다.

 

 

2.2 Fine-Tuning and higher resolution  

일반적으로 large dataset에서 ViT를 pre-train 하고 , 더 작은 dataset에 fine tuning하여 사용합니다. 이를 위해, pre-train된 prediction head를 제거하고 0으로 초기화된 D * K(class num) 크기의 feedforward layer를 붙혀줍니다.  pre train시 더 높은 resolution을 가지는 이미지로 fine-tuning하는 것이 더 좋은 성능을 낸다고 하네요! 

하지만 positional embedding이 효과가 없어지는 단점이 발생하기도 해서 input 이미지크기의 patch 위치에 맞게 position embedding 값을 2차원 보간하여야 합니다. 이러한 resolution 조정과 patch 추출이 ViT에 image-specific inductive bias를 추가해 줄 수 있습니다!

 

 

3. Experiments  

3.1 Comparison to state or the art  

JFT로 pre-train 후 , 각 데이터 셋에 대해 transfer learning을 진행한 결과입니다.

 

3.2 Pre-training data requirements  

앞서 설명했듯이 ViT는 inductive bias가 CNN보다 부족하기 때문에 pre-train시 사용하는 데이터 집합의 크기가 클 수록 좋은 성능을 냅니다. 이번 실험에서는 ImageNet(1.3M), ImageNset21K(14M), JFT(300M)에 대해 pre-train한 후 ImageNet 데이터 집합을 분류 했을 때의 성능입니다.