[논문 리뷰] ViT, An Image is Worth 16x16 Words

Vision transformer 논문에 대한 간단한 리뷰를 정리해보겠습니다.

 

Paper : 

https://arxiv.org/abs/2010.11929

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

 

- Background

Attention mechanism과 transformer model은 자연어 처리에서 이미 주류를 이루고 있었습니다. 

 

본 논문에서는 NLP 분야에서의 transformer의 활약에 영감을 받아 vision 영역에 transformer를 적용한 내용을 다루고 있습니다.

현재는 vision 영역에서도 transformer 기반 모델들이 상당히 많이 등장했고, 많은 SOTA 모델이 trasnformer 기반으로 알고 있습니다.

 

최대한 NLP 에서 사용되는 transformer model을 그대로 사용하는 방식으로 실험을 진행했는데,

'An Image is Worth 16x16 words'라는 논문의 제목에서도 볼 수 있듯이 기존 NLP 분야의 transformer에서 문장의 각 단어를 token으로 다룬 것처럼 이미지의 각 영역을 token으로 다루는 방식을 사용했습니다.

 

덕분에 NLP의 transformer 구조를 거의 그대로 사용하고 있습니다.

(이전에도 transformer를 적용하려는 연구들이 있었지만, pixel 단위로 attention을 계산하는 등의 방식은 비효율적이었고 본 논문의 방식이 가장 좋은 결과를 보였다고 합니다.)

 

또한 NLP 분야에서 large dataset에 training한 후 smaller task에 fine-tuning하는 방식을 주로 사용한 것처럼, 

ViT도 large dataset에 pre-train을 수행했을 때 비로서 좋은 성능을 보이고 SOTA를 달성하게 되는 결과를 보입니다. 

 

이에 대해 CNN과 달리 ViT에는 inductive bias가 부족하기 때문이라고 설명하고 있는데, 여기서 inductive bias란 간단히 '이미지 처리에 특화된 가정'이라고 얘기할 수 있습니다. 

 

직관적으로 이해하면 CNN은 이미지 처리를 위해 개발된 구조이지만 transformer는 그렇지 않다고도 이해할 수 있겠고, 

locality나 translation equivariance 등에 대한 자세한 설명, 수식에 대해서는 아래의 링크에 자세한 설명이 나와있습니다.

https://seoilgun.medium.com/cnn%EC%9D%98-stationarity%EC%99%80-locality-610166700979

CNN과 이미지가 찰떡궁합인 이유

결론적으로 transformer는 inductive bias가 부족하여 적은 수의 데이터로는 일반화가 잘 이뤄지지 않지만, 

충분히 큰 데이터셋에 pre-train 시켰을 때 CNN 기반의 모델들보다 좋은 결과를 보이게 됩니다. 

 

- Method

본격적으로 Vision Transformer가 어떻게 동작하는지 살펴보겠습니다. 

 

자세한 동작에 대해서는 

   - [Paper Review] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows,             (https://www.youtube.com/watch?v=2lZvuU_IIMA),

   - [DMQA Open Seminar] Transformer in Computer Vision,

(https://www.youtube.com/watch?v=bgsYOGhpxDc) 

 

위 두 영상의 설명을 많이 참고했습니다.

 

ViT 구조

ViT의 구조는 위와 같고, 크게 Embedding, Transformer Encoder, MLP head 3가지 부분으로 나눠서 살펴보겠습니다. 

 

1. Embedding

Input image를 token화시키고 transformer encoder의 input으로 만들어주는 과정입니다. 

N x N x C 크기의 image가 input으로 들어오면 각각의 크기가 h x w x c 크기인 patch들로 그림과 같이 나눠주고, 다시 각각의 patch들을 크기가 h*w*c인 vector로 만들어줍니다. 

 

각각의 vectors들을 D 차원으로 projection 시켜준 후, classfication token을 추가하고 positional embediing을 더해주면 그 결과가 transformer encoder의 input이 됩니다. 

 

Classification token은 classification을 위해 추가된 token으로, 해당 token에 대한 transformer encoder의 output을 classification에 사용하게 됩니다.

 

또한 positional embedding은 이미지에서의 위치 정보를 다시 추가하기 위해 더해주는 값이라고 볼 수 있겠습니다.

 

이때 classification token과 positional embedding 모두 trainable한 parameters입니다. 

 

2. Transformer Encoder

Transformer Encoder는 위와 같은 blocks가 L개 반복되는 구조입니다. 

LN - MSA - LN - MLP가 반복되고, residual connection이 추가되어 있습니다.

 

    ● Layer Norm

Norm은 Layer Normalization으로 batch 단위로 normalizatoin을 수행하는 BN과 달리, 각 데이터 별로 normalization을 수행합니다. 

 

    ● Multi-Head Attention

다음으로 Multi-Head attention의 연산 과정입니다. 앞서 생성한 N개의 D차원 vectors이 LN을 거친 결과가 input으로 전달되면 D x $D_h$ 크기의 query, key, value weight와 각각 곱해줍니다.

 

그렇게 생성된 N x $D_h$ 크기의 matrix가 각각 query, key, weight이고,

$A = QK^T$, N x N matrix에서 $ A_{ij}$가 i번째 token이 j번째 token에 얼마나 attention을 갖는지를 의미합니다. 

 

최종적으로 $ QK^T$를 $ \sqrt{D_h}$로 나눠준 값에 softmax를 취하고(각 row별로), value matirx와 곱해준 값이 head의 output이 됩니다.

 

따라서 각 head별 출력은 Nx$D_h$가 되고, 이때 $ D_h$는 보통 D / k(=number of heads)로 설정되어 k개의 heads의 output을 concat한 결과로 각 MHA의 출력은 NxD크기가 유지됩니다.

※ 수정) k*$ D_h$를 D 차원으로 projection한 결과가 최종 output입니다.

 

 

   ● MLP

MLP는 hidden layer가 하나인 two layer network이고, activation function으로 GELU nonlinearity를 사용합니다.

 

이렇게 LN - MHA - LN - MLP 구조가 L번 반복되고, 최종적인 출력은 0번째, classification token에 대한 output이 됩니다. (N x D matrix 중 첫번째 행)

 

   ● MLP Head

최종 classification을 수행하는 MLP로, cls token에 대한 transformer encoder의 output을 input으로 받아 classification을 수행합니다.

 

Hidden layer가 1개인 2 layer network입니다. 

 

- Results

Pre-training dataset의 수에 대한 영향 등에 대한 추가적인 실험들이 있지만, 해당 실험에 대해서는 생략하고 최종 성능만 살펴보겠습니다. ResNet이나 EfficientNet 등의 SOTA models과 비교해도 더 좋은 성능을 보입니다. 

 

개인적으로 transformer나 attention mechanism을 ViT를 통해서 처음 접해봤기 때문에, ViT의 처리과정은 이해가 가더라도 왜 좋은 성능을 보이는지 이해가 잘 가지 않았는데, 동작에 대한 해석을 살펴보겠습니다. 

● Principal Components of Embedding filters

먼저 embedding filters에 대한 해석입니다.

위 이미지는 embedding filters(projection matrix)에 PCA를 적용한 결과로, 마치 CNN에서 초반 layers의 filters와 유사한 것을 살펴볼 수 있습니다.

 

● Positional Embedding

다음으로 position embedding의 유사도를 나타낸 그림입니다. 

Position embedding이 학습된 결과 row와 column이 같고 가까울 수록 높은 유사도를 보입니다.

즉, transformer가 이미지에 대해 잘 이해하고 있다고 해석할 수도 있을 것 같습니다.

 

● Attention

먼저 좌측의 그림은 attention map을 시각화한 결과입니다. 

Attention map은 모든 layers의 attention을 종합한 것이라고 보면 되고, 결과적으로 model이 classification에 중요한 정보를 담고 있는 영역에 집중하는 것을 확인하실 수 있습니다.

 

우측의 그림은 mean attention distance를 나타낸 그림으로 CNN에서 receptive field가 점차 global한 영역으로 확대되듯이, ViT에서도 attention이 global한 영역으로 점점 확대된다는 것을 확인할 수 있습니다.