Vision Transformer: “An Image is Worth 16x16 Words”

Abstract

NLP에서는 Transformer 구조가 사실상의 표준으로 자리잡았으나, CV에서는 응용이 제한적

CNN에 결합해 사용하거나, CNN의 기본 골격을 유지한 채 일부분을 transformer로 대체하는 역할 정도에 그침

본 논문은 이러한 CNN에 대한 의존이 필수적이지 않으며, image patch에 직접 pure transformer를 적용해도 classification 문제에 좋은 성능을 보임을 확인

Vision Transformer(ViT)는 적은 계산량에도 불구하고 여러 벤치마크에서 state-of-the-art CNN에 비교해서도 좋은 성능을 보임

Introduction

Transformer를 위시한 self-attention 기반 아키텍처는 NLP의 주류로 사용되고 있음, 일반적으로 거대 corpus로 훈련 후 fine tuing하는 방식으로 사용

Transformer의 계산효율성과 scalability덕분에 전례없던 크기의 (100B 이상) 파라미터 수를 가진 모델들도 훈련이 가능해짐. model과 dataset이 점점 커짐에도 성능이 saturate될 기미를 보이지 않음

그러나 CV에서는 아직 CNN 기반 아키텍처가 주류. self-attention을 결합하려는 시도는 있었으나 ResNet-like architecture들이 여전히 SOTA의 자리를 차지

본 연구에서는 NLP에 영감을 받아 image에 최소한의 수정을 가한 standard Transformer를 적용해 실험.

이미지를 패치로 쪼갬 → linear embedding의 sequence → Transformer의 입력으로 사용.
Image patch가 NLP의 token과 같이 다뤄짐. 지도학습으로 훈련

ImageNet과 같은 mid-size dataset을 strong regularization없이 훈련시켰을 때에는 비슷한 크기의 ResNet보다 수 퍼센트 낮은 성능.

Transformer는 CNN에 내재된 “inductive bias”가 없음. (translation equivariance, locality, …) → 충분한 데이터가 없으면 쉽게 generalize되지 않음

그러나 큰 데이터셋으로 (14M - 300M images) 훈련시 이를 극복 가능.

sufficient scale로 pre-training 후 transfer시키면 좋은 성능
이미지 인식 벤치마크들에서 SOTA 이상의 성능

Transformer는 NLP에서 매우 널리 사용, large corpora에서 pre-train 후 fine-tune하여 사용하는 경우 많음.

BERT : denoising self-supervised pretraining task
GPT: language model을 pretraining task로

Self-attention을 naive하게 적용시 pixel 수의 제곱에 비례 → 비현실적

Transformer을 image processing 과정의 중간에서 적용하려는 많은 시도들
- e.g. Parmar et al. (2018): 각 query pixel의 근방에 있는 픽셀들에 대해서만 self-attention을 적용 (local multi-head dot product self-attention block) → convolution을 완전히 대체 가능
- 좋은 성능을 보였으나, HW 가속기상에서 효율적으로 구현하려면 복잡한 engineering 요구

본 연구와 가장 연관성 높은 모델: Cordonnier et al. (2020)

input image에서 2x2 patch를 추출하여 full self-attention을 적용
ViT와 매우 유사하지만 본 연구에서는 여기에 더해 large scale pre-training으로써 vanilla transformer가 SOTA CNN 이상으로 경쟁력있어진다는 점을 보여주려 함.
Cordonnier et al. (2020)은 2x2 patch 때문에 small resolution에만 적용가능했다는 한계 극복.

본 논문과 깊은 연관을 가진 또 다른 예시: image GPT(iGPT). image resolution과 color space를 축소한 후, Transformer를 image pixel에 적용.

비지도학습, generative model로써 훈련 후 fine tuning / probe linearly

기본 ImageNet보다 dataset size가 커지면(ImageNet-21k, JFT-300M) CNN의 성능이 얼마나 더 좋아지는지에 대해서는 이미 연구가 많이 진행 → 이러한 데이터셋을 이러한 연구들에서 사용된 ResNet 기반 모델 대신 Transformer에 적용

Method

최대한 original Transformer에 가깝도록 설계됨

Vision Transformer (ViT)

2D image의 처리를 위해서,

$\mathbb{R}^{H\times W\times C}$ 모양의 이미지를 $\mathbf{x}_p \in \mathbb{R}^{N \times (P^2 \cdot C)}$로 reshape ($P\times P$: patch의 resolution, $N=HW/P^2$: patch의 개수)
모든 layer에서 latent vector size는 $D$로 같음
- 이를 위해서 input에서는 patch를 flatten → traininable linear projection을 통해 D차원으로 바꿔줌 (patch embedding)
patch embedding에 1D positional embedding을 더해줌
- 2D-aware positional embedding이 더 큰 효과가 있지는 않음
Transformer Encoder는 Attention is All You Need에서와 똑같은 형태. MLP에는 GELU(Gaussian Error Linear Unit) non-linearity가 적용

Inductive Bias

ViT는 CNN보다 image-specific inductive bias가 훨씬 약함 (locality, translation equivariance). ViT에서 local, translationally equivariant한 구조는 MLP layer밖에 없음. self-attention은 global
- locality: 2D상에서 인접해 있는 구조
- positional embedding도 2D에 대한 정보를 담지 않으므로 patch간의 spatial relation에 대한 정보는 0에서부터 배우게됨

Hybrid Architecture

raw input patch 대신에, CNN의 feature map을 image sequence로 넣는 방법도 있음

Fine-Tuning and Higher Resolution

일반적으로 ViT를 large dataset에서 훈련 후 smaller downstream task에 적용

이를 위해서는 pre-trained prediction head를 제거한 후 zero-initialized $D\times K$ FF layer를 추가 (K: downstream class의 개수)

더 높은 해상도의 이미지에 적용하고 싶을 때는 patch size를 유지, sequence length가 늘어나도록 함. (임의의 sequence length에 대해 적용 가능)

positional embedding의 (원본 이미지상의 위치에 따른)2D interpolation을 사용
이는 image에 대한 inductive bias가 ViT에 인위적으로 주입된 유일한 측면임

Experiments

ResNet, ViT, hybrid의 representation learning 능력을 평가.

Setup

Dataset

Pretraining: ImageNet (1k classes, 1.3M images), ImageNet-21k (21k classes, 14M images), JFT (18k classes, 303M high-resolution images)
이후 여러 benchmark들로 transfer
- Natural(pets, CIFAR) / Specialized(medical, satelite) / Structured(localization과 같은 geometric understanding이 필요한 작업들)

Model Variant

BERT에서 사용된 configuration에 기반해 설정. B / L / H

Baseline CNN으로는 ResNet을 사용하되, BatchNorm을 GroupNorm으로 대체, standardized convolution을 사용
- transfer 성능을 높임 → ResNet (BiT)로 부름 (Big Transfer)
- BiT: Big Transfer
Hybrid model의 경우 intermediate feature map을 patch size 1(pixel)로 ViT에 넣어줌.

Training & Fine-tuning

Adam, batch size = 4096, weight decay = 0.1, linear lr warmup & decay
Fine tuning시: SGD with momentum, batch size = 512,

Metrics

downstream dataset에 대한 결과를 few-show과 fine-tuning accuracy로 보고.
- few-shot accuracy: train image의 부분집합으로부터 얻은 feature map으로 regularized least-square regression을 통해 얻음

Comparison to State of the Art

비교 대상

large ResNet에 supervised transfer learning을 하는 BiT → ImageNet 제외 SOTA
Noisy Student: ImageNet과 JFT-300M에 label을 지워서 EfficientNet을 semi-supervised로 훈련 → ImageNet의 SOTA

ViT-L/16 model을 JFT-300M dataset으로 사전학습시킨 것이 모든 태스크에서 BiT-L을 이기면서도 computing 자원은 덜 차지

더 큰 모델인 ViT-H/14는 (특히 더 어려운 dataset에서) 더 높은 성능

Pre-training Data Requirements

데이터셋 크기에 따른 성능 비교. ResNet이 가진 inductive bias 없이 훈련시켰다는 특성이 있으므로 dataset size가 어느 정도로 중요한지를 따져봄

앞에서 매우 큰 사이즈의 JFT-300M dataset으로 pretraining 시켰을 때는 좋은 성능을 보이는 것을 확인.

ImageNet < ImageNet-21k < JFT- 300M
각각을 통해 pre-training 시킨 후 ImageNet으로 fine-tuning. 각각에 대해서 regularization(weight decay, dropout, label smoothing)을 최적화
ImageNet에서는 ViT-B > ViT-L. ImageNet-21k에서는 비슷한 성능. JFT-300M에서는 ViT-B < ViT-L. 즉, 큰 모델의 장점은 dataset도 커야 살릴 수 있음

JFT-300M의 부분집합 (9M, 30M, 90M)으로 pretranining

regularization을 최적화하지 않고 같은 hyperparameter 유지
early stopping으로 best val accuracy를 저장 → few-shot linear accuracy 측정
작은 dataset에서는 ResNet과 비슷한 계산자원을 먹을 때 ResNet보다 overfitting이 심함. 반면 큰 dataset에서는 높은 성능
즉, 작은 데이터셋에서는 CNN의 inductive bias가 유용하게 사용되지만 데이터셋 크기만 커지면 data로부터 이러한 pattern을 직접 학습하는 ViT의 방식도 충분함

Scaling Study

model size에 따른 transfer performance의 비교. JFT-300M에서 훈련

performance / training time trade-off에서, ViT가 ResNet보다 우위
computation budget이 충분하지 않을때에는 Hybrid > ViT. 그러나 모델이 커지면 차이는 없어짐
- convolutional local feature processing이 ViT를 도와줄거라는 일반적인 예상과 반대
ViT는 실험에서 사용된 범위 내에서는 scaling에 따라 saturate가 되지 않았음. → 더 model size를 키우면 성능이 개선될 여지

Inspecting Vision Transformer

ViT가 어떻게 이미지를 분석하는지 이해하기 위해, internal representation을 분석

ViT의 첫 번째 layer는 flattened patch를 low-dimensional space로 projection

그림은 learned embedding fiter의 시각화

Projection 후에는 learned positional embedding이 patch representation에 더해지게 되는데, 이는 image 내 patch의 위치를 자동으로 학습

서로 가까운 patch들이 비슷한 positional embedding, 같은 row/column에 있는 패치들끼리도 유사성을 나타냄
larger grid에서는 sinusoidal structure가 나타남
이렇게 2D structure를 자동으로 학습하기 때문에, 2D-aware embedding은 효과가 없음

self-attention은 ViT가 가장 낮은 layer에서부터 모든 정보를 통합적으로 사용하도록 함

image space 상에서 image의 얼마나 멀리 떨어진 정보를 같이 사용하는지 average distance를 계산 (attention distance) → CNN의 receptive field size와 유사한 역할
attention head 중 일부가 이미지 전체적인 정보를 사용하고, 일부는 highly localized → convolutional layer와 비슷한 역할
classification과 semantically relevant한 영역들에 attention

Self-Supervision

트레이닝에 필요한 데이터가 너무 많기 때문에, 일일이 라벨링을 하기 힘들 수 있음 → BERT와 같이 masked patch prediction을 통해 pre-training하는 방법

supervised pretraining에 비해서는 4% 낮은 정확도

Conclusion

Transformer의 image recognition으로의 응용 방법을 살펴봄

CV에서 self-attention을 도입한 선행연구와 달리, patch extraction을 제외하면 inductive bias 없이 학습. NLP의 Transformer에서 변형을 최소화

large dataset에서의 pretraining을 통해 SOTA를 상회하는 결과. pretraining 또한 상대적으로 낮은 비용으로 가능

향후 연구 과제