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

Attention Is All You Need 논문에서 제안된 Transformer은 이후 NLP 분야에서 매우 폭넓게 사용됐습니다(현재까지도). inductive bias가 적다는 특징이 있으며, 데이터셋이 매우(구글의 300M과 같이...) 큰 환경에서 pretrain, 혹은 pretext task 학습을 진행후 downstream task 등에서 활용하기에 용이합니다. BERT, GPT 등의 Transformer 기반 NLP 모델들이 많이 연구되었으니, 찾아보시길 권합니다.

 

오늘 설명할 논문은 구글의 An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale 이라는 논문입니다. 구글 리서치, 브레인 팀 소속의 연구자들이 집필한 논문으로 2020년에 발표됐습니다. 기존의 NLP task에 적용된 transformer에서 벗어난, computer vision 분야에 transformer가 활용된 사례로서, Vision Transformer(ViT)로 불립니다.

---

1. Introduction

 

이전에도 언급했듯이, transformer는 매우 우수한 성능을 보여주며 RNN 계열의 모델보다 빠르게 학습될 수 있지만, 학습에 큰 크기의 데이터셋이 필요하다는 특징이 있습니다. 실제로 많은 NLP transformer는 매우 큰 text corpus에서 pretrain된 이후, 보다 작은 task에 적용되는 방식으로 사용됩니다.

 

Computer Vision 분야에서는 명실상부한 ResNet 기반의, skip-connection 및 convolution 기반의 모델들이 강세를 보입니다. CNN을 기반으로 하지 않는 SOTA 모델을 2012년 이후 존재하지 않았습니다... 저자들은 NLP 분야에서 transformer의 성공을 보고 Vision 분야에 적용하기로 합니다(Attention Is All You Need 논문 말미에도 사실 vision 및 다른 분야에서 적용가능할 것이라는 언급이 존재합니다). 하나의 걸림돌은 바로 transformer는 CNN과 비교했을때, 비교적 inductive bias가 적다는 점입니다. CNN는 translation invariance, locality, stationarity, compositionality 등의 inductive bias를 가집니다... 그럼에도, transformer는 매우 큰 데이터셋으로 학습될 때, 매우 좋은 성능을 보였다는 것이 이 논문의 요지입니다.

---

2. Related Work

 

Machine translation에서 활용되는 transformer, 특히 핵심적인 구성요소인 self-attention은 각 단어(token)별로 self-attention을 가능하게 합니다. 이를 naive하게 vision에서 적용하려면, pixel-wise로 self-attention을 구성할 수 있는데 이는 비현실적일뿐더러 계산비용이 과다합니다. Attention을 vision 분야에 적용한 다른 연구들은 근처 local neighborhood에 적용(global하지 않도록)하는 등으로 attention 개념을 특화시켜 사용했습니다. 혹은, CNN을 사용하면서, feature map에 적용하거나 하는 등으로 attention을 응용했었습니다.

 

이 논문에서는 먼저, global하게 full-sized image에 대해 self-attention을 사용하는 transformer를 선택합니다. 물론 raw image를 사용합니다.

---

3. Method

 

분야가 vision임에도 original Transformer의 구조를 아주 closely하게 따라갑니다. image가 주어질 때, 이를 fixed-sized의 작은 패치(patch)로 분할하고 각각을 linear하게 embedding시킵니다. 동시에 NLP 분야에서처럼 위치정보를 포함시키기 위해서 positional embedding을 추가합니다. 이를 transformer encoder에 태우게 됩니다. 이 과정에서 'classification token'이라고 불리는 학습되는 것이 추가됩니다.

Vision Transformer(ViT) Overview

 

좀 더 구체적으로, NLP에서의 input과 다르게 2D input 처리를 위해서 patch embedding을 진행합니다. 이는 기존의 HxWxC 차원을 가진 이미지를 Nx(P^2*C) 의 차원으로 바꾸어줍니다. N은 하나의 이미지가 분리되는 여러 패치의 개수이고, PxP는 각 패치의 resolution입니다. C는 본래 이미지의 channel(RGB)입니다. 이러한 N개의 vector들은(각 P^2*C 길이) 다시 embedding을 거치며 D dimensionality를 가지게 됩니다. positional embedding 차원도 동일하게 D가 됩니다.

 

이후 이러하게 변형된 여러개의 패치와, 한개의 classification head를 붙인 것을 z_0으로 부릅니다. 이를 transformer의 multi-head self-attention을 거치게 하고, position-wise MLP(1D-CNN과도 비슷합니다)를 거치게 합니다. 각 단계 사이에는 layer normalization과 skip-connection을 사용합니다. 최종적인 output, 즉 class prediction은 L번째 encoder block을 거친 것의 첫번째인 classification head를 통해 계산합니다.