[논문리뷰] An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale

이번 포스트는 Vision Transformer라고 불리는 논문입니다. Transformer가 NLP에서의 게임 체인저로 중요한 역할을 했기에, 이러한 방법을 Computer Vision에서도 사용하는 방법에 대한 논문이라고 볼 수 있습니다.

An Image Is Worth 16X16 Words:Transformers for Image Recognition at Scale

1. Abstract

• Computer Vision에서는 Attention은 CNN이 함께 적용되어 사용되었습니다.
• 이 논문은 CNN 없이 sequence of image에 적용되어 Image Classification task를 잘 수행함을 보여줍니다.
• ViT(Vision Transfomer)는 CNN에 비해 뛰어난 성능을 보임과 동시에 필요한 연산량이 상당히 적습니다.

2. Introduction

• Transformer model 로 인해 NLP (Natural Language Processing) 분야에서 100B 이상의 매개변수를 가지는 매우 큰 모델 학습이 가능해졌습니다.
• 여전히 Large-Scale 이미지에 대해선 ResNet이 최신 기술이었습니다.
• NLP 에서 Transformer Scaling 성공하여 이미지에 적용하였습니다. 이미지를 patch로 분할하고, 이를 sequence of linear embedding로 input으로 사용합니다.
• ViT는 상당히 큰 규모로 pre-trained 되고, 적은 data 에서 전이학습 시 상당히 좋은 결과를 가져옵니다.
• 가장 좋은 모델은 ImageNet-88.55%, ImageNet_Real-90.72%, CIFAR_100-94.55% 등에서 상당한 정확도를 보여줍니다..

3. Related Work

NLP

• Transformer : 기계 번역 모델로 제안 되었으며, 현재 많은 NLP Task에서 가장 좋은 성능을 보여주는 구조입니다.
• BERT : 노이즈를 제거하는 self-supervised pre-training을 수행합니다.
• GPT : 언어 모델을 pre-training task로 사용합니다.

Computer Vision

• self-attention을 이미지에 단순 적용하는 방식은 해당 픽셀과 모든 픽셀 간의 attention weight을 구해야 하기 때문에 계산비용이 pixel 개수 n에 대하여 $O(n^2)$의 복잡도를 가지게 됩니다.
• Parmar et al. (2018): local neighborhood에만 self-attention을 적용합니다.
• Sparse Transformers (2019), Weissenborn et al. (2019): attention의 범위를 scaling하는 방식으로 self-attention을 적용하고자 하였습니다..

→ 이러한 특별한 방식의 attention의 경우에는 하드웨어 가속기에서 연산을 효율적으로 수행하기에는 다소 번거로운 작업이 포함되어 있는 경우가 많습니다.

• Cordonnier et al. (2020): 이미지를 2x2의 패치로 쪼갠 후, 이에 self-attention을 적용함. 위와 같은 점에서 ViT와 매우 유사하지만, 이미지 패치가 매우 작으므로 저해상도 입력 이미지에만 적용이 가능하다는 단점이 있음. ViT의 경우에는 중해상도 이미지를 다룰 수 있다는 점, Vanilla Transformer 구조를 차용해 기존의 SOTA CNN보다 더 좋은 성능을 증명해냈다는 점에서 해당 연구보다 우위를 가집니다.
• image GPT (Chen et al., 2020): 이미지 해상도와 color space를 줄인 후, image pixel 단위로 Transformer를 적용한 생성 모델입니다.

ViT의 차별점

• 표준 ImageNet 데이터셋보다 더 큰 크기의 데이터셋에서 image recognition 실험을 진행하였고, 더 큰 크기의 데이터셋에서 학습시킴으로써 기존의 ResNet 기반 CNN보다 더 좋은 성능을 낼 수 있었습니다.

4. Method

Vision Transformer(ViT)

ViT의 모델의 구조

• 우선 input Image를 patch로 분리합니다. 분리된 것을 Linear Projection 한 것과 Position Embedding(아래의 그림)을 Transformer Encoder에 넣어줍니다.
• Standard Transformer는 input으로 token embedding의 1D sequence를 받습니다. 따라서 2D image를 다루기 위해 image를 Flattened 2D patch들을 reshape 해줍니다.

$$
\mathbf{x}\in\mathbb{R}^{H\times{W}\times{C}} \space\space\space\to\space\space\mathbf{x}_p\in\mathbb{R}^{N\times{(P^2\cdot{C})}}
$$

(P, P) : the resolution of each image patch

N : the resulting number of patches(패치 수)

• Transformer는 일정한 Latent vector 사이즈 D를 모든 layer에 사용하는데, 이를 통해 패치들을 Flatten 시킨 후, D차원으로 mapping 시킵니다. 이 projection의 output을 patch embedding이라 합니다.
• Classification head는 MLP에 의해서만 수행되는데 pre-training에서는 one hidden layer, fine-tuning에서는 single linear layer가 사용됩니다.
• Position Embedding은 위치 정보를 위해 patch embedding과 함께 추가되는데, standard learnable 1D position embedding이 사용되었습니다. (자세한 설명 Appendix D.4)

• 앞서 말한 것들에 대한 구조들이 수식으로 정리되어 있는 것으로 y가 계산되는 과정입니다.

Hybrid Architecture

• 각 Patch를 linear projection이 아니라 CNN을 활용하여 나온 feature map을 input으로 활용합니다.
• 특별한 경우 spatial size 1X1을 가질 수 있는데, input sequence가 feature map의 spatial dimension을 flatten 시키고, Transformer dimension으로 projetion을 시켜 나온 것을 의미합니다.
• 그리고 Classification input embedding과 position embedding은 위와 같은 방법으로 추가됩니다.

Fine-Tuning and Higher Resolusion

• pre-trained prediction head를 제거하고 zero-initialized D x K feedforward layer를 추가합니다. (K는 해결하고자 하는 task의 class 수 입니다.)
• pre-trained 보다 높은 resolusion의 image로 fine-tuning하는 것이 더 좋다고 합니다. 하지만 이 경우 기존의 Positional Embedding이 의미가 없어지게 됩니다.
• 이 때, Input image의 크기 내 patch 위치에 맞게 positional embedding도 2D interpolation을 적용하여 값을 채워줍니다.

Limitations

• 많은 data 로 pre-train 해야한다는 점 입니다.
• Fine-Tuning 시 Pre-Train 이미지보다 높은 해상도를 사용하면 성능 향상에 도움이 된다고 합니다.

5. Experiments

실험을 위해 ResNet, ViT 그리고 Hybrid(CNN feature map)을 평가했습니다. 다양한 크기의 dataset으로 평가를 진행하였는데 pre-training에서의 비용이 ViT가 매우 효율적이었으며, SOTA 수준의 성능을 달성하였습니다.

5.1 Setup

Dataset

• Pre-training
  • 모델의 스케일 능력을 실험하기 위해 ImageNet dataset, ImageNet-21K(21K 개의 클래스와 1400만 개의 image) JFT(18K 개의 class와 303M의 고해상도 image)를 사용했습니다.
• Transfer Learning
  • 전이 학습이 사용된 dataset으로는 ImageNet, CIFAR 10/100, Oxford-IIIt Pets 등이 있습니다.

Model Variants

• BERT의 구성을 base로 실험을 진행하였는데, ViT는 ‘Base’,’Large’,’Huge’ model로 Table 1을 참고하면 될 것으로 보입니다.

Training & Fine-tuning

• Pre-training
  • Adam Optimizer$(\beta_1=0.9, \beta_2=0.999)$
  • Batch Size = 4096
  • weight decay = 0.1
  • linear learning rate warmup과 decay도 사용하였습니다.
• Fine-tuning
  • SGD에 momentum과 함께 사용하였고, Batch size는 512입니다.

5.2 Comparison to SOTA

우선 ViT variants 모델들 중 ViT-H/14와 ViT-H/16을 CNN모델들과 우선 비교를 했습니다.

위의 그림에서 확인할 수 있듯이 기존 CNN 구조들 높은 정확도를 보여주고 있으며 TPU에서는 낮은 비용이 들었음을 확인할 수 있습니다.

아래의 그림은 Visual Task Adaptation BenchMark로 ViT-H/14 모델이 다양한 task에서 높은 정확도를 보여줍니다.

5.3 Pre-training Data Requirements

그래프에 따르면 ImageNet과 같은 중간 사이즈의 Dataset은 ResNet보다 약간 낮은 수치의 정확도를 보여주었습니다. 이는 Transformer가 CNN에 내재되어 있는 inductive bias가 부족함을 의미합니다. 그래서 중간 사이즈의 dataset은 학습시키기에 충분하지 않음을 확인할 수 있습니다.

5.4 Scaling Study

• ViT는 performance/compute trade-off에서 ReNet보다 성능이 좋습니다.
• Hybrid 모델은 적은 computing cost에서는 ViT보다 성능이 좋지만, Cost를 늘리면 큰 차이가 없어집니다.
• ViT는 Saturate되지 않으며 스케일링이 가능합니다. 이는 모델의 성능이 더 좋아질수 있음을 의미합니다.

5.5 Inspecting Vision Transformer

• Left : Embedding의 Filter를 확인할 수 있으며, 각 Filter의 기능이 저차원의 CNN filter 기능과 유사합니다.
• Center : Position Embedding간의 유사성. 가까운 패치 간의 유사도가 높다는 것은 Input Patch 간의 공간정보가 잘 학습되는 것을 의미합니다.
• Right : Self-Attention을 활용한, 전체 이미지 정보의 통합 가능 여부를 확인합니다.
  • Attention의 weight를 기반으로 Image Space의 평균 거리를 계산
  • “Attention Distance”는 Receptive Field를 의미
  • 낮은 Layer의 Self-Attention Head는 CNN처럼 ‘Localization’효과를 보여줍니다.

5.6 Self-supervision

기존의 NLP task에서 transformer는 BERT 등과 같이 Self-supervision pretext task와 같이 엮어서 많이 사용되고, 좋은 성능을 보여주고 있습니다.

이것과 유사하게, ViT도 Patch를 masking하고 이를 예측하는 방식으로 supervision task를 적용하여 학습시켰습니다.(JFT-300M dataset, batch size= 4096, epoch=14로 학습)

마스킹 비율은 50% 정도로 사전학습 시 수행하고 ImageNet에 fine-tuning 했을 때, ViT-B/16 모델로 79.9%의 정확도를 보여주었습니다.

지도학습에 비해 4%정도 낮은 정확도를 보여줍니다.

6. Conclusion

---

Contibution

• Image recognition task에 Transformer를 직접적으로 적용한 첫번째 사례입니다.
• CNN 기반 모델과 달리 image를 patch로 만들고, NLP에서 사용하던 Transformer Encoder를 사용하였습니다. 이는 대규모 Dataset으로 pre-train 했을 때, 높은 성능을 보여주었습니다.
• ViT는 pre-train 비용이 상대적으로 낮으며, Image Classification Dataset에서 높은 성능을 보여주었습니다.

Challenge

• Segmentation, detection과 같은 Task에도 적용하는 것
• self-supervised와 관련한 pre-training을 지속적으로 연구하는 것
• 모델 개선을 통해 성능을 더욱 향상시키는 것