[CV 논문 리뷰 스터디 / 6주차 / 박민규] An Image is Worth 16x16 Words, Transformer for Image Recogniton at Scale, 2020

INTRODUCTION

1. Transforemer vs CNN

CNN의 경우엔 멀리 떨어진 정보를 통합하기 위해서 여러 개의 layer를 통과해야 하지만, Transformer는 self-attention layer 하나만 거쳐도 멀리 떨어진 정보를 통합할 수 있다.

2. Inductive Bias

새로운 데이터에 대해 좋은 성능을 내기 위해 모델에 사전적으로 주어지는 가정;

Training에서 보지 못한 데이터에 대해서도 적절한 귀납적 추론이 가능하도록 하기 위해 모델이 가지고 있는 가정들의 집합을 의미한다.

CNN은 지역적인locality 정보 그리고 RNN은 순차적인sequentiality 정보이므로 이에 대한 inductive bias를 가진다. 반면 transformer는 CNN과 RNN에 비해 inductive bias가 낮다는 특징을 지닌다.

ViT에서 MLP(FFN파트)는 patch 구조로 인해 locality와 translation equivariance가 있지만, multi-self attention은 global하기 때문에 CNN보다 image-specific inductive bias가 낮다. 이를 해결하기 위해 2가지 방법을 시도한다.

• Patch extraction : cutting the image into patches
• 이미지를 여러개의 patch로 분리한 후 순서대로 넣는다. position embedding으로 위치 정보 사라지는 것을 방지한다.
• Resolution adjustment : adjusting the position embeddings for images of different resolution at fine-tuning사전 훈련때보다 더 높은 해상도로 fin-tuning 하는 것이 성능에 도움이 된다.
• fine-tuning을 진행할 때 다른 해상도의 같은 이미지로부터 position embedding를 튜닝한다.

Vision Transformer, ViT

NLP에서 사용되는 standard Transformer를 이미지에 그래도 적용하여 이미지 분류에 좋은 성능을 도출한 Vision Transformer ViT를 제안하였다.

ViT는 이미지를 패치로 분할한 후, 이를 NLP의 단어로 취급하여 각 패치의 linear embedding을 순서대로 Transformer의 input으로 넣어 이미지를 분류한다.

ViT의 patch 정보를 추출하는 transforemr encoder 부분은 약간의 조정을 거쳐 그대로 사용하고, 뒷단은 zero-initialized feedforward layer로 대체한다. 이로써 특정 task output을 도출해 내기 위한 head가 된다.

STEP1 : EmBedding Vector

전체 이미지의 크기는 HWC = 48*48*3 이다.

하나의 patch의 크기는 16163 이다. 따라서 전체 이미지에 대해 총 N=9개의 patch가 만들어진다.

하나의 patch를 기존 Transformer에서 하나의 단어라고 생각하자. Embedding을 위해 Trainable linear projection을 통해 $x_p$의 각 패치를 flatten한 벡터를 D차원으로 변환한다. 해당 그림에서 $x_p$를 flatten하게 되면 3*16*16=768이 된다. 따라서 각 patch의 embedding vector는 768차원이 된다.

STEP2 : Add Embedding vector, Class Token with Position Embedding

Class Token : Classification을 위해 사용되는 token으로 BERT에서 CLS token과 같다.

Position Embedding : patch에 위치정보를 제공해준다.

Class Token과 Position Embedding 모두 learnable하며 학습을 통해 결정되는 파라미터이다.

앞서 만든 embedding vector 각각에 해당하는 patch embedding과 위의 그림과 같이 더해지고, Class token은 PE0과 더해진체 젤 앞단에 위치하게 된다.

STEP3 : Transformer Encoder

ViT에선 기존의 Vanilla Encoder에서 수정된 형태이다. 기존의 Transformer encoder는 Attention과 FFN을 통과한 후 Residual block과 Layer 정규화 과정을 거치게 되는데, ViT Transformer에서는 Layer 정규화 과정을 마친 후 Attention과 FFN을 통과하고 Residual block과 더해지는 형태를 띈다.

그 이유는 아래에 링크된 논문에서 자세히 다루고 있다. “learning deep transformer models for machine translation”

빨간 영역 : 1번 patch입장에서 다른 embedding vector들과의 유사도, attention score를 나타냄. 각각의 attetion score와 그에 맞는 value를 곱한 후 이를 더해서(가중합) 최종 score를 구한다.

하나의 head는 64차원이므로 input vector의 768차원과 output 차원을 동일하게 맞춰주기 위해 Self-attention을 12번 수행한다.

ViT Transformer의 FFN 과정은 다음과 같다.

FFN은 MLP로 생각할 수도 있다. c(768)차원에서 4c(4*768)차원으로 들어났다가 다시 c(768)차원이 된다.

STEP4 : Classification

Transformer encoder layer를 12층 쌓아 최종 output를 계산한다. 각각의 patch에 대한 embedding이 나오는데, 이미지 전체에 대한 embedding을 필요로 한다.

Classification token을 도입하여 classification token+PE0로부터 나온 최종 embedding이 image 전체의 embedding이라고 가정하고 모델을 설계한다면, Classification token의 embedding만을 활용하여 Classification이 가능하다.

$Z_{0,0}$을 정규화한 후 MLP를 통과하여 최종적인 classification을 수행하면 최종 output의 차원은 클래스 개수만큼이 된다. (해당 연구에서는 1000개의 class 존재)

Experiments

1. Advanced

• Hybrid Architecture

ViT는 raw image가 아닌 CNN으로 추출한 raw image의 feature map을 활용하는 hybrid architecture로도 사용할 수 있다.

Feature map은 raw image의 locality 정보를 지니고 있으므로 hybrid architecture 패치 크기는 1by1으로 설정하고, 별도의 Positional Embedding을 더하지 않아도 된다.

1by 1 크기의 패치를 사용할 경우 feature map의 공간 차원을 flatten하여 각 벡터에 linear projection을 적용하면 된다.

• Fine-Tuning & Higher Resolution

fine-tuning 할때 pre-training 때와 동일한 패치 크기를 사용한다. fine-tuning 때 higher resolution image를 사용하기 때문에 patch 개수가 많아져 seqeunce도 늘어난다.

ViT는 가변적 길이의 패치들을 처리할 수 있지만, pre-trained된 position embedding은 의미가 사라지므로, pre-trained position embedding을 원본 이미지의 위치에 따라 2D interpolation하여 사용하는 resolution adjustment를 사용하여 inductive bias를 주입한다.

2. Transfer Learning 성능비교

구체적으로 pre-training 때 lower resolution image를 사용하여 학습하고, fine-tuning 때 pre-training때보다 higher resolution image를 사용하여 튜닝한다. 이렇게 된다면 효율적인 사전 학습이 가능해져서 모델의 성능이 높아진다고 알려져 있다.

“Fixing the train-test resolution discrepancy, 2019”

ViT-B, ViT-L, ViT-H는 다음과 같다. Base에서 Huge로 갈수록 embedding vector의 차원이 커지게 되므로 파리미터 수도 늘어난다.

Base CNN은 transfer learning에 적합한 Big transfer(BiT) 구조의 ResNet을 사용하였다.

 

14by14 patch size를 사용한 ViT-Huge와 16by16 patch size를 사용한 ViT-Large의 성능을 baseline으로 비교하였다.

row는 Fine-tuning dataset이고, column은 pre-training dataset이다. 그리고 pre-training dataset 아래의 괄호에 있는 것은 사용한 모델이름이다.

• ViT(Ours-JFT)과 BiT(BiT-L)를 비교해보았을 때 전반적인 성능은 비슷하지만 ViT가 TPU코어를 더 적게 사용하기 때문에 성능이 좋다는 것을 알 수 있다.

3. Category에 따른 성능비교

19-task VTAB classification suite를 3개의 category로 나누어 실험을 진행하였다.

여기서 VTAB은 전체 데이터셋, Natural은 Pets/CIFAR와 같은 과제, Specialized는 의학과 위성 사진, Structured는 지리학적 이해가 필요한 과제이다.

전체 VTAB 데이터 뿐만 아니라 각각의 카테고리에 대해서도 ViT-H/14가 가장 좋은 결과를 도출해냈다.

4. Pre-training Data Requirements

본 실험에서는 pre-training dataset의 크기에 따른 fine-tuning 성능을 확인하였다.

• 데이터가 클수록 ViT가 BiT보다 성능이 좋았고, 크기가 큰 ViT 모델이 더 성능이 좋다.
• JFT를 각각 다른 크기로 random sampling한 dataset을 활용하여 실험을 진행한 결과
• 작은 데이터셋에서 CNN의 inductive bias가 효과가 있어 BiT의 성능이 더 좋았지만, 큰 데이터셋에서는 데이터로부터 패턴을 학습하는 것만으로 충분하였기에 ViT의 성능이 대체적으로 더 좋았다.

5. Inspecting Vision Transformer

다음 그림과 같이 첫번째 patch의 이미지는 position encoding으로 인해 linear projection을 통과하였지만, 위치정보가 그대로 남아있어 유사도가 높은 것을 알 수 있다.

Linear projection을 통과한 output이 CNN의 convolution filter를 시각화한 것과 유사하다.