ViT : An Image Worth 16 x 16 Words : Transformers for Image Recognition At Scale

기존에 이미지에 attention을 취하는 등의 작업은 이루어졌지만, 주로 혼합되어 사용되었습니다. 본 논문은 Transformer를 이미지에 적용하여 처음으로 우수한 성과를 낸 모델이자 시초입니다.

Introduction

Inductive bias란?

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

DNN이 기본적인 요소들의 inductive bias는 아래와 같습니다.
1) Fully connected : 입력 및 출력 element가 모두 연결되어 있으므로, 구조적으로 특별한 relation inductive bias를 가정하지 않습니다.
2) Convolutional : CNN은 작은 크기의 kernel로 이미지를 지역적으로 보며, 동일한 kernel로 이미지 전체를 훑어 본다는 점에서 locality와 transitional invariance 특징을 가집니다.
3) Recurrent : RNN은 입력한 데이터들이 시간적 특성을 가지고 있다고 가정하므로 sequentiality와 temporal invariance 특성을 가집니다.

=> 이러한 가정들로 인한 각 "귀납 편향"이 있기에 특정 역할을 비교적 잘 해낼 수 있는 것 입니다.
하지만! Transformer는 CNN 및 RNN 보다 상대적으로 inductive bias가 낮습니다.

ViT 개요

기존 NLP에서 사용되던 standard transformer를 이미지에 그대로 적용하여 분류에 좋은 성능을 보인 ViT를 제안합니다.

How?) 이미지를 patch로 분할 후 NLP에서 사용되는 토큰과 같이 각 패치의 linear embedding을 순서대로 Transformer의 input으로 넣어 이미지를 분류합니다.

이 ViT를 ImageNet-1k로 학습했을 때, 비슷한 크기인 ResNet보다 낮은 정확도를 보였기에, ViT가 CNN보다 inductive bias가 낮음을 알 수 있습니다.
=> 하지만 large scale data set을 학습했을 때는 SOTA성능을 보였습니다. 이는 데이터 셋의 크기가 커질 수록 CNN보다 낮은 inductuve bias로 인해 낮은 성능을 보였던 것이 해소됨을 확인할 수 있는 것 입니다!

Method

ViT는 위와 같으며, 5가지 step으로 구성됩니다.

1) 기존 NLP 표준 transformer는 "1D 토큰 임베딩 시퀀스"를 입력으로 받습니다. 그렇기에 2D이미지를 flatten 2D 패치 시퀀스로 reshape 해줍니다. 이미지 x∈R^(H×W×C)를 xp∈R^(N×(P^2ㆍC))로 변형해줍니다. 각 패치 sequence xp를 만들어내는 것 입니다.
ex) (300,300,3) 이미지를 (100,100,3) 이미지 9개로 자름

(H,W) : 원본 이미지의 해상도
C : 채널의 수
(P,P) : 각 이미지 패치의 해상도
N = HW/P^2 : 최종 패치의 개수

2) 트랜스포머는 내부의 모든 레이어들의 latent vector의 사이즈를 D로 통일시켰습니다. 그렇기에 2차원 patch들을 1차원으로 flatten하고 (:NㆍP^2ㆍC size), 이를 trainable한 linear projection을 거쳐 D차원 벡터로 매핑시켜줍니다.
이 D차원 벡터가 patch embeddings들 입니다.
ex) (100,100,3) -> (1,100ㆍ100ㆍ3) 인 1D input값(하나의 patch embedding)으로 생성

3) 여기서 Class token방법이 나오는데, 이는 BERT의 CLS token과 유사합니다. 본 논문에서는 하나의 학습 가능한 Class token을 임베딩된 패치들의 맨 앞에 하나의 학습 가능한 class token embedding vector를 추가했습니다.
:이 Class token embedding은 Encoder의 여러 층을 거쳐 최종 output이 나왔을 때, 이미지에 대한 1차원 representation vector로써 역할을 수행합니다.

'Learnable class 임베딩'과 '패치 임베딩'에 1D learnable position embedding을 더합니다. positional embedding을 더해줌으로써 inductive bias를 조금 주입해줍니다!

여기서 나오는 positional embedding은 아래의 네 가지 시도 후, 가장 좋은 1D position embedding을 사용한 것 입니다.

raster order : 왼쪽 위 부터 오른쪽 아래까지 순번을 매김
Rel.Pos.Emb : 패치들 간의 거리를 고려한 임베딩

4) 이 임베딩된 1D input값을 vanilla Transformer의 encoder에 넣어 Multi-head Self-Attention(MSA) 을 여러번 수행 후, 마지막 layer에서 class embedding에 대한 output인 image representatoin을 도출합니다. 앞에서 말했듯 Class token이 트랜스포머의 여러 인코더 층과 layer normalization을 거쳐 최종 output y가 나왔을 때, 이 y가 이미지에 대한 1D representation vector로 역할을 수행하는 것 입니다.

5) 위의 MLP는 y를 flatten 후 단순 linear layer를 통과시켜 classification을 수행합니다. 다시 말해, MLP에 image representation을 input으로 넣어 이미지의 클래스를 분류하는 것 입니다.

Transformer Encoder 구성

ViT는 Mutil-head Self Attention(MSA) 와 MLP block으로 구성되어 있습니다.
MLP는 2개의 layer를 가지며, GELU activtion ft을 사용합니다.
각 block의 앞에는 Layer Norm(LN)을 적용(차이점 : 기존은 각 block 뒤에 적용했었음)하고, 각 block뒤에는 residual connection을 적용합니다.
z0 : D차원으로 매핑시켰던 vector에 positional임베딩을 더한 것
이를 LN->MSA->RC => LN->MLP->RC => LN 통해 최종적인 y값을 구성하게 되는 것 입니다.

Hybrid Architecture

ViT에 이미지를 넣을 때, raw image가 아닌 CNN으로 훑어서 추출한 feature map을 활용하여 이를 패치로 분할하고 넣는 Hybrid Architecture 방식으로도 사용할 수 있습니다. Feature map에는 이미 raw image의 공간적 정보를 포함하고 있기에 바로 사용해도 되는 것 입니다.

Fine-tuning and Higher Resolution

위의 두 가지 방법 중 하나를 선택하여 Large scale로 pre-training을 진행한 후, 해당 모델을 downstream task에 맞게 Fine-tuning을 진행하게 됩니다.

ViT를 Fine-tuining할 때, pre-training을 진행할 때 사용했던 인코더를 나와 통과하는 MLP를 downstream task에 맞게 zero-initialized feedforward layer로 대체합니다. 이렇게 구성한 후 어디까지 freeze시킬 지 정하고 fine-tuning을 진행하게 되는 것 입니다.

pre-training과 동일한 패치 크기를 사용하기에, 고해상도 이미지로 fine-tuning을 한다면 sequence의 길이가 더 길어집니다. 물론 ViT는 가변적인 길이의 패치들을 처리할 수 있지만, pre-trained position embedding의 의미가 사라집니다. 그렇기에 이럴 때는 pre-trained position embedding을 원본 이미지의 위치에 따라 2D interpolatoin하여 사용합니다. : Resolution Adjustment 장치
이는 그마나 있었지만 사라질 뻔한 inductive bias를 조금이나마 주입할 수 있는 것 입니다!

즉, Fine-tuning시에는 MLP head만 변경을 하여 모델을 구축하고, 만약 고해상도 이미지를 사용하게 된다면 positional embedding의 의미가 사라짐을 방지하기 위해 2D interpolation인 Resolution Adjustment 장치를 사용하여 inductive bias를 주입해주는 것 입니다!

Experiments

Datasets

총 세 개의 데이터셋을 pre-training에 사용했습니다.

Model Variants

아래와 같이 3가지 volume에 대해 실험을 진행했습니다. 이와 비교할 Baseline CNN은 batch norm layer를 group norm으로 변경하고 standardized conv layer를 사용하여 transfer learning에 적합한 Big Transformer(BiT) 구조의 ResNet을 사용하였습니다.

실험 결과

본 실험에서는 JFT데이터 셋으로 pre-training을 진행한 14x14 패치 크기를 사용한 ViT-Huge와 16x16 패치 크기를 사용한 ViT-Large의 성능을 baseline과 비교했습니다.
모든 downstream task에서 ViT-L/16의 성능이 BiT-L의 성능을 능가하였고, 이에 더하여 사전훈련에 드는 계산상의 cost는 확연하게 낮음을 확인할 수 있습니다.
더 큰 모델인 ViT-H/14는 더욱 향상된 성능과 낮은 pre-train 비용을 보여줍니다. 특히 어려운 데이터셋인 ImageNet, CIFAR-100, VTAB 등에서 두드러집니다.

두 번째 표는 19가지 task(19-task VTAB classification suite)를 3개의 그룹으로 나누어 추가 실험을 진행한 결과입니다.(Natural, Specialized, Structured)
이도 ViT-H/14가 가장 좋은 결과를 도출하였습니다.

다음 실험으로는 Pre-training dataset의 크기에 따른 fine-tuning 성능 비교입니다. 각 dataset에 대해 pre-training한 ViT를 ImageNet에 전이학습을 한 정확도를 확인한 결과, 아래에서 보시다시피 데이터가 클 수로 ViT의 성능이 BiT보다 우수하며 크기가 큰 ViT모델이 효과가 있음을 볼 수 있습니다. 이는 작은 데이터셋에서는 CNN의 inductive bias가 효과가 크지만, 큰 데이터셋에서는 큰 영향을 미치지 않음을 알 수 있습니다.

다음은 JFT pre-train 기반으로 cost대비 transfer 성능을 검증하여 모델드르이 scaling study를 진행한 것 입니다. ViT가 "성능<->cost trade-off"에서 ResNet(BiT)보다 우세하며, Cost가 증가할 수록 Hybrid와 ViT의 "성능<->cost trade-off"차이가 감소함을 확인할 수 있습니다.

마지막으로 ViT가 어떻게 이미지를 처리하는지 이해하기 위한 실험을 진행하였습니다.
맨 왼쪽의 사진은 flatten 패치들을 2차원으로 projectin한 후에 그 행렬에 대한 주성분을 분석한 것 입니다.
가운데 사진은 패치 간 position 임베딩의 유사도를 통해 가까운 위치에 있는 패치들의 position 임베딩이 유사한지 확인한 것입니다.(같은 행/열에 있는 패치는 유사한 임베딩을 갖습니다)
우측 사진은 ViT에서의 self-attention이 얼마나 수용할 수 있는지 확인하기 위하여 attention distance를 확인한 결과입니다. 초반 layer에서도 attention을 통해 이미지 전체의 정보를 통합하여 사용함을 알 수 있으며, layer가 깊어질 수록 attentino거리는 증가하며 모든 이미지에 attention하고 있는 것을 확인할 수 있습니다. 이는 정보를 global하게 integrate할 수 있는 능력을 보여줍니다. 즉, 특정 부분은 전체적으로 보고 일정 부분은 특정 요소로 파악하는 것입니다.

Conclusion

본 논문에서는 Transformer를 이미지 인식 task에 직접적으로 적용하였습니다. CV분야에서 self-attention을 사용한 적은 있지만, 본 논문은 patch extaction(positional embedding)과 Resolution Adjustment을 제외하면 귀납 편향을 반영하지 않고도 우수한 성능을 보였습니다. 이렇게 단순하면서도 확장성 있는 방법은 large scale dataset에서 사전학습과 더불어 사용되었을 때 좋은 성능을 보입니다. Vision Transformer가 많은 이미지 분류 데이터셋의 SOTA모델들과 필적하거나 뛰어 넘으며, 이에 더해 cost 또한 절감할 수 있었습니다.

이제는 이미지 분류를 넘어 ViT를 detection이나 segmentation에도 적용하는 시도가 필요하다고 합니다.

참고 :https://velog.io/@xuio/%EB%85%BC%EB%AC%B8%EB%A6%AC%EB%B7%B0-CV-ViT-An-Image-Worth-16-x-16-Words-Transformers-for-Image-Recognition-At-Scale
https://velog.io/@choonsik_mom/Vision-TransformerViT-%EB%85%BC%EB%AC%B8-%EB%A6%AC%EB%B7%B0
https://zeuskwon-ds.tistory.com/95
https://www.youtube.com/watch?v=0kgDve_vC1o&t=536s