(2020)AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

0. Abstract

Transformer architecture는 NLP task에서는 표준이 되었지만 computer vison task에서는 아직 제한적이다. Computer vision task에서 attention은 CNN과 함께 적용되거나 전체구조를 유지하면서 CNN의 특성 요소를 대체하는데 사용된다.

본 논문에서는 CNN에 대한 이러한 의존이 필요하지 않으며 image patch의 sequence가 transformer에 적용될 때 image classification task에서 잘 수행될 수 있음을 보여준다.

많은 양의 데이터에 대해 사전 학습을 수행하고 여러가지 recognition benchmark(ImageNet, CIFAR-100, VTAB 등)에 대해 transfer learning을 수행하면 Vision Transformer는 훨씬 적은 computational resource를 가지며 동시에 SOTA인 CNN과 비교하여 더 우수한 결과를 얻을 수 있다.

1. Introduction

computer vision에서는 CNN이 여전히 많이 사용된다. NLP 성공에 영감을 받아 여러 연구들에서 CNN과 유사한 architecture를 self-attention과 결합하려고 시도하며 일부는 CNN을 완전히 대체한다. 이론적으로는 효율적이지만 specialized attention pattern을 사용하기 때문에 최신 하드웨어 가속기에서는 아직 효과적으로 사용하기 어렵다. 따라서 large-scale image recognition task에서 ResNet-like architecture는 여전히 SOTA이다.

NLP의 Transformer 성공에 영감을 받아, 가능한 최소한의 수정으로 Transformer를 image에 직접 적용하는 실험을 한다. 이를 위해 image를 patch로 분할하고 이러한 patch의 linear embedding sequence를 Transformer에 대한 입력으로 feed한다. Image patch는 NLP 애플리케이션의 token(word)과 동일한 방식으로 처리된다.

이러한 model은 ImageNet과 같은 중간 규모의 dataset에서 학습할 때 적당한 결과를 산출하여 비슷한 크기의 ResNet보다 조금 아래의 accuracy를 달성한다. 이는 겉보기에 실망스러운 결과를 예상 할 수 있다. Transformer는 translation equivariance 및 locality와 같은 CNN 고유의 inductive bias가 없기 때문에 불충분한 양의 dataset에 대해 학습할 때 일반화가 잘 되지 않는다.

그러나 large-scale dataset(14M-300M Images)에서 model을 학습하면 결과가 다르다. Large-scale training이 inductive bias를 능가한다는 것을 보여준다. Transformer는 충분한 규모로 pre-train되고 더 적은 dataset으로 transfer learning 될 때 좋은 결과를 얻는다.

JFT-300M 데이터 세트에 대해 사전 훈련 된 Vision Transformer는 여러 Image Recognition Benchmark에서 SOTA에 접근하거나 이를 능가하여 ImageNet에서 88.36%, ImageNet-ReaL에서 90.77%, CIFAR-100에서 94.55%, 77.16%의 accuracy를 달성했다.

2. Related Work

Transformer는 NMT(Neural Machine Translation)을 위한 것이며 많은 NLP task에서 SOTA를 달성했다. Large Transformer-based model은 일반적으로 large-scale corpus에 대해 pre-train을 수행하고 해당되는 task에 fine-tuning을 수행한다. BERT는 denoising self-supervised pre-training task를 사용하는 반면 GPT 계열은 language modeling 방식으로 pre-train을 수행한다.

Self-attention을 image에 naive하게 적용하려면 각 pixel이 다른 모든 pixel에 attention해야 한다. 이는 pixel수의 quadratic cost를 가지고 실제 input size로 확장되지 않는다. 따라서 image generation 측면에서 Transformer를 적용하기 위해 시도된 몇 가지 연구들이 있다.

가장 최근의 연구인 iGPT는 image resolution과 color sapce를 줄인 후 image에 대해 transformer를 적용한다. Model은 unsupervised 방식으로 학습되고 이후 fine-tuning을 수행하거나 linear를 수행하여 ImageNet에서 72%의 accuracy를 달성하였다.

3. Method

1) Vision Transformer(ViT)

Image용 Transformer는 NLP용으로 설계된 architecture를 따르며 위 그림과 같다. Standard Transformer는 token embedding의 1D sequence를 입력으로 받는다.

Image를 처리하기 위해 $\text{x}\in\text{R}^{H\times W\times C}$를 flatten된 2D patch $x_p\in \text{R}^{N\times(P^2\cdot C)}$ sequence로 재구성한다. $(H,W)$는 원본 image의 resolution이고 $(P,P)$는 각 image patch의 resolution이다. $N=HW/P^2$는 Transformer의 sequence length이다. Transformer는 모든 layer에 일정한 width를 사용하므로 학습 가능한 linear projection은 각 vectorized patch를 dimension $D$에 mapping하며, 그 결과를 patch embedding이라고 한다.

BERT의 $[class]$ token과 유사하게 Transformer encoder의 output state $(z_L^0)$가 image representation $y$로 사용되는 embedded patch의 sequence $(z_0^0=x_{class})$ 앞에 학습 가능한 embedding을 추가한다. Pre-train 및 fine-tuning 중에 classification head는 $(z_L^0)$에 추가된다.

Position embedding은 위치 정보를 유지하기 위해 patch embedding에 더해진다.

Transformer Encoder는 Multi-headed self-attention 및 MLP block으로 구성된다. Layer norm(LN)은 매 block 이전에 적용되고 residual connection은 이후에 적용된다.

$\text{z}_0=[\text{x}_{\text{class}};\ \text{x}_p^1\text{E};\ \text{x}_p^2\text{E};\ \cdots\ ;\ \text{x}_p^N\text{E}]+\text{E}_{pos},\ \text{E}\in\text{R}^{(P^2\cdot C)\times D},\ \text{E}_{pos}\in\text{R}^{(N+1)\times D}$

$\text{z}_l'=\text{MSA}(\text{LN}(\text{z}_{l-1}))+\text{z}_{l-1},\ l=1\ldots L$

$\text{z}_l=\text{MLP}(\text{LN}(\text{z}_l'))+\text{z}_l',\ l=1\ldots L$

$\text{y}=\text{LN}(\text{z}_L^0)$

Image를 patch로 나누는 대신 ResNet의 중간 feature map에서 input sequence를 형성할 수 있다. Hybrid Model에서 patch embedding projection $\text{E}$는 ResNet의 early stage로 대체된다.

ResNet의 중간 2D feature map중 하나는 sequence로 flatten되고 transformer dimension에 projection된 다음 transformer의 input sequence로 feed된다.

Classification input embedding 및 position embedding은 위에서 설명 하듯이 transformer에 대한 input에 추가된다.

2) Fine-tuning and Higher Resolution

일반적으로 large-scale dataset에 대해 ViT를 pre-train하고 downstream task에 대해 fine-tuning을 수행한다. 이를 위해 pre-trained prediction head를 제거하고 0으로 초기화된 $D\times K$ feedforward layer를 추가한다. 여기서 $K$는 downstream class의 개수이다.

Pre-train 보다 높은 resolution으로 fine-tuning하는것은 종종 도움이 된다. 더 높은 resolution의 image를 feed할 때 patch 크기를 동일하게 유지하므로 sequence length가 더 길어진다. Vision Transformer는 임의의 sequence length를 처리할 수 있지만 pre-trained position embedding은 의미가 없을 수 있다. 따라서 원본 image에서의 위치에 따라 pre-trained position embedding의 2D interpolation을 수행한다. Resolution 조정 및 patch 추출은 image의 2D 구조에 대한 inductive bias가 Vision Transformer에 수동으로 주입 되는 유일한 지점이다.

4. Experiments

ResNet, Vision Transformer(ViT) 및 Hybrid에 대해 representation learning 실험을 수행한다. 다양한 크기의 dataset에 대해 pre-train하고 benchmark를 수행한다.

pre-training computational cost을 고려할 때 ViT는 더 낮은 cost로 대부분의 recognition benchmark에서 SOTA를 달성하였다.

1) Setup

Datasets.

- Pre-train dataset

• ILSVRC-2012 ImageNet dataset(ImageNet) - 1k class 및 1.3M images
• ImageNet-21k - 21,000 class 및 14M images
• JFT - 18k class 및 303M image

- Transfer Learning dataset

• ImageNet 및 ReaL labels
• CIFAR 10/100
• Oxford-IIIT Pets
• Oxford Flowers-102
• 19-task VTAB classification suite

Model Variants.

• BERT에 사용되는 구성을 기반으로 함.
• 접미사는 “B”(Base), “L”(Large), “H”(Huge) 를 뜻함.
• 예를 들어 ViT-L/16 의 경우는 Large 사이즈 이며 16x16 의 patch크기를 가짐.

Training & Fine-tuning.

- Pre-train

• Adam optimizer, $\beta_1=0.9,\ \beta_2=0.999,\ batch\ size=4096$
• weight decay: 0.1

- Fine-tuning

• SGD w/ momentum, $batch\ size=512$
• Using linear learning rate warmup and decay
• Higher resolution: 512 for ViT-L/16 and 518 for ViT-H/14

2) Comparison to State of the Art

모든 model은 TPUv3에서 학습되었으며 pre-train에 소요된 TPUv3-core-days를 확인할 수 있다.

ViT-L/16 model은 모든 dataset에서 BiT-L과 일치하거나 더 좋은 성능을 보여줌과 동시에 훨씬 적은 computational resource를 필요로 한다. 더 큰 model인 Vit-H/14는 ImageNet 및 CIFAR-100과 VTAB에서 성능이 더욱 향상되었다.

VTAB task에 대해 각각의 그룹으로 분할하고 이전 SOTA와 비교하였다.

3) Pre-training Data Requirements

ViT는 large-scale JFT-300M dataset에서 pre-train 하였을 때 좋은 성능을 보여준다. ResNet보다 vision에 대한 inductive bias가 적을 때 dataset의 크기가 얼마나 중요한지 실험을 수행한다.

가장 작은 dataset에 대해 pre-train을 수행한 경우 ViT-Large model은 normalization에도 불구하고 ViT-Base보다 성능이 떨어진다. 그러나 ImageNet-21k dataset을 사용하면 성능이 비슷하다. JFT-300M dataset에서는 ViT가 BiT보다 더 좋은 성능을 보여줌을 확인할 수 있다.

결과적으로 더 작은 dataset의 경우 convolutional inductive bias가 유리하지만 큰 dataset의 경우 관련 패턴을 학습하는 것이 충분하고 이점이 있다고 판단할 수 있다.

4) Scaling Study

JFT-300M dataset에서 transfer performance에 대해 다양한 model의 확장 연구를 수행한다. Dataset 크기는 model 성능에 병목 현상이 없으며 각 model의 pre-train cost 대비 성능을 평가한다. 위 그림은 transfer performance vs pre-training compute에 대해 나타나있다.

• ViT는 performance/compute trade-off 에서 ResNet보다 효율적이다. ViT는 동일한 성능을 달성하기 위해 약 2배 적은 computational cost가 들어감.
• Hybrid는 적은 computational cost에서 ViT를 능가하지만 큰 computational cost에서는 차이가 사라진다.
• ViT는 시도된 범위 내에서 saturate되지 않는 것으로 보이며 향후 확장이 가능해 보인다.

5) Inspecting Vision Transformer

ViT가 image를 처리하는 방법을 이해하기 위해 분석을 수행한다. ViT의 첫 번째 layer는 flatten patch를 더 낮은 차원의 space에 projection한다.

위 그림에서 왼쪽은 학습된 embedding filter의 구성 요소를 보여준다. 구성 요소는 각 patch내 미세 구조의 low-dimensional representation에서 그럴듯한 basic function과 유사하다.

Projection 이후 학습된 position embedding이 patch representation에 추가된다. 위 그림에서 가운데는 model이 position embedding의 유사성에서 image 내 거리를 encoding하는 방법을 학습함을 보여준다. 즉, 더 가까운 patch는 더 유사한 position embedding을 갖는 경향이 있으며 행-열 구조가 나타난다.

Self-attention을 통해 ViT는 가장 낮은 layer에서도 전체 image에 대한 정보를 통합할 수 있다. Self-attention의 weight를 기반으로 정보가 통합되는 image 공간의 평균 거리를 계산한다. 위 그림에서 오른쪽 attention distance는 CNN의 receptive field size와 유사하다.

일부 head는 이미 최하위 layer에 있는 대부분의 image에 attention을 하여 global하게 model에서 사용됨을 보여준다. 또한 network 깊이에 따라 attention distance가 증가한다. model이 의미 상의 분류와 관련된 image 영역을 담당한다는 것을 알 수 있다.

6) Self-supervision

Transformer는 NLP task에서 인상적인 성능을 보여주었다. 그러나 대부분의 성공은 확장성 뿐만 아니라 self-supervised pre-training에서 비롯된다. 또한 BERT에서 사용되는 MLM(Masked Language Modeling)을 모방하여 self-supervision을 위한 Masked Patch Prediction에 대한 preliminary exploration을 수행한다.

Self-supervised pre-training을 수행한 ViT-B/16 model은 ImageNet에서 79.9%의 정확도를 달성하여 scratch로 부터 train을 수행한 것보다 2%가 향상되었지만 supervised pre-training보다는 4% 떨어졌다.

5. Conclusion

Image recognition에서 Transformer를 직접 적용하는 방법을 제안했다. Computer Vision에서 Self-attention을 사용하는 이전 연구들과 달리, 본 논문에서는 architecture에 image-specific inductive bias를 사용하지 않았다. 대신 이미지를 patch로 해석하고 NLP에서 사용되는 standard transformer encoder로 처리한다.간단하면서도 확장 가능한 전략은 large-scale dataset에 대한 pre-train과 결합될 때 좋은 결과를 보여줬다. 따라서 Vision Transformer는 많은 image classification dataset에서 SOTA를 능가하거나 동시에 pre-train에 대한 부담도 적다.

이러한 결과는 고무적이지만 많은 challenge가 남아있다. 첫째, detection 및 segmentation과 같은 computer vision task에 ViT를 적용하는 것이다. 둘째, pre-training method에 대해 연구하는 것이다. 초기 실험에서 self-supervised pre-training이 개선된 점을 보여주었지만, supervised pre-training을 능가하지만 못하였다. 셋째, model의 크기가 증가하여도 saturate 상태가 아닌 것으로 보이기 때문에 ViT를 더 확장할 필요가 있다.