[논문 리뷰] ViT, Vision Transformer (An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale)

[딥러닝 논문 리뷰 시리즈]

노션에서 작성한 글을 옮겼으며, 아래 노션에서 더 깔끔하게 읽으실 수 있습니다.

>>노션 링크<<

ViT (Vision Transformer) | Notion

 

Abstract & Introduction

• 2017년 Transformer의 발표 이후, NLP 분야에서는 지금까지 Transformer 구조가 표준으로 사용되고 있다.
  • 대부분은, Large Dataset으로 pre-train한 이후, task에 맞는 dataset으로 fine-tuning하는 방식으로 이용하고 있다.
  • Transformer 이후 계산적 효율 및 확장성으로 인해, 모델의 크기(# of param)가 굉장히 커짐.
• 현재(2021)까지 Computer Vision 분야는 CNN 구조에서 벗어나지 못하고 있었다.
  • CNN에 Attention을 융합하여 사용하려는 시도가 있었으며,
  • transformer와 같이 CNN 구조를 attention을 이용하여 완전히 대체하려는 시도 또한 있었다.
  • 그러나 성능 및 효율성이 아직 수준에 미치지 못해, ResNet과 같은 구조가 여전히 SOTA 모델로써 사용되고 있다.

 

• 이 논문에서는 최소한의 수정만을 거쳐, Image를 그대로 Transformer 구조에 적용했을 때의 실험 결과를 보여준다.
  • 이미지를 여러 개의 patch로 분할한 후,
  • 이 patch의 linear embedding 결과를 그대로 transformer의 input으로 사용한다.
  • 다른 과정은 기존 transformer와 거의 동일하다.
• 이 구조는 중간 크기의 dataset으로 학습한다면, ResNet보다 성능이 좋지 않다.
  • 강한 regularization 없이는, ImageNet과 같은 mid-size dataset으로 학습했을 때 성능이 좋지 않다.
  • 이는 translation equivariance, locality와 같은 CNN에서의 Image의 Inductive bias가 존재하지 않기 때문.
• 그러나 큰 dataset으로 학습했을 때는 성능이 좋아진다.
  • Large Dataset으로의 training을 통해 Inductive bias가 없더라도 충분한 성능을 낼 수 있음을 확인했다.
  • ViT를 ImageNet-21k, JFT-300M과 같은 Large Dataset으로 pre-train한 이후, fine-tuning하였을 때, 다른 SOTA 모델들보다 Image Recognition 성능이 우수했다.

 

(사견)
결국 기본적인 RNN 구조 - Seq2Seq, LSTM.. 과 같은 구조들이 지배적인 모델일 시기에는, (물론 데이터의 양도 중요했지만) 모델의 구조가 어떻느냐에 따라 성능이 좌우되는 경향성이 어느 정도 있었다.
그러나 Transformer의 등장 이후, 학습에서의 병렬처리가 가능해짐에 따라 Scalability, Efficiency가 폭발적으로 확장되었고.. 대규모 dataset이 생겨남에 따라, 작금에 들어서는 얼마나 많은 데이터로 pre-train했느냐가 성능을 좌우하는 큰 요소가 되었겠구나 하는 생각이 든다.

 

Method

• ViT(Vision Transformer)의 구조는 위 그림(그림1)과 같다.
  • 모델 구조는 의도적으로 Transformer와 거의 유사하도록, 간단하게 구성하였다.
  • 이를 통해 NLP에서의 Transformer의 scalability와 efficiency를 가질 수 있었다.

 

Vision Transformer (ViT)

• 2D Image를 Transformer의 Input (1D seq of token embedding)에 맞추기 위해, 여러 개의 patch로 나누어 각각 flatten시킨다.
  • $\mathbb{R}^{H\times W\times C}$ 의 Image를 $\mathbb{R}^{N\times (P^2\cdot C)}$로 Flatten한다.
  • 이 때 $H\times W$는 Image의 해상도, $C$: Channel(RGB), $N$: patch의 개수 ($HW/P^2$), $P$: patch의 해상도 (16x16의 patch라면, P=16)

 

• Tranformer에서 Latent Space vector size를 일정하게 유지시켰던 것처럼 동작하기 위해, Flatten한 patch들을 D차원으로 Linear Projection(Mapping)한다.
  • 위 식(Eq 1.)에서의 $\mathbf{E}$가 Linear Projection을 수행하는 matrix.
  • $\mathbf{E} \in \mathbb{R}^{(P^2\cdot C)\times D}$ 이므로, 각 Patch $\mathbf{x}_p^n \in \mathbb{R}^{P^2\cdot C}$ 와 행렬곱되어 $D$차원의 vector로 변환된다.
  • 여기서의 Linear Projection Matrix $\mathbf{E}$는 trainable한 parameter이다.
  • 이를 마친 것을 patch embedding이라 한다.

Transformer에서의 Latent Space vector size는, Input Embedding 이후 Encoder-Decoder를 거치면서 일정하게 유지된 $d_{model}$을 의미한다. ($d_{model} = 512$)

 

• BERT에서의 Class token과 비슷하게, patch embedding sequence 맨 앞에 class token을 추가해 준다. ($\mathbf{x}_{\text{class}}$)
  • 다른 patch embedding과 동일한 D차원으로 초기화된다.
  • class token을 이용하여 결과를 내어 학습할 때, pre-train 시점에는 1개의 hidden Layer를 가진 MLP를, fine-tuning 시점에는 Single Linear Layer를 거치도록 하여 학습한다.
    (class token을 LayerNorm → MLP/Single Linear → Softmax → Classification.)

Class token의 역할은 다른 모든 patch embedding input들로부터 모든 정보를 집약하는 것이다. (Seq2Seq에서의 Context vector와 비슷한 맥락이라고 생각하면 될 것 같다.)
여러 번의 Self Attention을 통해 업데이트되며, Class token이 Image 전체의 정보를 통합하도록 학습한다.
이를 통해, L번의 Self-Attention Layer (Encoder)를 거친 이후, 최종적으로 classification을 수행할 때, update된 class token($\mathbf{z}_L^0$) 만을 사용하여 결과를 낸다. (Eq 4.)

 

• Linear Projection을 마친 이후, Transformer와 마찬가지로 Positional Embedding을 더해준다. ($\mathbb{E}_{pos}$)
  • Image에 대한 Positional Embedding이기 때문에, 2D positional embedding을 사용해야 하는 것 아닌가- 라고 생각했지만,
  • 실험 결과 1D positional embedding에 비해 눈에 띌 만한 개선이 보이지 않아, 1D positional embedding을 사용한다.
  • positional embedding까지 더한 결과를 Encoder의 Input으로 사용하게 된다.

 

• Transformer Encoder는 Multi-head Self Attention(MSA) block과 MultiLayer Perceptron(MLP) block으로 구성한다. (그림 참고, Eq 2, Eq 3 참고.)
  • 각 block (MSA, MLP) 이전에는 모두 Layer Normalization을 적용한다.
  • 각 block 이후에는 모두 Residual Connection을 더해준다.
  • MLP는 activation function으로 GELU를 이용한다.

 

Inductive bias

• ViT는 CNN보다 Inductive bias가 훨씬 적다.
  • CNN에서는 2D locality나 equivariance가 모델 전역에 걸쳐 보존된다.
  • 이에 반해, ViT에서는 이러한 locality, equivariance가 MLP에서만 존재한다.
• 그렇기 때문에, positional embedding을 통해 이를 보완하지만, 2D locality는 굉장히 적게 사용된다.
  • fine-tuning 시, 다양한 해상도의 image에 대해 학습하여 positional embedding을 조정한다.
  • positional embedding 초기화 시, 2D의 공간적 정보는 들어 있지 않으며, 학습 과정에서 2D의 spatial한 정보가 학습된다.

여기에서 알 수 있듯, Transformer와는 달리 ViT에서의 positional embedding은 trainable하다!
고정된(sine, cosine을 이용한) 원래의 positional embedding을 사용하는 것과 trainable한 positional embedding을 사용하는 것이, 실험적으로 거의 유사하다는 것이 입증되었다.
(https://arxiv.org/abs/2102.11090)

 

Hybrid Architecture

• Raw Image에 대한 patch 대신, CNN을 함께 적용하여 CNN에서의 feature map으로 patch를 만들어 사용하는 방식도 사용될 수 있다.
• 이러한 feature map을 사용하는 경우, patch의 크기를 1x1로 설정하여 학습을 진행하는 경우도 있다.

 

Fine-tuning and Higher Resolution

• 일반적으로 ViT는 Large Dataset으로 pretrain한 이후, fine-tuning하는 식으로 학습한다.
  • fine-tuning시, pre-train에 사용했던 prediction head를 제거하고 다시 새로운 Feed-Forward Layer로 대체한다 (zero-initializing)
• Pre-train에서 사용했던 이미지보다 고해상도의 이미지로 fine-tuning하는 것이 효과적이다.
  • 이 때, 같은 patch size를 사용하면 고해상도의 이미지에서 patch sequence의 길이가 길어진다.
  • 이 과정에서 학습한 positional embedding이 의미가 없어질 수 있다.
  • 따라서, 원본 이미지의 위치에 따라 2D interpolation을 수행하여 처리한다.
  • (여기에서 Inductive bias가 추가된다.)

 

Experiments

• ResNet, ViT, Hybrid 모델을 다양한 크기의 Dataset으로 Pretrain한 후, 여러 benchmark를 통해 평가한다.
• 평가 결과에 따르면, 대부분의 경우 ViT가 더 낮은 비용으로 더 높은 성능을 낸다.

 

Setup

• Model Scalability를 확인하기 위해, 여러 크기의 Dataset을 통해 pretrain하여 실험하였다.

 

• BERT에서 쓰인 것과 같이, 위와 같이 변형된 Transformer 구조를 사용한다.
  • 예를 들어, ViT-L/16이라면, ViT-Large 구조를 사용하며, 16x16 patch size를 사용했다는 것.
• 비교군인 CNN은 ResNet을 사용한다.
  • Batch Normalization Layer를 Group Normalization으로 대체해 사용하며, standardized convolutions를 사용하였다.
• Hybrid는 feature map을 가져와 1x1 pixel 크기의 patch를 사용하였다.
  • 다양한 seq length를 위해, ResNet의 여러 단계에서 feature map을 추출하여 사용했다.
• 학습 과정에서는 Adam optimizer를 사용했다.
  • $\beta_1 = 0.9, \beta_2 = 0.999$, batch size=4096, high weight decay 0.1.
• linear learning rate warmup / decay 사용.
• Fine-tuning 시에는 SGD with momentum (batch size 512) 사용.
• Dataset들에 대해 fine-tuning accuracy 또는 few-shot accuracy Metric을 이용한다.
  • fine-tuning accuracy는 각 dataset에 대해 fine-tuning 이후 accuracy를 측정한다.
  • fine-tuning에 cost가 과다하게 드는 경우, 대안책으로 few-shot accuracy를 이용한다. 이 방식에서는 pretrained model에서 얻은 frozen representation을 바탕으로 Least Squares regression problem의 해를 구하여 측정한다.

 

Comparison to State of the Art

• 우선 이 표에서는 ViT의 가장 큰 모델 (ViT-H/14, ViT-L/16)을 이전의 SOTA 모델들 (ResNet을 이용한 BiT-L, Noisy Student)와 비교한다.
  • BiT-L (ResNet)은 Large ResNet을 이용하여 Supervised Transfer Learning을 수행한 모델.
  • Noisy Student는 EfficientNet 구조로, label을 제거한 ImageNet/JFT-300M Dataset을 이용하여 Semi-Supervised Learning을 수행한 모델.
  • 현재(논문 시점) ImageNet에 대해서는 Nosiy Student가, 나머지 Dataset에서는 BiT-L이 SOTA.
• ViT-L/16 (JFT-300M) 모델이 BiT-L보다 학습 cost가 더 적음에도 모든 task에서 성능이 뛰어남을 확인할 수 있다.
• ViT-H/14 (JFT-300M) 모델은 이보다 더 나은 성능을 보인다. 그럼에도 기존의 SOTA 모델들보다 학습 cost가 더 적다.
• ViT-L/16 (ImageNet-21k) 모델은 훨씬 적은 학습 cost가 들어감에도, 괜찮은 성능을 보인다.

 

• VTAB task에서도 이전 SOTA 모델들보다 더 나은 결과를 확인할 수 있다.

VTAB: pre-trained model이 여러 downstream task에서 trasnfer learning / few-shot learning 등을 통해 얼마나 좋은 성능을 내는가를 측정하는 지표.

 

Pre-training Data Requirements

• ViT는 이전의 CNN 모델들보다 Inductive bias가 적기 때문에, Dataset의 크기에 영향을 크게 받는다.
• 아래에서는 두 종류의 실험을 통해 ViT가 Dataset의 크기에 따라 얼마나 영향을 받는 지를 보여 준다.

 

(Figure 3, 4)

• 왼쪽 표 (Figure 3)에서는 모델들을 ImageNet / ImageNet-21k / JFT-300M의 세 Dataset으로 pre-train한 결과를 보여준다.
  • ImageNet은 크기가 작기 때문에, weight decay, dropout, label smoothing과 같은 regularization 적용.
• 작은 크기의 dataset(ImageNet)으로 pre-train했을 때, 표에서 볼 수 있듯, Base Model(ViT-B/16)의 성능이 오히려 Large Model(ViT-L)보다 나은 것을 볼 수 있다. (regularization을 적용했음에도!)
• 큰 크기의 모델이 충분한 이점을 내려면 JFT-300M 정도 크기의 Dataset은 되어야 한다는 것을 확인할 수 있다.

 

• 오른쪽 표 (Figure 4)에서는 JFT-300M의 subset을 이용하여 10M, 30M, 90M, 300M으로 pre-train한 결과를 보여준다.
  • 여기에서는 추가적인 regularization 없이 모든 조건 동일하게 적용.
  • early stopping 적용, 최고 성능 epoch의 val_acc 사용.
  • fine-tuning accuracy 대신 few-shot linear accuracy 사용.
• 결과를 보면, ViT는 작은 dataset에서 CNN보다 훨씬 overfitting이 심하게 일어나는 것을 확인할 수 있다.
• 그러나 90M 이상으로 가면, ViT의 성능이 올라오는 것을 볼 수 있다.
  • 이것으로 미루어 보아, CNN에서의 Inductive bias가 작은 dataset에 대해서는 효과적이지만, 큰 dataset에서는 많은 양의 data로부터 relative한 pattern을 충분히 학습할 수 있다는 것이 확인된다.

 

(이하 실험결과 생략, 논문 참고.)

 

Conclusion

• Image Recognition에 Transformer를 직접 적용해 보았다.
• 이전과 달리 Inductive bias를 활용하지 않고, 이미지를 patch로 잘라 Transformer의 Encoder Input으로 활용하였다.
• 이러한 간단하고 scalability가 높은 방식은 Large Dataset에 대해 Pre-train하였을 때 좋은 성능을 보이며, 또한 적은 비용으로 학습이 가능하였다.

 

• Detection, Segmentation 등과 같은 Vision Task에 ViT를 적용해보는 확장 가능성이 남아 있다.
• 또한 Self-Supervised Learning에 대해 간단하게 실험해 보았지만, 아직 Supervised Learning에 미치지 못하는 성능을 보였으므로, 발전시킬 여지가 남아 있다.

 

리뷰 후기

아직까지 많은 곳에서 기본적인 Vision 처리 모델로 활용되는 ViT 모델을 이제야 정리해 본다.

Abstract와 Conclusion만 보아도 알 수 있듯, Image를 그대로 Transformer 구조에 넣어버림으로써 Scalability와 Efficiency를 챙기고, 잃어버린 2D Image의 locality와 translation equivariance를 Large Dataset으로 Pre-train함으로써 보완하는 것이 핵심 내용.

RNN의 틀을 깨고 NLP에서의 scalability를 챙긴 Transformer에 이어,
Vision에서의 scalability를 열어준 ViT 이후로 작금에 이르러선 모든 분야에서 더 크고 더 많은 Dataset이 너무나도 중요해져 버린 것 같다.