Vision Transformer (AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE) 리뷰하기

Original paper : https://arxiv.org/abs/2010.11929

 

 

NLP task에서 Transformer가 등장하고 나서 Vision에서도 transformer가 적용된 논문이다. 요즘 sota에 대해서 왠만한 분야 다 transformer를 사용하는 것 같아서 이번 기회에 ViT 모델에 대해서 알아보고자 리뷰를 하게 되었다.

 

1. introduction

Transformer가 Self-Attention 기반 구조를 토대로 이루어져있는데, 이 모델의 장점이라 하면 계산이 효율적이고 확장에 용이하다는 점이다. 그러나 Vision에서는 CNN 기반 구조들이 중점인데, layer가 클 수록 모델 파라메터가 엄청나게 커지는 것을 학습하면서 다들 확인했을 것이다. 그래서 Image task에 transformer를 적용시킨 것이 Vision Transformer인데, 여기서 input을 하나의 이미지를 여러개의 패치로 나누어서 linear embedding을 통하여 들어가게 된다. 이 패치들이 결국 하나의 토큰이 되는 것이라고 한다.

 

 

---

 

모델 설명

 

 

위 모델이 ViT의 전반적인 구조인데, 여기서 중요한 점은 첫 번째 토큰의 입력으로 class에 대한 토큰을 넣는다는 점이다. 이는 Bert에서 착안한 컨셉이라고 한다. 그 외에는 우리가 아는 기본적인 Transformer 구조를 띄고 있다.

 

Input으로 2D image를 이제 어떻게 들어가게 하는 가가 중요한데, 기존의 NLP에서의 Transformer는 1D sequence of token embedding으로 들어가지만, 여기서는 이제 원본 이미지를 패치로 나누어서 이를 flatten 시켜서 들어가게 되는 것이다 ($x_{p} \in R^{N \times (P^{2} \dot C)}$).  여기서 N은 패치의 갯수를 나타내는데 H*W / P^2 의 값을 가진다.

예를 들어서, 3 by 3의 이미지를 이용해서 패치로 나눈다면, p = 3으로 할때 N은 9 개가 된다. 이렇게 하면 위의 모델 구조 입력에 나와있는 것처럼 구성이 될 것이다. 그리고 여기서 중요한 점은 input 이미지가 정사각형 형태로 구성이 필요하다는 것이다. 아마 이래야 패치 구성이 쉬워서이지 않을까 싶다. 마지막으로 Transformer는 모든 레이어에서 고정된 벡터 크기 D를 사용하기 때문에 이미지 패치는 펼친 다음 D차원 벡터로 linear projection 시키게 된다.

 

위치정보는 패치 임베딩한 것에 추가하는데, 2D positional embedding으로는 효과를 보지 못해서 1D position embedding을 사용한다. 

Transformer encoder는 전체적으로 두개의 Norm, multi-head attention, MLP로 구성되어 있는데, Norm 에는 LayerNorm, MLP에는 GELU와 two layer를 구성하고 있다. 그리고 residual connection을 이용하고 있다.

 

 

그리고 Transformer가 CNN, RNN보다 Inductive bias가 약하다라고 나와있는데, Inductive bias란 학습 시 보지 못한 입력에 대한 출력을 예측할 때 사용된 가설이라는 뜻을 가지고 있는데, 이는 즉 CNN은 spatial한 정보를 볼 수 있기에 가정을 할 수 있으며, RNN 역시 인접 및 이전 혹은 다음 시퀀스에 대한 정보와 연관이 있기 때문에 여러 가정을 할 수 있으나, Transformer의 경우 이미 positional embedding과 같은 여러 정보가 주어지기때문에, 설득력 있는 결과를 얻기 위해서는 충분한 데이터가 필요하다. 그래서 Transformer에서 inductive bias가 존재하는 부분이라 함은 이미지를 패치 임베딩 시킬때의 부분 정도로만 볼 수가 있다고 한다.

 

 

- Fine tuning and higher resolution

ViT는 보통 큰 데이터 셋에 대해 pre-train하고 나서, downstream task에 대해 파인 튜닝을 하게 되는데, 이를 위해서 pre-trained 된 prediction head를 제거하고 D x K의 feedforward layer를 추가한다고 한다 (K : number of downstream classes). 그리고 고해상도에서 학습할 때가 더 효율적일때가 있다. 고해상도 이미지를 다룰때는 패치 사이즈는 동일하게 하나 패치 시퀀스 길이를 더 크게 한다고 한다.

사전 훈련된 positional embedding은 크게 의미가 없기에, 원본 이미지의 위치에 따라 2D - interpolation을 수행한다.

 

 

Model Variants

 

 

모델 사이즈에 따라서 Base, Large, Huge로 표현을 하는데, 모델 파라메터 사이즈를 보니까 각 크기마다 사이즈 차이가 많이 나긴 하지만, 다른 모델들보다 비교적 작게 나오는 것 같다. 그리고

Transformer의 시퀀스 길이는 패치 크기의 제곱에 반비례하므로 패치 크기가 작은 모델은 계산적으로 더 비쌉니다 라고 하는데, 패치 크기가 클 수록, sequence가 짧기 때문에 더 효율적이라는 뜻인 것 같다.

 

 

 

Vit performance

 

ViT-H/14의 경우 이미지 분류의 여러 dataset에서 좋은 성능을 보인다. 그리고 pretrain에 있어서 계산적으로 덜 소모된다.  그러나 작은 데이터 셋에 대해선는 큰 모델일 수록 성능이 좋지 않다는 걸 보인다.

그리고 훈련시, random subset의 형태로 훈련을 하는데, 추가적인 정규화는 하지 않고 하이퍼 파라메터는 동일하게 적용을 한다고 한다. 이를 통해서 모델 내부의 특성을 평가할 수 있다고 한다. 

---

 

INSPECTING VISION TRANSFORMER

 

ViT의 첫 번째 layer는 저차원 공간으로 flatten한 patch들을 투영시키는데, Figure 7에서 왼쪽의 이미지와 같다.

 

 

그리고 가운데는 position embedding 간의 유사도를 나타낸 이미지다. 이를 통해서 이미지 내 인코딩을 하는 방법을 학습하는 것을 보여준다. Attention weight를 기반으로 해서 image space 간 평균 거리를 계산한다. 여기서 attention distance는 CNN에서의 receptive field와 같은 효과를 한다. 일부 head 에서는 최하위 레이어에 있는 대부분의 이미지에 attention 을 적용하여 global 하게 모델을 사용할 수 있다는 것을 보여준다. 또한 신경망 깊이에 따라 attention distance 가 증가한다.