ViT Vision Transformer(2021) 논문 리뷰

ViT (Vision Transformer) 모델의 논문 이름은 An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale다. (링크)

 

저자는 Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby다. 

 

NLP에서 쓰던 Transformer를 컴퓨터 비전 분야에 성공적으로 적용한 모델이 나온 논문이다.

 

Abstract

트랜스포머 구조는 NLP 분야에서 실질적으로 (de facto) 표준적인 모델일 때 Comupter vision 분야에서는 제한적으로 사용되었다. 기존에는 attention와 convolution을 함께 사용했다. 하지만 본 논문에서는 attention만 가지고도 신뢰할 수 있는 이미지 분류 결과를 얻었음을 보인다. 이는 attention가 이미지의 패치 (patch)의 시퀀스를 학습하게 함으로써 달성한다. 우리는 이 모델을 ViT (Vision Transformer)라고 부른다.

 

 

3. Method

 

 

ViT의 기본 구조는 Figure 1과 같다.

 

Input Size and Tensor Shape

본래의 표준적인 트랜스포머는 토큰 임베딩으로 구성된 1D 시퀀스를 입력으로 받는다.

 

하지만 이미지는 2D이므로 image $\textbf{x} \in {\mathbb{R}}^{H \times W \times C}$를 다음과 같은 과정으로 정제한다.

 

이미지 패치를 flatten하여 다음처럼 크기를 변환한다.

 

${\textbf{x}}_p \in {\mathbb{R}}^{P^2 \cdot C}$

 

이때 $(P, P)$는 이미지 패치의 해상도이며 $(H, W)$는 원본 이미지의 해상도다. 그리고 $C$는 채널의 수다.

 

따라서 패치의 수는 $N = HW / P^2$로 계산된다.  

 

이때 이미지의 로 데이터 (raw data)를 사용하는게 아니라 하나의 학습가능한 linear projection 레이어를 통해서 패치 임베딩을 생성하고 이를 Transformer encoder에 넣는다. 

 

ViT에서는 모든 latent vector와 layers에 대해서 constant vector size인 $D$를 사용한다.

 

Special Token

BERT와 유사하게 sequence의 맨 앞에 [class] 토큰을 덧붙인다. 

 

이 class head에 대한 MLP는 pre-training 단계에서는 숨기고, fine-tuning 단계에서 하나의 레이어로 [class] 토큰에 대한 라벨을 예측한다.

 

 논문에는 말이 좀 어렵게 쓰여져있는데 ViT 설명 한국어 (링크)에 보면 그냥 Transformer encoder를 통과한 다음 [class] 토큰에 MLP를 적용해서 클래스를 예측한다는 이야기다.  Figure 2만 보는게 차라리 더 이해하기 쉬운거 같다. 

 

 

 

Positional embeddings는 1D 포지션 임베딩을 적용했다.

2D-aware 포지션 임베딩도 수행했으나 성능적인 향상을 볼 수 없었으며 이는 Appendix D에 첨부했다.

 

Patch Embedding과 Position Embedding을 더해서 Transformer Encoder에 보낸다.

 

Final Input

 

Flatten, projection, adding special token, positional embeddings를 모두 사용한 최종 input ${\textbf{z}}_0$은 다음과 같다.

 

$\textbf{E}$는 embedding matrix를 나타낸다. $\textbf{E} \in {\mathbb{R}}^{(P^2 \cdot C) \times D}$ 그리고 $\textbf{E}_{pos} \in {\mathbb{R}}^{(N + 1) \times D}$

 

 

$ {\textbf{z}}_{0}$ = [ $ \textbf{x}_{class}, \textbf{x}_{p}^{1} \textbf{E} $ ; $ {\textbf{x}}_{p}^{2} \textbf{E} $ ; $ \cdots $ ; $ \textbf{x}_{p}^{N} \textbf{E} $ ] + $ \textbf{E}_{pos} $

 

 

Transformer Encoder Block

Encoder block은 Multi-Head Self-Attention (MSA)과 FFN (MLP)로 이루어진다.

FFN의 경우 MLP와 GELU로 구성된다.

그리고 Layernorm (LN)은 모든 블록에 적용되며 residual connections는 모든 블록을 거친 다음에 적용된다.

 

 

 

 

Inductive Bais

CNN은 2차원에서 인접한 locality를 학습하지만 ViT의 경우 self-attention이 global feature를 학습하고 MLP가 local을 학습하기 때문에 다르다.

 

Fine-tuning

ViT는 대규모 데이터셋에 pre-train한 다음 상대적으로 작은 downstream tasks에 파인 튜닝을 수행한다. Pre-train 후 Higher resolution에 파인튜닝을 하는데 이때도 패치의 크기는 동일하다. 하지만 pre-trian에서 수행했던 position embeddings의 의미가 없어진다. 이를 해결하기 위해서 pre-trained 포지션 임베딩에 2D interpolationd을 수행한다.

 

4. Experiments

Datasets

ILSRVRC-2012 ImageNet에는 총 1k class가 있으며 1.3M 이미지가 있다.

JFT 데이터는 총 18k의 클래스가 있고 303M high-resolution 이미지가 있다. 

 

Downstream 데이터는 다음과 같다. ReaL labels, CIFAR-10 / 100, Oxford-IIIT Pets, Oxford Flower-102, VTAB classification tasks다.

 

Model Variants

 

ViT는 Base, Large, Huge의 세 가지 사이즈가 있다.

자세한 내용은 아래 Table 1에 명시되어 있다.

 

 

Training and Fine-Tuning

Adam을 사용했으며 $\beta_1$은 0.9, $\beta_2$는 0.999다. 배치 사이즈는 4096이며 weight decay는 0.1로 설정했다.

 

 

 

 

 

ViT-Huge를 JFT에 대해 pre-train 했을 때 좋은 성능임을 알 수 있다.

 

 

VTAB에 대한 ViT의 크기별 성능이다. 클수록 성능이 좋다.

 

 

Figure 3, 4, 5는 ViT와 ResNet의 성능을 비교한다.

 

 

Appendix D

 

D. 4 Positional Embedding

 

No positional Embeddings: 패치들의 묶음

 

1D positional Embeddings: 하나의 시퀀스 

 

2D positional Embeddings: D/2 사이즈의 X축 임베딩과 Y축 임베딩을 concat하여 사용

 

Relative positional Embeddings: 모든 패치의 쌍 (pairs)에 대해서, 구체적으로는 (q, k/v) 페어들이다. (q_1, k_2), (q_1, k_3), ..., (q_1, v_2), ... (q_1, v_3), ... ), 에 대해서 offset $p_q - p_k$을 구한다. 저자들은 단순히 추가적인 attention을 수행했다고 한다.

 

포지션 임베딩들의 비교 실험 결과는 아래의 Table 8에 있다.

 

 

1D를 기준으로 했을 때 포지션 임베딩이 없는 경우를 제외하고는 큰 차이가 없음을 알 수 있다.

 

저자들은 원본 이미지 224 x 224 같이 위치의 차이가 큰 경우 포지션 임베딩이 중요하지만,

패치의 경우 16 x 16과 같이 위치의 차이가 작기 때문에 포지션 임베딩의 방법이 상대적으로 덜 중요하다고 추측하고 있다.

 

 

 

 

References:

https://daebaq27.tistory.com/108

https://discuss.pytorch.kr/t/vision-transformer-a-visual-guide-to-vision-transformers/4158

https://blog.mdturp.ch/posts/2024-04-05-visual_guide_to_vision_transformer.html?