[Paper review] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

오늘 리뷰할 논문은 An Image Worth 16x16 Words: Transfomers for Image Recognition at scale이다. ICLR 2021에 등재된 논문으로, 비교적 최근 논문이라고 말할 수 있겠다. 논문제목을 보면 이미지는 16x16 word의 가치를 가진다고 돼있다. 이게 무슨 뜻일까? 16x16은 단순히 숫자이므로 무시하면, image와 word는 서로 같은 것이다라고 해석할 수 있다. 상식적인 선에서는 이해하기 어렵지만, 이는 논문에서 다루고 있는 Vision Transfomer의 핵심을 관통한다. 

 

Vision Transfomer가 나오게 된 배경을 이해하기 위해서는 Transfomer를 알 필요가 있다. 이 글에서는 자세히 설명하지는 않겠지만, 2017년 NLP 분야, 특히 기계번역 분야에서 하나의 모델이 발표되는데 이게 바로 Transfomer이다. RNN과 같은 Recurrent method을 아예 사용하지 않고, 오직 Attention이라는 핵심적인 원리로 작동하는 이 모델은 기존의 줄기차게 사용해온, 모델들의 성능을 넘어서는 sota모델로써 군림하게 되었다. 이후 Transfomer의 구조는 이후에서 더욱 다른 모델, BERT 등에 응용되어 더욱 발전하게 되었다. 한편 ComputerVison분야에서도 이러한 Transfomer의 뛰어난 성능을 할용을 시도해보기위한 연구가 진행됐다. 지금 리뷰하고자하는 논문도 이러한 연구들의 결과물이라고 할 수 있다.

논문에서 제시하는 Vision Transfomer의 핵심은 두가지로 설명할 수 있다.

 

첫번째, 이미지를 단어로 생각한다.

즉 Transfomer가 NLP에서 최고의 성능을 내기때문에, 이미지를 word로 적용할 수 있다면 좋은 성능을 기대할 수 있다는 생각에서 출발한다. 그렇다면 Vision Transfomer는 이러한 idea를 어떤식으로 구현했을까? 모델의 구조를 통해 이해해보자.

논문 Figure 1 참조

이미지는 2D(HxWxC)이기 때문에, 이를 patch로 먼저 분할해 1D(PXPXC, P는 patch 길이)로 만들어 주는 것으로 시작한다. 여기서 분할된 patch의 개수는 곧 Transfomer의 sequence length에 해당하는 값이라는 것을 파악할 수 있다. 이후 이 patch들은 위 그림에서 Linear Projection이라는 모듈을 통과하게 되는데, 이는 Transfomer가 사용하는 dimension으로 patch들의 dimension을 바꿔주는 과정이다. 논문에서는 이 과정을  patch embedding이라 한다. 위 그림에서 쌩뚱맞게 하나의 블록(index 0)을 확인할 수 있는데, 이는 class token의 기능을 하는 learnable embedding이다. 쉽게 말해서, classification을 위해 만들어 장치로써, 특히 Vision Transfomer에서는 classfication head가 이 class token이 Transfomer를 거치고 나온 output에만 attach 되어있다는 점에서 중요한 역할을 수행한다고 할 수 있다. 이후 과정은 Transfomer와 거의 유사하다. Positional Embdding을 거치고 encoder 블록을 거친다. (여기서는 decoder는 사용하지 않음) 그리고 조금전 언급했듯이, class token 대응하는 encoder output이 classification head와 연결돼, classification 작업을 수행하는 것이다. 

두번째, 큰 규모의 데이터를 통해 pretrain한 후, finetuning하는 방식을 사용한다. 

여기서 핵심은 큰 데이터를 사용한다는 점이다. 논문에서는 데이터셋의 크기를 바꿔가면서, pretrain했고 small size나 middle size의 데이터셋을 통해 pretrain는 한 모델을 fine tuning했을때 SOTA를 달성하기는 커녕, 기존의 CNN 기반에 모델에 비해 심각한 과대적합의 문제가 발생했다. 하지만 큰데이터셋, 논문에서는 JFT-300M을 사용했을때는, SOTA를 달성할 수 있었다. 이는 Vision Transfomer가 큰 데이터를 통한 pretrain의 중요하다는 것을 보여준다. 다음은 논문에서 보여준 성능 지표이다. 

특히 여기서 흥미로운 점은 VTAB의 결과인데, VTAB은 기존 SOTA모델들도 좋은 성능을 내지 못하는 low-data로 이루어져있는 데이터셋이다. 이러한 데이터셋에서도 Vision Transfomer는 SOTA를 달성했다. 가히 놀라운 성능이다.

논문의 결론부분에서는 이러한 논문에서 제시한 Vision Transfomer의 한계를 시사하고 있다. 첫번째는 아직 모델을 다른 이미지작업(segmentation, detection 등)에 적용해보지 않았다는 점이다. 이러한 작업들에서도 성능이 잘 나올지는 미지수이다. 두번째는 self-supervised에 관한 것이다. 이는 label 직접 주어지는 supervisor learning과 다르게, label 없이 학습하는 unsupervisor learning 중 하나이다. 이 방법을 통해 NLP 영역에서 Trasfomer가 더욱 비약적으로 발전했다. 실제로 논문에서도, 이러한 Transfomer의 self-supervised pre-training에 영감을 받고, Vision Trasfomer에 적용했다. 하지만 큰 데이터셋을 단순히 pretrain했을 때보다 좋은 성능을 내지 못했다. 이 부분 역시 잘 연구된다면, NLP 분야와 같이 성능면에서 개인적으로 기대가 된다. 

 

이번 논문을 읽으면서, 적잖이 충격을 받지 않을 수가 없었다. 먼저 NLP와 ComputerVision의 경계가 서서히 사라지고 있다는 것을 확실하게 느낄 수 있었다. 예전에 교수님께서 어떤 맥락에서 이러한 말씀을 하셨는지 확실히 이해가 됐다. 그리고 ComputerVision을 주름잡고 있었던 CNN을 사용하지 않고도, SOTA를 달성할 수 있다는 점에서 이미지 프로젝트나 작업을 할때, 항상 CNN 모델만을 생각한 나의 시야가 한층 넓어진 느낌이다. 

 

아래는 리뷰한 논문 주소이다.

https://arxiv.org/abs/2010.11929