[5주차 / 임채명 / 논문리뷰] An Image Is Worth 16X16 Words: Transformers For Image Recognition At Scale

Introduction

NLP 분야에서 트랜스포머의 사용은 성능 저하 없이 상상을 초월하는 크기의 데이터 학습을 가능하게 하였지만, CV 분야에서는 아직도 CNN 기반의 모델이 주로 사용되고 있습니다. 본 논문에서는 NLP에서 사용되는 트랜스포머의 구조를 최소한의 수정을 거쳐 이미지에도 적용하려는 시도를 하고 있습니다. 이미지를 patch로 나누고 이 patch의 linear embedding sequence을 트랜스포머의 input으로 사용하고 있습니다. 모델 학습에 있어 supervised 방식을 채택하고 있습니다.

 

Mid-sized 데이터셋을 사용하여 모델을 학습시킨 결과, 비슷한 크기의 데이터셋을 ResNet 기반의 모델에 학습시켰을 때보다 낮은 성능을 보여주었습니다. 이는 트랜스포머가 translation equivariance나 locality 같은 inductive bias가 CNN 구조에 비하여 떨어지기 때문에 예상가능한 결과였습니다. 그러나 사이즈가 매우 큰 데이터셋으로 학습하는 경우, 트랜스포머는 inductive bias를 이기고 훌륭한 결과를 보여줍니다. 

 

Inductive bias란 학습 시에는 만나보지 않았던 상황에 대하여 정확한 예측을 하기 위해 사용하는 추가적인 가정을 의미합니다. 

인접한 픽셀끼리 관련성이 높고 픽셀 사이의 거리가 멀어질 경우 관련성이 적어지는 경우 locality가 있다고 할 수 있습니다.

translation invariance란 input의 위치가 달라져도 동일한 output으로 인식하는 특성을 말합니다.

 

Method

Attention Is All You Need(2017) 트랜스포머 구조

 

본 논문에서는 모델 구조를 기존 트랜스포머의 구조와 최대한 비슷하게 사용하였습니다.

 

기존 트랜스포머는 1차원의 token embedding sequence를 input으로 받기 때문에 2D 이미지 $x$ (HxWxC)를 다루기 위해서 본 논문에서는 이미지를 flattened 2D patch의 sequence $x_{p}$ 로 만들어줍니다. 이때 $x_{p}$는 $ N\times (P^{2}\cdot C) $의 shape를 갖습니다. 

 

(H, W)는 원래 이미지의 크기이며 C는 채널 수를 의미합니다. (P, P)는 패치 이미지의 크기이며 N은 $HW / P^2$로 이미지 패치 개수이자 트랜스포머에 들어갈 input의 sequence length입니다. 트랜스포머는 일정한 크기의 잠재 벡터 D를 모든 layer에서 사용하기 때문에 linear projection(학습 가능)을 사용하여 patch와 map을 D 차원으로 flatten 시켜야 합니다. flatten 시킨 결과를 patch embedding이라고 합니다.

 

$L$은 encoder의 개수를 의미합니다.

 

BERT에서 사용하는 class token처럼, 본 논문에서는 학습 가능한 임베딩을 patch embedding sequence 앞에 추가해줍니다. class token은 classification head에 input으로 들어가는데, classification head는 MLP로 구현되며, pre-training에서는 one hidden layer이며 fine-tuning에서는 single linear layer를 사용한다. 식 4

 

위치 정보(positional information)을 보존하기 위해 position embedding도 patch embedding에 추가해줍니다. 2차원 이상의 임베딩을 사용해도 성능이 뚜렷하게 향상되지 않기 때문에, 1차원 position embedding을 사용합니다. 이렇게 만들어진 임베딩 벡터의 sequence가 encoder의 input으로 사용됩니다. 식 1

 

트랜스포머의 encoder는  multiheaded self-attention 블록(식 2)과 MLP 블록(식 3)으로 이루어져있습니다. Layernorm은 블록 이전에 매번 적용되며 residual connection은 블록 이후에 매번 적용됩니다. MLP 블록은 2 layer로 구성되며, 활성화함수로 GELU를 사용합니다. 

 

Inductive bias 

본 논문에서는 Vision Transformer가 CNN에 비해 inductive bias가 적다고 말합니다. CNN 기반 모델에서는 모든 layer가 locality, two-dimensional neiborhood structure, translation equivariance를 가집니다. ViT의 경우, MLP layers는 local하고 translation equivariance를 가지지만 self-attention layers는 global*합니다. 또한 2D neighborhood structure는 이미지를 패치로 나누는 과정과 fine-tuning 시에 position embedding을 조정하는 과정(아래에서 설명)에만 쓰이고, 초기화 단계에서 position embedding은 패치의 2D position에 대한 정보를 가지고 있지 않기 때문에 이러한 정보는 학습되어야만 합니다.

* A predictive model is called global when it is trained on many different datasets, each being the random outcome of its own stochastic process.

 

Hybrid Architecture

원본 이미지를 patch로 나누는 방법 대신에, CNN feature map을 사용하여 input sequence를 만드는 방법도 있습니다. 이때 patch embedding projection E는 CNN feature map에서 뽑아낸 patch에 적용됩니다. patch의 spatial 사이즈가 1x1인 특별한 경우도 있는데, 이는 단순히 feature map을 flatten하여 사용한 경우입니다. classification imput embedding(class token)과 position embedding은 위와 같은 방법으로 추가합니다.

 

Fine-tuning and Higher Resolution

ViT는 보통 큰 데이터셋으로 pre-train한 후 작은 task에 맞게 fine-tune하여 사용합니다. Fine-tune하기 위해 우리는 pre-trained된 prediction head(MLP classification head)를 제거하고 0으로 초기화된 D x K feed forward layer를 추가합니다. 여기서 K는 새로운 task에서의 클래수 수입니다. 

보통 pre-training 때보다 해상도가 더 높은 이미지를 사용하여 Fine-tuning하는 것이 더 좋은데, 이러한 경우에 원래의 position embedding이 의미가 없어지기 때문에 원본 이미지에서의 위치에 기반하여 position embedding에 2D interpolation을 수행합니다.

 

Experiments

JFT 데이터셋을 사전학습에 사용한 ViT-L이 (같은 데이터셋으로 사전학습을 한) ResNet 기반 BiT-L보다 모든 task에서 더 좋은 성능을 보입니다. 또한 더 큰 ViT-H가 ViT-L보다 더 뛰어난 성능을 보인다는 것을 확인할 수 있습니다. 

 

하이브리드 모델은 모델 사이즈가 작은 경우에는 일반 ViT보다 성능이 향상된 결과를 보이지만, 모델 사이즈가 커지면 ViT와 성능 차이가 없음을 확인할 수 있습니다.

 

Inspecting Vision Transformer

위 사진은 flattened image patches를 linear projection하는 Vision Transformer의 첫 layer에서 주요 28개의 filter를 보여줍니다. 이는 CNN의 앞쪽 레이어를 시각화한 것과 닮았습니다. 즉, low-dimensional한 feature를 잘 포착하고 있음을 보여줍니다.

 

위 사진은 특정 행과 열을 가지는 patch의 position embedding와 나머지 patch의 position embedding 간의 코사인 유사도를 나타내고 있습니다. 거리가 가까운 patch들끼리 더 유사한 position embedding을 가지고 있음을 확인할 수 있고, 이를 통해 이미지 내의 거리를 position embedding이 잘 포착하고 있음을 알 수 있습니다. 또한 같은 행과 같은 열끼리 비슷한 embedding을 가지고 있는 것으로 보아 1D position embedding만으로도 2D 정보를 충분히 잘 포착하고 있음을 확인할 수 있습니다.

 

위 사진은 신경망의 깊이에 따라 관심 영역(attention distance)의 크기가 어떻게 변하는지 보여주고 있습니다. 앞쪽 layer의 경우 attention distance가 넓은 head(global한 특징을 포착)도 있고 attention distance가 좁은 head(local한 특징 포착)도 있는 반면, 네트워크가 깊어질수록 head들이 점차 넓은 attention distance를 가진다는 것을 확인할 수 있습니다. 

 

attention distance는 CNN의 receptive field size와 유사