An Image Is Worth 16x16 Words (ViT)

 

논문 링크

https://arxiv.org/abs/2010.11929

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

 

ViT의 기본 구조가 되는 transformer에 대한 설명

https://kk-eezz.tistory.com/97

Attention Is All You Need (Transformer) 논문 요약

 

 

1. Introduction

• 이 논문에서는 CNN이 없이 transformer 구조만으로도 image classification task에서 좋은 성능을 낼 수 있음을 보여준다.
• 이를 위해, 이미지를 작은 patches로 분할하고 이러한 patch들에 대해 linear embedding을 적용하여 transformer의 입력값으로 사용하도록 한다.

• 이렇게 학습된 transformer를 ImageNet과 같은 중간 크기를 갖는 데이터셋에 대해 학습시키면 그렇게 좋지 못한 결과를 얻게 되는데, 이는 transformer가 CNN과는 달리 내제되어 있는 inductive bias가 부족하기 때문이다.
  • inductive bias에 대한 설명은 밑에서 좀 더 자세히 나와있다.

 

• 그러나, 큰 크기의 데이터셋에서는 inductive bias가 문제가 되지 않았다. ViT는 충분한 크기의 데이터셋에 대해 pretrain 한 후에, 전이 학습을 수행한 경우, 굉장히 좋은 퍼포먼스를 보였다.

 

 

 

2. Method

• Transformer는 1차원의 token embedding을 입력으로 받는다. 저자들은 2차원의 이미지들을 처리하기 위해, 이미지를 patch로 쪼개서 처리한다.

• 이 부분을 그림으로 나타내면 다음과 같이 나타낼 수 있을 것 같다.

• 또한 patch embedding이란 과정을 통해 이러한 patch들을 trainable linear projection을 통해 D-차원으로 매핑한다.
• 이렇게 임베딩된 patch의 맨 앞부분에 추가적인 embedding인 [class] token을 덧붙이는데, 이 class token은 나중에 인코더의 마지막 레이어를 통과한 후에, image representation y가 된다. 이 y는 다시 classification head의 input으로 들어가게 되고, 최종적인 결과값이 된다.
• 또한, patch embedding에 더하여 위치 정보를 삽입하기 위하여 position embedding이 사용되는데, 이 논문에서는 learnable 1D position embedding을 사용하였다고 한다.
• 이러한 과정을 수식화 하면 아래와 같이 표현할 수 있다. 여기서 MSA는 (Multi-head self attention이며, LN은 layer normalization, X_class는 class token을 뜻한다)
  • 참고로, 원본 transformer 논문에서는 sine, cosine 함수를 사용하여 위치 정보를 인코딩하였다.

 

Inductive bias

• inductive bias란 학습 시에는 만나보지 않았던 상황에 대하여 정확한 예측을 하기 위해 사용하는 추가적인 가정을 의미한다. CNN을 예로 들면, 커널을 사용하면 지역적인 특성을 잘 학습할 것이라는 가정이 이에 해당된다고 말할 수 있을 것 같다.
• locality 뿐만 아니라, CNN은 two-dimensional neighborhood structure, translation invariance 등의 특성이 각 레이어마다 존재하여 전체 모델에 대해 이러한 특징을 가지 된다.
• 반면에, ViT는 MLP 레이어가 유일하게 locality와 translation equivalent 특성을 가지므로 CNN과 비교했을 때, 굉장히 적은 inductive bias를 갖는다고 말할 수 있다.

 

Hybrid Architecture

• 입력 이미지를 input으로 받는 기본 ViT 모델 이외에도 저자들은, raw image 대신에 CNN의 output feature를 입력으로 받는 hybrid architecture 또한 제시한다.

 

Fine-tuning and higher resolution

• Fine-tuning시에는 pre-train된 prediction head를 D*K 크기의 feedforward 레이어를 대체하였다고 한다.
• 또한, fine-tuning시에 더 높은 화질을 갖는 이미지에 대해 학습시킬때, patch 사이즈를 그대로 유지하면 입력 시퀀스의 길이가 늘어나게 된다. transformer의 경우 입력 받을 수 있는 길이가 정해져있는 것이 아니기 때문에 상관이 없지만, positional encoding의 경우 미리 학습되기 때문에 이 부분이 문제가 될 수 있다. 따라서 이를 해결하기 위해 입력 이미지에서의 위치를 통해 2D interpolation을 수행하였다고 한다.

 

 

 

3. Experiments

• 각각, Base, Large, Huge 크기를 갖는 세가지 버전이 있으며, ViT-L/16으로 표시된 모델은 16*16 크기의 patch를 사용한 ViT Large를 의미한다.

 

(그당시) SOTA 모델과의 비교

• ViT-L/16 모델은 BiT-L 모델에 비해 학습 시에 비교적 적은 양의 컴퓨팅 파워를 요구하면서도 더 높은 성능을 보였다. 여기서 TPUv3-core-days는 학습에 사용된 TPU core의 개수와 소요된 날짜를 곱한 값이라고 한다.
• 또한, 그보다 큰 모델인 ViT-H/14 모델 또한 더 적은 수의 컴퓨팅 파워를 요구하면서도 더 좋은 성능을 보인다.

• 아래의 그림은 각각의 세부적인 VTAB task에 대한 성능을 나타낸다.

 

Pre-training Data Requirements

• 아래는 각각 증가하는 크기의 데이터셋에 대해 모델의 성능을 비교한 그래프이다.

• ViT는 데이터셋이 증가할 수록 좋은 성능을 보였는데, 이는 inductive bias가 작은 크기의 데이터셋에 대해서는 효과를 보일지 몰라도, 큰 크기의 데이터셋에 대해서는 데이터로부터 학습하는 것만으로도 충분하다는 것을 보여준다.

 

Scaling Study

• 아래는 Performance-compute trade-off 측면에서 ViT/ViT-hybrid, ResNets을 비교한 그래프이다.

• ViT가 ResNets을 능가하는 것을 확인할 수 있으며, small computation의 경우에는 hybrid 모델이 ViT를 능가했지만, 그렇지 않은 경우에는 그 차이가 없어지는 것을 확인할 수 있다.