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

1. Introduction

Self-attention 기반 구조는 NLP 분야에서 많이 사용되어 왔다. 가장 지배적인 방식은 대규모 text corpus에서 사전 학습하고, 작은 task-specific 데이터셋에서 fine-tuning하는 BERT 방식이다. Transformer의 계산 효율성(computational efficiency)과 확장성(scalability) 덕분에, 전례 없는 크기의 모델을 학습할 수 있게 된다.

⭐ fine tune : 기존 학습된 모델 기반으로 변형하고 이미 학습된 모델 가중치로 학습을 업데이트하는 방법

컴퓨터 비전 분야에서는 여전히 Convolutional 구조가 우세하게 사용되고 있다. NLP에서의 성공으로부터 영감을 받아, 다양한 연구에서 CNN과 self-attention을 결합하고자 해왔다.

Replace the convolutions entirely (Ramachandan et al., 2019; Wang et al., 2020a)

최근 모델들은 이론상으로는 효과적이었지만, 특수한 attention 패턴 때문에 현대의 하드웨어 가속기에서 효과적인 스케일링을 기대하기란 어려웠다.

NLP에서의 Transformer의 스케일링이 성공한 것에서부터 영감을 받아, 이미지에 직접적으로 표준 Transformer를 적용하는 방안을 실험해냈다. 이를 위해서,

1) 이미지를 패치로 분할하고, (split an image into patches)

2) 이러한 패치들의 linear embeddings sequence를 Transformer의 input으로 전달한다.

이미지 패치는 마치 NLP에서의 ‘토큰’처럼 여겨지며, supervised 방식으로 이미지 분류에 학습을 시켜 실험하였다.

ImageNet과 같은 중간 크기의 데이터셋으로 regularization 없이 학습했을 때, 이 모델은 ResNet보다는 약간 낮은 정확성 수치를 보였다. 즉, Transformer에게 CNN에 내재되어 있는 inductive biases(translation equivariance, locality)가 부족함을 의미한다.

이 때 translation invaraince란, translation invariance의 반대 개념으로, 입력 위치가 변함에 따라 출력이 변한다는 의미이다.
locality = 국소성 (작은 지역에 한정됨)

이 때문에 중간 사이즈의 데이터셋은 이 모델을 학습시키기에 충분하지 않으며 일반화 성능도 좋지 못하다.

그러나, 모델이 더 큰 데이터셋 (14M-300M개 이미지)을 학습하게 되면 다른 양상을 띤다. 대규모 학습은 inductive bias를 찍어눌러 버린다.

inductive bias : 주어지지 않은 입력의 출력을 예측하는 것.

Vision Transformer (ViT)는 충분한 규모에서 사전학습되고 task로 전이 학습될 때 훌륭한 결과를 보여준다.

(ViT를 사용하면 image-specific inductive biases를 주입 않고도 SOTA 달성 가능)

예를 들어,ImageNet-21k 데이터셋이나 in-house JFT-300M 데이터셋에서 사전학습되었을 때, ViT는 최신 멀티 이미지 인지 기술을 뛰어넘는다. 정확한 수치로는, ImageNet에서는 88.55%, ImageNet-Real에서는 90.72%, CIFAR-100에서는 94.55%, VTAB에서는 77.63%의 정확성을 보인다.

2. Related Work

self-attention을 이미지에 단순 적용하기 위해서는, 각 픽셀이 모든 픽셀에 attend할 것을 요구한다. 이러한 방식은 픽셀 수에 대해 quadratic한 복잡도를 가지게 되고, 이는 현실적인 입력 사이즈에 확장될 수 없다. 그리하여 이미지 처리 과정에서의 Transformer를 적용하기 위해서 과거에는 근사(approximation)를 사용하였다.

Parmar et al.(2018)은 globally(전역적으로) 적용시키는 대신 각 쿼리 픽셀의 local neighborhood에만 self-attention을 적용시켰다.
위와 같은 local multi-head dot-product self attention block은 convolution을 완전히 대체할 수 있다. (Hu et al., 2019; Ramachandran et al., 2019; Zhao et al., 2020)
Sparse Transformers는 이미지에 적용하기 위해 global self-attention에 대한 scalable approximations을 이용한다. (Child et al., 2019)
block의 다양한 사이즈를 적용해 attention을 확장한다. (Weissenborn et al., 2019)

이런 많은 specialized attention 구조는 컴퓨터 비전 문제의 유망한 결과를 보여주지만, 하드웨어 가속기에서 효율적으로 구현하기 위해 복잡한 엔지너링이 필요로 하다.

가장 유사한 모델은 Cordonnier et al. (2020)의 모델이다. 공통적으로, 입력 이미지로부터 $**2\times 2$ 사이즈의 패치를 추출**하고 top에 full self-attention을 적용한다. 그러나 저해상도 이미지에만 적용 가능한 Cordonnier et al.의 모델과는 달리, 이 논문에서 다룰 모델은 medium-resolution 이미지까지도 다룰 것이다.

(중략)

3. Method

모델 설계 시, Vaswani et al이 제시한 기존의 Transformer을 최대한 수정 않고 그대로 적용할 수 있도록 하였다.

3.1. Vision Transformer (ViT)

표준 Transformer는 1차원 token embedding sequence를 입력으로 받는다. 2D 이미지를 다루기 위해서, 이미지 $\mathbf{x} \in \R^{H\times W \times C}$를 2차원 패치 $\mathbf{x}_p \in \R^{N\times (P^2 \cdot C)}$로 flatten한다.

이 때, $(H,W)$는 원본 이미지의 해상도이고, $C$는 채널 개수, $(P,P)$는 각 이미지 패치의 해상도이다. $N=HW/P^2$는 패치의 최종 개수로, Transformer Encoder에 입력되는 유효 시퀀스 길이로써 바라볼 수 있다. Transformer 내부의 모든 layer에서 latent vector의 크기가 $D$로 통일되어 있어서 trainable linear projection을 거쳐 $D$차원으로 매핑한다. 이 논문에서는 이 projection의 결과를 patch embedding이라 할 것이다.

BERT의 [class] 토큰과 유사하게, 임베딩된 패치들의 맨 앞에 학습 가능한 하나의 class token embedding vector를 추가한다. ($\mathbf{z}^0_0 = \mathbf{x}_{\mathsf{class}}$) 이 임베딩벡터는 Transformer encoder의 출력 $(\mathbf{z}_L^0)$으로 나올 때 방정식 4번에서처럼 이미지 표현 벡터 $\mathbf{y}$로의 역할을 수행하게 된다. pre-training (사전 학습)과 fine-tuning 둘 다 하는 동안, classification head는 $\mathbf{z}_L^0$, 즉, 이미지 표현 벡터에 부착된다.

위치 정보를 유지하기 위해서 patch embedding에 position embedding을 추가한다. 이 때, 표준 학습 가능한(learnable) 1D position embedding을 사용한다. 그 이유는, 2D position embedding을 사용했어도 더 나은 성능을 보이지 않았기 때문이다. 이 임베딩 벡터의 결과 시퀀스가 인코더의 입력으로 전달된다.

Transformer encoder는 multiheaded self-attention (MSA) layer들과 MLP block으로 이루어져 있다. Layernorm(LN)이 매 block 이전에 적용되며, residual connection은 block 이후에 적용된다.

1) Inductive bias

Vision Transformer (ViT)는 CNN보다도 image-specific inductive bias가 훨씬 적다. ⭐

CNN에서는 locality, 2D neighborhood structure, translation equivariance가 전체 모델에 거쳐 각 레이어에 존재했었다. 그러나 ViT에서는 오직 MLP 레이어만이 local하고 translation equivariance가 있으며, self-attention은 global하다. 2D neighborhood structure는 매우 적게 쓰인다. :

• 모델 처음에, 이미지를 패치로 나눌 때
• fine-tuning 시 서로 다른 해상도의 이미지에 대해 position embedding을 적용할 때

2) Hybrid Architecture

Raw image 패치를 대신해 CNN의 feature map으로부터 입력 시퀀스를 얻어낼 수 있다. 하이브리드 모델에서는, patch embedding projection $\mathbf{E}$가 CNN feature map으로부터 추출한 패치에 적용된다. 즉, 패치가 $1\times 1$ spatial size를 갖는데, 이는 input sequence가 단순히 feature map의 차원을 flatten하고 Transformer 차원으로 투영함으로 얻어질 수 있다는 것을 의미한다. 동일하게, classification input embedding과 position embedding이 추가된다.

3.2. Fine-Tuning and Higher Resolution

ViT를 큰 규모의 데이터셋에서 사전학습하고, 더 작은 규모의 downstream task 데이터셋에 fine-tune한다. 이를 위해 사전 학습된 prediction head를 제거하고 0으로 초기화된 $D\times K$feedforward layer를 추가한다.($K$는 downstream 클래스 개수이다.)

고해상도 이미지를 입력했을 때, 패치 크기 ($P,P$)를 유지하고는, 유효 시퀀스 길이 $N$을 크게 만든다. ViT는 메모리 제한 내에서 아무 길이의 시퀀스를 처리할 수 있지만, 사전 학습된 position embedding은 더 이상 의미가 없게 된다. 그러므로 pre-trained position embedding에 2D interpolation이 적용된다. 이런 해상도 조정과 패치 추출은 이미지의 2차원 구조에 대한 inductive bias가 ViT에 수작업으로 주입되는 유일한 과정이다.

4. Experiments

4.1. Setup

1) Datasets

모델의 scalability를 비교하기 위해, ImageNet(1k classes, 1.3M images), ImageNEt-21k (21k classes, 14M images), JFT(18k classes, 303M high resolution images)와 같은 여러 데이터셋을 사용한다.

task를 3개 그룹으로 나눈다. : Natural, Specialized(의학, 인공위성), Structured(localization을 이해하는 geometry)

2) Model Variants

“Base”와 “Large”모델은 BERT와 동일하지만, “Huge”는 더 큰 규모로 추가하였다. ViT-L/16은 16X16 입력 패치 크기를 갖는 Large Model을 의미한다.

단, Transformer의 시퀀스 길이는 패치 사이즈에 반비례하므로, 패치 사이즈가 작아질수록 계산 비용이 더 커진다.

baseline CNN으로 ResNet을 사용했는데, Batch Normalization layer를 Group Normalization으로 대체했고, standardized convolution을 사용했다. 이런 수정을 통해 전이 학습을 향상시켰고, 수정된 모델을 ResNet(BiT)라고 표기한다.

3) Training & Fine-tuning

ResNet을 포함한 모든 모델을 학습시킬 때, Adam optimizer를 이용했다. 하이퍼파라미터는 다음과 같이 설정한다. $\beta_1 = 0.9, \beta_2 = 0.999$, 배치 크기 = 4096, weight decay는 0.1로 설정하였다. linear learning rate warmup and decay를 사용한다.

✅ Linear Learning rate warmup and decay : 업데이트하기 위해 linear하게 learning rate를 증가시켰다가, linear하게 감소시키는 learning rate schedule.

fine-tuning 시 모멘텀으로 SGD를 사용하고, 배치 사이즈는 512이다.

4) Metrics

데이터셋에 대한 결과를 few-shot accuracy와 fine-tuning accuracy로 기록했다.

✔️ Fine-tuning accuracy는 fine tuning이후 각 모델의 성능을 반영한다.

✔️ Few-shot accuracy는 훈련 이미지 부분집합 표현을 $\{-1,1\}^K$ 타겟 벡터로 매핑하는 least-squares regression의 값을 구하는 과정에서 얻을 수 있다.

주로 fine-tuning accuracy에 집중하지만, fine-tuning의 비용이 너무 비쌀 경우, 가볍게 평가할 수 있는 few-shot accuracy를 사용한다.

4.2. Comparison to State of the Art

4.3. Pre-Training Data Requirements

ViT는 대규모 JFT-300M 데이터셋을 학습했을 때 성능이 좋았다. ResNet과 비교해 inductive bias가 더 적은데, 이 때 데이터셋의 크기가 얼마나 중요할까? 이를 확인하기 위해 2가지 실험을 시도했다.

첫째, ViT 모델을 데이터셋의 사이즈를 증가시키면서 사전학습시킨다. : (ImageNet, ImageNet-21k, JFT-300M 순.) 작은 데이터셋에서의 성능을 향상시키기 위해서 3가지 기본 regularization 패러미터를 최적화한다. : 1️⃣ weight decay 2️⃣ dropout 3️⃣ label smoothing

가장 작은 데이터셋인 ImageNet에 사전학습했을 때, ViT-Large 모델은 ViT-Base 모델보다도 성능이 좋지 않았다. (regularization을 했음에도 불구하고) medium size dataset인 ImageNet-21k를 사전학습했을 때는 성능이 비슷했다. 그런데 large size dataset인 JFT-300M을 사전학습했을 때에는 더 큰 모델을 사용했을 때 더 큰 효율을 보였다.

작은 데이터셋으로 사전 학습 시, BiT가 앞서지만, 데이터셋 규모가 커지면 커질수록 ViT가 능가한다.

둘째, 랜덤 부분집합 (9M, 30M, 90M)에 대해 학습시켜보았다. 작은 부분집합에 대해서는 추가적인 regularization을 하지 않았고, 모든 하이퍼패러미터 설정을 유지했다. 이런 방법으로 모델의 본질적인 성질을 평가했다. 계산 비용을 절약하기 위해서, full fine-tuning accuracy 대신 few-shot linear accuracy를 기록했다.

상대적으로 작은 데이터셋에 대해 사전 훈련했을 때에는, BiT가 성능이 더 좋았으며, ViT는 소규모 데이터셋에 대해 overfitting이 발생했다. 예를 들어 ViT-B/32(Base model with 32X32 patch size) 모델은 ResNet50보다는 살짝 빨랐지만 소규모 데이터셋 9M subset에서는 좋지 않았다. 그러나 90M 이상의 subset에 대해서는 더 나았다. 이러한 결과는 ⭐ convolutional inductive bias가 소규모 데이터에 대해서는 유용하지만, 대규모 데이터에 대해서는 데이터로부터 적절한 패턴을 학습하는 것만으로도 충분함을 의미한다.

4.4. Scaling Study

결과 Table6을 보면 여러 패턴을 확인할 수 있다.

첫째, ViT 모델은 performance/compute trade-off 관점에서 ResNet을 압도한다. 같은 성능을 달성하기 위해 필요한 계산 비용이 ViT가 2~4배가 더 적다.

두 번째 하이브리드는 적은 계산비용으로는 ViT를 앞지르지만, 대규모 모델에 대해서는 그 차이가 거의 없다.

셋째, ViT는 실험한 범위 내에서는 saturate(포화)되지 않았다.

4.5. Inspecting Vision Transformer

Vision Transformer가 이미지 데이터를 어떻게 처리하는지 이해하기 위해서, 내부적으로 어떻게 표현되는지를 (representation) 분석한다.

ViT의 첫 번째 층은 flattened 패치를 저차원 공간으로 선형 투영한다. 위의 Figure 7. 왼쪽 사진을 보면, 이 층에서 학습된 임베딩 필터의 상위 주성분을 보여준다. 이 성분들은 각 패치 내의 미세 구조의 저차원 표현을 위한 basis function을 닮았다.

projection 이후에, 학습된 position embedding은 patch representation에 추가된다. Figure 7. 가운데 사진은 모델이 position embedding의 유사도 내에서 이미지 내부의 ‘distance’를 인코딩하는 방법을 학습함을 보인다. 다시 말해, 가까이 있는 패치일수록, 더 유사한 position embedding을 가진다. 더욱이, row-column 구조를 나타내는데, row-column 구조란, 같은 행,열에있는 패치가 서로 유사한 임베딩을 가지는 구조이다.

Self-attention은 ViT로 하여금 최하위층에서도 전체 이미지 정보를 통합할 수 있게 해준다. 이 논문의 저자들은 이 네트워크가 얼마나 광활한 수용력을 잘 이용하는지에 대한 정도를 조사했다. 구체적으로, Figure 7.의 오른쪽 사진과 같이 attention weight를 기반으로, 정보가 통합되는 이미지 공간에서의 평균 거리를 측정한다. 이 “attention distance”는 CNN의 receptive field size와 유사하다.

CNN의 receptive field란?
각 단계의 입력 이미지에 대해 하나의 필터가 커버할 수 있는 이미지 영역의 일부이다. 다시 말해, 필터가 한 번에 볼 수 있는 이미지의 영역의 개념이다. receptive field는 층(convolutional layer)이 깊어질수록 선형적으로 증가한다.

사진 보면, 층이 깊어질수록 attention distance도 증가하는 성질을 보이는데, 이런 것도 receptive field와 유사하다. 심지어는 최하위층 레이어에서도 몇몇의 attention head가 이미지의 대부분에 attend됨을 발견할 수 있었는데, 이는 정보를 global하게 integrate할 수 있는 능력이 실제로는 모델이 사용하고 있음을 보여준다. 다른 attention head는 하위층에서 일관적으로 작은 attention distance를 가진다. 이렇게 고도로 국소화된 attention은 Transformer 이전에 ResNet을 적용한 하이브리드 모델보다는 덜하지만, 이는 이 local한 attention이 CNN의 초기 Convolutional layer들과 유사한 기능을 함을 보인다.

더불어, 저자들은 ViT가 분류에 semantically relevant한 이미지 영역에 attend함을 발견하였다. Figure 6 사진을 보면, input에서 의미적으로 관련 있는 영역에 대해 attend하는 것을 볼 수 있다. (ex: 강아지, 비행기)

4.6. Self-Supervision

Transformer는 NLP 문제에 대해서 좋은 성능을 보였다. 그러나, 이런 성공은 뛰어난 확장성 뿐만 아니라, 대규모의 self-supervised pre-training으로부터 비롯되었다. 저자들은 BERT에서 사용된 masked language modeling을 모방하여, self-supervision을 위해 masked patch prediction을 실험했다. 그 결과, ViT-B/16 모델은 ImageNet에서 79.9%의 정확성을 보이며, scratch로 학습시키는 것보다 2%라는 더 유의미한 향상을 보였다. 그러나, 여전히 supervised pre-training에 비해서는 4%정도 뒤쳐졌다.

5. Conclusion

컴퓨터비전에서 self-attention을 쓰던 다른 연구와는 달리, ViT는 그 어떤 이미지 특화된 inductive bias를 추가하지 않았다. 대신, 이미지를 패치 시퀀스로 보았고, NLP에서 쓴 표준 Transformer encoder 부분을 이용해 처리하였다. Transformer encoder를 최대한 수정하지 않고 가져오게 되어 비교적 간단하면서도 확장성이 좋기 때문에, 이런 전략은 큰 데이터셋에 대해 사전 훈련할 경우 훨씬 좋은 성능을 보였다. ViT는 사전 훈련 비용이 상대적으로 저렴하면서도 SOTA와 견주었을 때 비슷하거나 더 향상된 성능을 보였다.

그렇지만, 여전히 과제들이 남아 있다.

1) detection이나 segmentation과 같은 다른 컴퓨터비전 과제에도 ViT를 적용할 수 있어야 한다.

2) self-supervised pre-training 방법을 계속해서 연구해야 한다.

: 이 연구에서 self-supervised pre-training로 좋은 결과를 보였지만, 여전히 large-scale supervised pre-training과의 격차가 컸다.

3) 성능 향상을 위해 ViT의 추가적인 확장이 필요하다.