[논문리뷰] AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

  오늘 리뷰할 논문은 2021년 Google에서 발표된AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE라는 논문입니다. Vision Transformer 혹은 ViT라는 모델을 제안한 페이퍼입니다. NLP 분야에서 엄청난 영향을 끼친 transformer 모델을 vision 분야에 성공적으로 적용한 페이퍼입니다. 제목에서도 볼 수 있듯이 수많은 픽셀 값으로 이루어진 이미지를 NLP에서의 token처럼 바꾸어 사용하기 위해 이미지를 16x16 patch로 쪼개서 사용하였습니다. 

(https://arxiv.org/abs/2010.11929)

(해당 글은 개인적인 기록을 목적으로 잘못된 내용이 있을 수 있음을 알립니다.)

 

0. Abstract

  해당 페이퍼의 개제일을 기준(ICLR2021)으로 Vision분야는 CNN기반의 구조가 지배적으로 차용되고 있습니다. 그런데 CNN구조는 attention mechanism을 완전하게 적용하기가 힘든 구조입니다. 저자들은 NLP분야에서 이제는 업계표준처럼 사용되는 transformer 구조를 Vision분야로 가져와서 방대한 양의 데이터를 통해 pretrain한 뒤 task에 맞는 비교적 작은 데이터셋에 fine-tuning을 하는 Vision Transformer(ViT)를 제안하였습니다. 결과적으로 최신 CNN기반 SOTA과 비교해도 손색 없는 우수한 성능을 보인다고 합니다. 

 

 

1. Introduction

  NLP분야에서 압도적인 선택을 받고 있는 transformer는  거대 데이터셋을 가지고 pretrain을 하는 패러다임이 제시되면서 성공적인 성능을 보이기 시작했습니다. 또한 CNN계열의 모델과는 다르게 파라미터의 크기를 성능저하 없이 효과적으로 늘릴 수가 있었습니다. 덕분에 100B사이즈의 모델도 학습이 가능했다고 합니다. 이러한 현상이 attention mechanism에서 기인했다고 판단해서 CNN 형태에서도 attention mechanism을 적용하려 했지만 성공적이진 않았다고 합니다. 

 

  해당 페이퍼에서 저자는 이미지 도메인에서도 transformer를 NLP에서 처럼 사용하기 위해서 이미지를 여러 개의 patch로 쪼개서 각각의 patch를 하나의 token처럼 사용하기로 합니다. 그리고 transformer에서는 CNN보다 inductive bias가 적기 때문에 더 많음 양의 데이터셋으로 먼서 모델을 pretrain할 필요가 있었다고 합니다. 구체적인 성능수치는 뒤에서 다루도록 하겠습니다. 

 

 

2. Related Work

해당 부분은 넘어가도록 하겠습니다. 

 

 

3. Method

ViT 구조

ViT는 위와 같은 구조로 되어있습니다. 아래에서 자세하게 설명을 해보겠습니다. 

 

3.1. Vision Transformer (VIT)

  이미지는 위 처럼 여러 개의 겹치지 않는 격자 patch로 만들어줍니다. 그리고 patch를 일열로 펴서 patch마다 linear projection을 해줍니다. 물론 이 linear projection layer는 모든 token들이 하나의 동일한 linear projector를 사용합니다. 실제 코드 상으로는 linear projection layer에 fully connected layer 대신 CNN layer 하나를 사용하는 경우가 많습니다. 연산 상 더 효율적이지만 성능 하락은 거의 없어서 그런 것으로 알고 있습니다. 그리고 patch 맨 앞에는 0으로 채워진 cls token이 들어가게 됩니다. 이렇게 num_patch+1개 만큼의 patch가 만들어지고 여기에 sinusoidal positional embedding을 더해줍니다. Positional embedding은 transformer가 샘플(patch)을 입력으로 받을 때, 해당 patch가 전체 이미지 중 어디 쯤에 위치하는지는 파악할 수 있게하기 위함입니다. 

 

  Cls token은 위 figure에서 MLP Head라고 써져있는 위치로 연산 결과가 나오게 됩니다. 연산된 cls token가지고 MLP Head에 넣어서 class classification을 하게 됩니다. ViT의 backbone은 transformer encoder로 구성이 되어있는데 이는 transformer encoder가 들어오는 여러 token들을 가지고 attention 연산을 수행할 수 있기 때문입니다. 여기서 사용된 transformer encoder block은 보이시는 것처럼 MHA(multi-head attention)과 MLP가 순서대로 위치해있고, 각 모듈 앞에 residual connection과 Layer norm이 들어가 있습니다. Transformer block 내부의 구조는 항상 동일하진 않고 사용을 하는 사람이 원하는 대로 조금씩 바뀔 수가 있습니다. 이 부분은 어느정도 튜닝의 영역인 것 같습니다. 

  

 

3.2. Fine-tuning and Higher Resolution

  새로운 데이터셋에 튜닝을 해야하는 경우에는 backbone의 파라미터는 그대로 두고 0으로 초기화된 MLP head를 새로 붙여줍니다. 새롭게 학습해야할 이미지의 크기가 더 큰 경우에는 patch size를 통일하게 유지하고 token 갯수가 늘어난 채로 넣어주면 됩니다. 이는 transformer가 token length가 가변적이어도 연산이 가능하기 때문에 그렇습니다. 하지만 이렇게 되면 기존의 학습된 positional embedding이 무의미해지게 되는데 이를 해결하기 위해 positional embedding을 2D interpolation 해준다고 합니다. 이 부분은 정확이 저도 positional embedding을 2D interpolation할 일이 없어서 확실히 말씀드리긴 어렵습니다 하지만 예를 들자면 position상 왼쪽 아래 구석이었던 positional embedding이 이미지 사이즈가 바뀌면 해당 position이 더 이상 왼쪽 아래 구석 위치가 아니게 되어버리는 문제라고 생각하시면 될 것 같습니다.

 

 

4. Experiments

 

4.1. Setup

ViT Model Variant

위에 따라 Base, Large, Huge type을 만드실 수 있습니다.

Batch size는 4096, optimizer는 Adam을 사용했습니다. SGD보다 성능이 살짝 더 좋았다고 합니다. 

 

4.2. Comparison to State-of-the-Art

Comparison

ViT를 JFT-300M과 ImageNet1K로 pretrain하고 각 데이터셋에 맞게 fine-tuning을 진행한 결과입니다. 

 

4.3. Pre-training Data Requirements

  앞에서 Transformer-based model은 pretrain을 할 때 많은 양의 데이터를 사용해야한다고 언급했었는데 이를 정량적으로 실험한 부분입니다. 왼쪽 figure를 보면 데이터셋의 크기가 커짐에 따라 성능의 상당히 증가하는 것을 볼 수 있습니다. 오른쪽 figure에서도 ResNet보다 성능 향상 폭이 큰 것을 볼 수 있습니다. 

 

5. Conclusion

나머지 부분이 궁금하시다면 페이퍼 뒤쪽을 참고하시면 좋을 것 같습니다!