[논문 리뷰] ViT / An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

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

** YAI 9기 박찬혁님이 비전논문팀에서 작성한 글입니다.

---

Introduction

  현재 Transformer와 같은 Self-attention-based 구조들이 NLP에서 주된 방식이다. 이 방법들은 주로 큰 모델에서 사전학습을 진행 후에 작은 데이터셋으로 Fine tuning하여 사용하는 방식으로 쓰인다.

  이 Transformer 방식을 이미지에 바로 적용시키는 것이 이 논문의 목적이다. 이미지를 patch라는 작은 단위로 나누어서 NLP의 token과 같은 방식으로 간단한 imbedding 후에 Transformer에 집어넣는다.

  이 방식은 기존 CNN이 이미지에 강했던 이유인 inductive bias가 거의 없는 형태로 학습을 진행하게 된다. Inductive bias에 대해선 아래에서 설명한다. 이 inductive bias가 없기 때문에 입력 이미지에서 대상의 위치가 달라진다거나 방향이 달라지는 것에 대해서 효율적으로 대처할수가 없게 된다. 따라서 ViT는 크기가 작은 dataset에 대해서는 성능이 그렇게 좋지 못한것을 볼 수 있다. 하지만 데이터셋이 커지면 커질수록 그 성능은 inductive bias가 있는 모델을 능가한다.

 

• Inductive Bias
• 학습때 만나지 못했던 상황을 해결하기 위해 대입된 추가적인 가정이다. CNN의 경우 Convolution Filter가 지역적으로 정보를 얻어서 만약에 이미지 상에서 대상의 위치가 바뀌었다고 하더라도 인식이 가능하다. 따라서 inductive bias가 크다고 볼 수 있다. 반면에 FCN같은 경우 모든 뉴런들이 연결되어있기 때문에 CNN에 비해 상대적으로 이미지의 위치가 고정되어있다고 볼 수 있다. 정리하면 모든 요소와 연결되어있으면 inductive bias가 낮다고 표현할 수 있다. 따라서 global한 정보가 많이 필요한 경우는 inductive bias가 낮은 transformer와 같은 방법이 좋고, 지역적인 정보가 많이 필요한 경우는 CNN과 같은 방법이 일반적으로 좋다고 할 수 있다.

---

ViT

  우선 ViT(Vision Transformer)의 구조는 아래와 같다.

  ViT는 기존의 Attention is all you need라는 논문에서와 거의 같은 모델을 사용한다. 다른점은 우선 자연어 처리가 아닌 이미지의 처리이기 때문에 입력되는 데이터의 형태가 다르다. 자연어 처리의 tokenizing과 같이 이미지를 patch로 쪼개고 그 이미지들을 모두 1차원으로 flatten한다.
  그래서 원본 이미지의 형태가 $ \mathbb x \in \mathbb R^{H \times W\times C} $이고, patch의 형태가 $ \mathbb x_{p'} \in \mathbb R^{P\times P\times C} $ 였다면 그 패치를 1차원으로 펼친게 $ \mathbb x_{p'} \in \mathbb R^{1\times (P^2\cdot C)}$ 형태이고 ViT에 입력되는 형태는 각 패치의 1차원 형태를 이어붙인 $\mathbb x \in \mathbb R^{N\times (P^2\cdot C)}$의 형태가 된다. 여기서 N은 패치의 개수이다 ($ N=HW/P^2 $).

위 과정 이후에 아래와 같은 식을 통과시킨다.

  위 식을 보면 각각 flatten된 patch에 E가 곱해지는 것을 볼 수 있다. E는 임베딩 벡터로써 $ P^2 \cdot C\times D $의 형태를 가지고 있다. 따라서 하나의 patch($ \mathbb x_{p'} \in \mathbb R^{1\times (P^2\cdot C)} $)와 곱해지게 되면 $1 \times D$의 형태가 된다. 제일 첫번째로 들어가는 $\mathbf x_{class}$라는 벡터는 classification이 첫번째 출력에서 이루어질 수 있도록 넣어주는 벡터이다. 따라서 총 벡터의 크기는 $\mathbb{ R }^{ \left( N+1 \right) \times D }$가 된다. 여기다가 $\mathbf E_{pos}$라는 parameter를 더해주게 되는데, 이는 입력의 위치 정보를 나타나게 해주는 parameter이다. 기존의 Transformer에서는 cos, sin으로 일정한 값을 넣어줬지만 여기서는 위치 벡터도 학습이가능한 파라미터로 표현된다.

  Transformer Encoder는 Multi-Head self Attention과 MLP로 이루어져 있다. 특징은 각 블럭의 전에 Normalization을 실행하고 블럭 뒤에는 residual connection을 적용시켰다. 여기서 MLP는 GELU를 activation function으로 가진 2개의 레이어를 포함한다.

그 후 출력의 첫번째인 $z^0_L$에 아래와 같이 Normalization을 해주면 image representation인 y가 만들어진다.

$$
\mathbf{y}=\mathrm{LN}\left( \mathbf{ z }_{ L }^{ 0 } \right)
$$

이렇게 만들어진 image representation을 마지막의 MLP Head에 넣으면 classification이 완료된다.

• Multihead Self Attention(MSA)
  Self attention은 아래와 같은 식으로 쿼리(Q), 키(K), 밸류(V) 사이의 관계성을 추출하는 과정이다.이때 각 Q,K,V는 아래의 식으로 만들어진다.이때 이 계산과 학습을 한번에 수행하지 않고 여러개의 head로 나누어서 병렬로 계산하고 학습시키는 방법이 MSA이다.
• $$
  Q=X\times W_Q,　 K=X\times W_K,　 V=X\times W_V
  $$
• $$
  Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_K}})V
  $$
• Hybrid Architecture
    모델의 입력을 이미지를 자른 patch로 넣는 것이 아닌 CNN feature map을 넣는다. 이때 map의 사이즈는 1*1이 될 수도 있고, 이 또한 flatten하여 입력한다. 이 논문에서는 ResNet50의 4stage의 결과와 레이어 개수를 같게 유지한 3stage의 결과의 두가지 경우를 사용하였다. 이때 3stage의 결과는 이미지의 크기가 크기 때문에 sequence가 더 길어진다.

Setup

Pre-train dataset

• ImageNet (1k classes, 1.3M images)
• ImageNet-21k (21k classes, 14M images)
• JFT (18k classes, 303M images)

Transfer learning dataset

• ImageNet
• ImageNet-Real
• CIFAR 10
• Oxford-lllT Pets
• Oxford Flowers-102
• 19-task VTAB classification suite

Model Variant

ViT의 성능 실험은 위와 같이 크게 세 모델로 나누어 진행하였다.

Pre-train에서는 베이스라인 모델인 BiT를 포함한 모든 모델을 아래의 조건들 하에서 학습시켰다.

• Optimizer: Adam( $\beta_1=0.9, \beta_2=0.999$)
• batch size: 4096
• high weight decay 0.1
• image resolution: 224

Fine-tuning 단계에서는 아래와 같은 조건을 사용했다.

• Optimizer: SGD with momentum
• batch size: 512
• image resolution: 512 for ViT-L/16, 518 for ViT-H/14
• Polyak & Juditsky averaging with factor 0.9999

  Fine tuning할 때 Pre training 에서 보다 높은 해상도의 이미지를 사용하는것이 더 좋다. 따라서 위와 같이 Fine tuning 단계에서는 더 큰 이미지를 사용했다. 이때 같은 패치 사이즈를 사용하면 입력 값의 갯수가 많아진다. 개수가 많아지는 것은 Transformer에서는 큰 문제 없이 수용이 가능하지만 입력때 더해주는 position embedding의 사이즈가 달라져야한다. 따라서 Position embedding에 이미지의 위치를 고려해서 2d interpolation을 해준다.

Experiment

  위 표는 ViT의 모델들과 baseline 모델인 BiT-L과 Efficient net인 Noisy Student 모델을 비교한 것이다. 각 모델들은 JFT 데이터셋으로 Pre-Training후 각각 위의 Transfer learning 데이터셋으로 Fine tuning을 한 결과이다. 결과를 보면 ViT모델 중에 가장 큰 모델인 ViT-H가 대부분의 데이터셋에서 성능이 가장 높은 것을 확인할 수 있다.

  아래의 TPUv3-core-days는 학습하는데 걸리는 시간을 나타낸 것이다. ViT-L/16같은 경우를 보면 같은 데이터셋인 JFT를 학습하는데 걸린 시간이 BiT-L에 비해 14~15배 정도 적게 걸린 것을 확인할 수 있다. 이처럼 ViT는 이전의 SoTA였던 BiT에 비해 성능이나 학습 시간에서 더 우월한 결과를 가지는 것을 볼 수 있다.

  위는 ViT모델끼리 Pre-train, Fine tuning의 데이터셋을 바꿔가며 비교한 결과이다. 학습 데이터셋이 많을수록 확실히 성능이 좋아지는 것을 확인할 수 있다.

  위 왼쪽 그래프는 Pre-training datase에 따른 성능이다. ImageNet은 데이터셋의 개수가 1k개, ImageNet-21k는 21k개, JFT-300M은 303M개이다. 그래프를 보면 데이터셋의 개수가 적은 ImageNet에 대해서는 기존 모델인 BiT가 우세한 모습을 보인다. 하지만 데이터셋의 크기가 커질수록 ViT기반 모델들이 성능이 계속 더 좋아지는 것을 볼 수 있다.

  오른쪽 그래프는 모두 같은 데이터셋인 JFT에서 pre-training에 쓰이는 데이터의 개수를 조절한 결과이다. 여기 또한 데이터의 개수가 낮을땐 BiT가, 데이터의 개수가 많을땐 ViT가 높은 것을 확인할 수 있다.

  이런 결과가 나오는 이유는 위에서 설명했듯이 ViT의 경우 resnet보다 inductive bias가 적기 때문에 더 많은 양의 데이터를 학습해야 다양한 형태의 test 데이터셋에 대응할 수 있다. 하지만 많은 데이터를 학습했을 경우에는 inductive bias가 큰 모델보다 더 좋은 결과를 보임을 위의 결과에서 확인할 수 있다.

  위 그래프는 사용된 컴퓨팅 자원에 따른 성능을 나타낸 그래프이다. 여기서는 ViT, BiT, Hybrid의 세 가지 모델을 비교하였다. 결과를 보면 BiT같은 경우는 계속 성능이 상대적으로 낮고 Hybrid가 초반엔 가장 성능이 높은 것을 볼 수 있다. 하지만 투입된 컴퓨팅 자원이 커질수록 Hybrid와 ViT가 거의 같은 성능으로 수렴하는 것을 볼 수 있다. 이를 해석하면 자원이 한정되어 있는 경우에는 Htbrid 모델이 더 높은 성능을 보인다.

위 그래프는 모델의 전이학습을 평가하는 데이터 세트인 VTAB에 대해 학습한 결과를 나타낸 것이다. 전이 학습에 대해서도 ViT가 BiT를 능가하는 성능을 보임을 알 수 있다.

Inspection Vision

논문에서는 ViT 내부 과정을 이해하기 위해 몇몇 분석을 시도했다.

  위의 왼쪽 사진은 내부의 embedding filter를 가시화한 것이다. 여기서 보통의 CNN과 같이 일정한 패턴들을 학습한 것을 볼 수 있다. 학습 자체는 전체적인 관계를 보는 Attention을 사용했지만 CNN과 비슷한 결과가 나타난 것으로 보아 각 embedding filter의 기능이 CNN의 필터와 유사하게 작동하고 있다는 것을 알 수 있다.

  오른쪽 사진은 Position embedding의 유사도를 이미지로 시각화 한 것이다. 결과를 보면 가까운 patch끼리는 embedding의 유사도가 높아 비슷한 형태로 학습이 되었다는 것을 알 수 있다. 이것은 Input patch간의 공간 정보가 잘 학습된 결과라고 볼 수 있다.

  위 그래프는 ViT와 Hybrid 모델에서 244*244의 이미지를 학습했을 때 네트워크 각 층의 head별로 attention의 평균 거리를 나타낸 결과이다. 왼쪽의 그래프를 보면 낮은 층일 경우 Attention의 거리가 작은 head가 존재하지만 층이 올라갈수록 평균 거리가 120정도까지 점점 커지는 것을 볼 수 있다. 이는 낮은 층에서는 거리가 가까운 픽셀들과의 특징을 추출하는 모습을 볼 수 있다.

  오른쪽 그래프는 Hybrid 모델의 결과인데 이때는 낮은 층에서 ViT모델보다 Attention 거리가 짧은 head들이 상대적으로 없는 것을 볼 수 있다. 이는 Hybrid의 경우 국소적인 특징을 추출한 Feature map을 입력으로 넣어서 원래 지역적인 특징을 추출하던 낮은 층의 head들이 인접 픽셀을 고려할 필요성이 사라졌다는 의미로 해석할 수 있다.

  위의 세가지 경우에서 ViT는 Inductive bias가 기존의 CNN구조보다는 적지만 position embedding이나 attention의 거리를 봤을 때 Inductive bias가 일부분 존재한다는 것도 알 수 있다.

Conclusion

  이번 ViT 논문은 ResNet에서 머물러있던 Computer Vision을 NLP에서 사용하던 Transformer를 가져와서 더 높은 성과를 냈다는 점에서 큰 변화를 보여주었다. 기존의 CNN이 가졌던 Filter의 특성 없이도 모든 부분을 보는 Attention만으로도 기존의 CNN 구조의 성능을 넘었다는 데에서 이미지 인식의 또 다른 방향성을 연 경우라고 생각된다.

  Transformer를 사용한 모델은 Pre train에 사용되는 학습 이미지는 다른 모델보다 많이 필요하지만 사전학습에 들어가는 시간이나 비용은 기존의 CNN 모델보다는 적다는 점에서도 의미있는 결과라고 할 수 있다. 또한 이 모델은 크기를 올렸을 때도 뚜렷한 saturation이 확인되지 않아서 크기 조절에 따른 성능 향상도 기대해 볼 수 있다고 한다.

  Self-supervised learning에서도 ViT의 labeled pre-training에는 미치치 못하지만 기존의 CNN 구조보다는 높은 성능을 보여주었다. 단순히 BERT를 모방한 경우에도 성능 향상을 보였기 때문에 앞으로도 비용이 높은 큰 데이터세트가 아니더라도 높은 성능을 얻을 수 있다는 가능성을 제기했다.