An Image is worth 16*16 words: Transformers for image recognition at scale

https://arxiv.org/pdf/2010.11929

 

ViT논문 드디어 리뷰한다!

 

Abstract

transformer는 자연어 처리에서 많이 사용되는 중, 컴퓨터 비전에서는 제한적 사용

주로 vision에선 attention은 conv net과 함께 적용되거나 전체 구조를 유지하면서 conv net의 특정 구성 요소를 대체하는 데 사용

여기서는 CNN에 의존하지 않고 이피미 패치 시퀀스에 직접 transformer를 사용해서 이미지 분류 작업에 우수한 성능을 보인다는 것을 보임

ImageNet, CIFAR-100, VTAB 등 여러 중형~소형 이미지 벤치마크로 ViT가 훨씬 적은 자원으로 CNN에 비해 우수한 결과를 내는 것을 확인했음

1. Introduction

self-attention 기반 아키텍처(특히 Transformer)은 NLP에서 주로 사용됐음

일반적으로는 대규모 pretrain 후, finetuning하는 방식

transformer의 계산 효율성과 확장성 덕에 100억 개가 넘는 파라미터로 전례없는 규모의 모델을 훈련할 수 있게 됨

모델과 데이터 셋이 증가함에 따라 성능의 한계가 안 보이는 지경

하지만 vision에서는 여전히 CNN이 지배적

NLP에서의 성공에 영감받아 CNN과 attention을 결합하려고 시도했으나 아직 완벽하게 확장되진 않았음

여전히 ResNetlike 아키텍처가 최신기술

최소한의 수정으로 표준 Transformer를 이미지에 직접 적용해보는 실험을 진행

이를 위해 이미지를 패치로 분할하고 이러한 패치의 linear embedding sequence를 transformer에 입력으로 제공

이미지 패치 =  NLP에서 토큰 과 같은 취급

지도 방식으로 이미지 분류 모델 훈련

강력한 regularization없이 ImageNet정도의 중간 규모의 데이터셋에 대해 학습할 경우 ResNet보다 낮은 수준의 정확도를 제공->transformer는 CNN 특유의 inductive bias(translation equivariance, locality)가 부족하기 때문

하지만!대규모 데이터셋에 대해서는 transformer의 inductive bias가 CNN을 뛰어넘는다!

ViT는 대규모 데이터셋으로 pretrain 후, data point가 적은 작업에 적용할 때 우수한 결과를 냈음

ImageNet-21K, JFT-300M 데이터셋에 대해 pretrain된 경우, ViT는 SOTA

2. Related Work

Transformer는 기계 번역을 위해 제안됐음

ex)BERT(denoising self-supervised pretraining task), GPT(language modeling as its pretraining task)

vision에서 self-attention을 naive하게 적용하려면 각 픽셀이 다른 모든 픽셀에 attention해야함

->픽셀 수를 생각하면 비용이 만만치 않다

과거 vision에 transformer를 적용하려는 시도들?

1. 각 쿼리 픽셀에 대한 local neighborhoods에만 self-attention 적용

- local multi-head dot-product self attention block은 CNN을 완전히 대체할 수 있었음

2. Sparse Transformer : vision에 적용하기 위해 global self-attention에 대한 scalable approximation을 적용

3. scale attention을 확장하는 방법 중 하나는 다양한 크기의 블록으로 적용

- 개별 축을 따라 적용하는 것

=>전체적으로 좋은 결과를 보여주긴 하나, hw에서 효율적으로 구현하기 복잡함

 

해당 논문에서 발표한 모델(ViT)과 가장 유사한 모델은 image에서 2*2크기의 패치를 추출하고 full self-attention을 적용하는 Cordonnier 등의 모델(2020)

->이 모델은 ViT와 매우 유사하지만 이 논문에서는 대규모 pretrain을 통해 vanila transformer가 최신 CNN과 동등하거나 그보다 더 낫다는 걸 입증

+ 이 모델은 2*2 픽셀의 작은 패치를 사용하므로 저해상도 이미지에만 적용할 수 있었음

하지만 ViT는 중간 해상도 이미지도 처리 가능

+이미지 분류를 위해 feature map을 증강하거나, self-attention을 사용하여 CNN의 출력을 추가로 처리하는 등

=>CNN과 self-attention과의 결합에 많은 관심

 

해당 논문에서 발표한 모델(ViT)과 가장 유사항 모델로 iGPT가 있음

iGPT?이미지 해상도와 색 공간을 줄인 후, 픽셀에 transformer를 적용하는 모델

->비지도 방식으로 생성 모델을 학습, 그 결과 representation은 분류 성능을 위해 finetuning되거나 linear probe되어 성능 높임

이 후 대규모의 이미지 데이터셋을 활용한 논문이 증가하는 추세

=>이 논문에서는 대규모 이미지 데이터셋으로 ResNet기반 모델 대신 Transformer 훈련!

3. Method

모델은 최대한 Transformer를 따르려고 함

이렇게 하면 확장 가능한 NLP transformer 아키텍처와 효율적인 구현을 거의 바로 사용할 수 있다는 점이 장점

3.1 Vision Transformer(ViT)

Model Overview

 

표준 transformer는 token embedding의 1D sequence를 입력으로 받음

2D image를 처리하기 위해 이미지를 flatten 2D patch sequence로 재구성

->(H, W)는 이미지 해상도, C는 채널 수, (P, P)는 각 패치의 해상도, N=HW/P^2는 patch의 수(=transformer의 입력 시퀀스의 길이), D는 latent vector(모든 레이어에서 일정한 크기)

->flatten patch+map to D dimensions(with lineasr projection)=>출력 결과를 patch embedding이라고 함

BERT의 [class] token과 유사하게 학습가능한 임베딩을 임베딩된 패치 시퀀스에 미리 추가하여 transformer encoder의 출력 상태가 이미지 representation y로 사용

pretraining, finetuning 둘 다 classification head는 z에 부착

classification head는 pretraining때는 하나의 hidden layer와 finetuning때는 하나의 linear layer가 있는 MLP에 의해 구현됨

position embedding은 위치 정보를 유지하기 위해 patch embedding에 추가됨

*여기서 고급 2D 인식 위치 임베딩을 사용해도 성능이 크게 향상되지 않아, 표준 학습 가능한 1D 위치 임베딩 사용

embedding vector의 결과 sequence는 encoder의 input으로 사용

transformer encoder는 multihead self-attention과 MLP block의 교대 layer로 구성

LayerNorm(LN)은 모든 블록 앞에 적용되고, residual connection은 모든 블록 뒤에 적용

MLP는 GELU nonlinear를 가진 두 개의 layer포함

Inductive bias

ViT는 CNN보다 image specific inductive bias가 훨씬 적음

CNN에서는 locality, 2D neighborhood 구조, translation equivariance가 전체 모델의 각 레이어에 영향

ViT에서는 MLP레이어만 그렇고, self-attention layer는 global layer

2D neighborhood 구조는 모델 초기에 이미지를 patch로 자르고 해상도가 다른 이미지의 position embedding을 조정하기 위한 fine tuning시에 매우 드물게 사용됨, 이 외에는 초기화 시점의 position embedding에는 patch의 2D position에 대한 정보가 없으며 patch간 모든 공간 관계를 처음부터 학습해야함

Hybrid Architecture

raw image patch의 대안으로 input sequence는 CNN의 feature map에서 형성가능

patch embedding projection E가 CNN feature map에서 추출한 patch에 적용됨

특수한 경우 patch가 1*1의 크기를 가질 수 있는데 이는 feature map의 dimension을 간단히 flatten시키고 transformer dimension으로 projection해서 input sequence를 얻을 수 있음

classification input embedding과 position embedding은 위에서 말한 대로 추가됨

3.2 Fine-Tuning and Higher Resolution

일반적으로 대규모 데이터셋에 대해 ViT를 pretrain하고 downstream task에 맞게 fine tuning함

이를 위해서는 pretrain된 prediction head를 제거하고 zero initialize된 D*K feedforward layer에 연결(K는 downstream class의 수)

pretrain보다 더 높은 해상도로 finetuning하는 것이 유리한 경우가 많음

더 높은 해상도의 이미지를 제공할 때 patch크기를 동일하게 유지하면 유효 sequence 길이가 더 길어짐

ViT는 임의의 sequence 길이까지 처리가능하지만 pretrain position embedding이 더 이상 의미가 없을 수도 있음

->원본 이미지의 위치에 따라 pretrain된 position embedding의 2D interporlation을 수행

해상도 조정+패치 추출->이미지의 2D 구조에 대한 inductive bias가 ViT에 수동으로 주입되는 유일한 지점

4. Experiments

4.1 Setup

Datasets