Contents

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

Vision Transformer (ViT)

자연어처리분야에서 Transformer 아키텍쳐가 등장하며 비약적인 발전이 있어왔다. Vision Transformer (ViT)로 알려진 이 논문은 2021년에 공개되어 Transformer를 어떻게 비전분야에 도입할 수 있는지를 보여주었다. 저자는 논문에서 가능한한 원래의 Transformer, 즉 언어모델을 처리하는 Transformer를 많이 변형하지 않고 사용하려고 했다고 한다. 그렇다면 어떻게 언어를 처리하던 Transformer를 이미지를 처리하는데 사용하였는지 알아보도록 하자.

개인적으로 ViT의 등장이 흥미로웠던 부분 중 하나는 시계열 데이터를 처리하던 아키텍쳐인 Transformer를 정적인 이미지의 영역으로 가져왔다는 점이다. 보통은 시간적인 속성이 없는 (non-temporal) 데이터보다 시계열데이터는 시간의 흐름을 반영해주어야 한다는 어려움이 추가되고 이 어려움을 해결하기 위해 RNN, LSTM, GRU에서부터 Transformer와 같은 모델로 발전해왔다. 따라서 Transformer는 태생적으로 시계열을 처리하는 모델이고 이는 모델의 input shape자체가 (Sequence Length, Batch Size, Feature Number) 꼴이라는 점에서도 확인할 수 있다. 제목에서 알 수 있듯 이 논문에서는 이미지를 16x16 개로 구성되는 단어처럼 해석하여 접근하는 흥미로운 관점을 취하고 있다.

Abstract

Vision분야에서는 그동안 CNN에 의존해 이미지의 특징을 처리해왔다. 여기서는 Transformer를 사용해 이미지를 일련의 patch(sequences of image patches)로 보고 이렇게 학습한 모델이 분류 모델에서 매우 잘 작동함을 보여준다. 또한 학습할때도 더 적은 리소스로 훈련이 가능함도 확인하였다고 한다.

Introduction

Transformer(Vaswani et al. 2017)의 등장이후 NLP에서 Transformer의 사용은 사실상 표준으로 여겨지고 대규모 text corpus를 Transformer로 학습시킨 뒤 task-specific dataset에 대해 fine-tune하는 방법이 많이 사용되고 있다. 대표적인게 유명한 BERT(Devlin et al., 2019)모델이다. BERT는 대규모 언어 데이터셋에 대해 masked language model(MLM)이라는 개념을 도입해 문장의 중간 중간 mask를 만들고 이를 맞추도록 하는 self-supervised 방식의 학습을 시켜 모델로 하여금 언어의 이해도를 올려놓은 뒤 풀고 싶은 문제에 대해 fine-tune하는 방식을 사용하였고 이는 실제로 상당히 효과적임을 보였다. 놀랍게도 기존 모델들과는 달리 Transformer는 매우 scalable한 성질을 보이고 모델을 키우면 키울수록 성능이 향상되는 경향을 보이고 있다. 즉 성능에 대한 saturation이 다른 모델에 비해 뚜렷하게 나타나지 않아 충분한 자원과 데이터만 있다면 성능을 계속 향상시킬 수 있다는 가능성을 보여주고 있다. 이와 관련해서는 scaling law의 키워드로 찾아보면 추가적인 정보를 얻을 수 있다.

워낙에 NLP에서 Transformer가 압도적인 성능을 보여주고 GPT모델들이 이전과는 차원이 다른 언어생성모델을 보여주다보니 vision에도 이를 도입하려는 시도는 예상해 볼 수 있는 것이었고 구글리서치팀에서 vision분야에 Transformer사용을 본격적으로 알리는 이 논문을 발표하였다. 이미지를 sequence로 처리하기위해서 저자는 이미지를 patch단위로 나누어 linear embedding을 사용해 Transformer에 입력하게 된다. Transformer는 BERT와 비슷하게 Transformer encoder는 이어 붙여서 사용한다. 이미지 patch 각각은 NLP에서 단어에 대응된다.

이렇게 학습하는 방식은 mid-sized dataset규모에서 성능이 낮은건 아니지만 ResNets와 같은 기존 vision모델보다 낮은 성능을 보여준다. 저자는 이에 대한 원인으로 inductive bias를 언급한다. Inductive bias란 학습하는 모델 (learner)에 적용되는 개념으로 모델이 갖는 추가적인 가정을 의미한다. 예를 들어 vision에서 CNN구조를 사용하는 것은 이미 그 자체로 “이미지를 이해하는데는 지역적인 정보가 중요하다"는 가정을 하는 것과 같다. CNN의 작동방식을 보면 receptive fields가 target pixel의 주변부로 형성되기 때문이다. 이런 inductive bias가 강할수록 모델은 가정을 만족할 때 더 적은 데이터에서 효율적으로 성능을 높일 수 있다. 이미 훈련하기 전부터 지역적인 정보(locality)가 중요하다는 것을 알고 학습하는 것과 같기 때문이다. 하지만 Transformer의 경우는 이런 inductive bias가 낮다. 모델구조를 생각해보면 모든 입력에 대해 attention을 사용할 뿐이지 이 attention에 대해 지역적인 정보가 중요하다는 내용을 알고 학습하지 않는다. 따라서 CNN과는 다르게 Transformer는 순전히 학습만으로 이미지에서 지역적인 중요성을 학습해야하며 더 많은 데이터를 필요로한다는 것도 이에 기인하게 된다. 즉 inductive bias가 낮은 Transformer이기에 충분한 데이터를 통해 학습할 수 있다면 일반화 측면에서는 이점이 있지만, 데이터가 충분하지 않다면 동일 조건의 CNN보다 성능이 낮을 수 있다는 것이다. 개인적으로는 내재적으로 sequence가 아닌 이미지를 patch를 나누었기 때문에 Transformer입장에서는 오히려 inductive bias측면에서 더 불리하게 학습을 시작한다는 생각도 든다.

저자가 수행한 실험에서도 이러한 inductive bias관점에서의 분석을 뒷받침한다. ViT는 데이터셋이 커지게되면 기존 vision모델들을 이기면서 더 좋은 성능을 보여준다.

Method

저자는 논문에서 여러차례 original Transformer를 최대한 따라가려고 했다는 것을 언급한다. 실제로, 모델을 간단하게 설명하면 이미지를 patch로 나누고 이를 token처럼 처리하기위한 linear projection을 해준 뒤 이후에는 기존 Transformer encoder와 완전히 동일한 구조로 입력해준다. 그리고 Transformer encoder의 출력 값 중 첫번째 token의 값에 MLP Head를 붙여서 classification과 같은 task를 수행하면 된다는 것이 ViT의 핵심이다.

Vision Transformer

/vit/figure1.png
Figure 1

이 논문에서 제일 중요한 부분은 이미지가 어떻게 encoder까지 처리되어 들어가는지라고 생각한다. 순서대로 살펴보자. 우선 이미지는 다음과 같이 표현된다. $\mathbf{x} \in \mathbb{R}^{H \times W \times C}$ 예를들어 224 x 224이미지라면 각각 가로, 세로 해상도로 $H,W = 224$가 될 것이고 RGB채널이므로 $C=3$이 될 것이다. 이를 patch로 나누기 위해서 우선 patch의 크기를 결정해야 한다. Patch의 크기가 16 x 16이라고 생각해보자. 그러면 기존 224 pixel을 16 pixel을 갖는 patch로 만들어야 하므로 가로, 세로가 각각 $224/16 = 14$가 된다. 다시 말해, 이렇게 한다면 원본 이미지는 가로와 세로 각각 14등분되어 196개의 patch로 구성된다. 식으로 나타내면, $N = \frac{HW}{P^2}$이 된다. 여기서 $N$이 sequence길이로서 Transformer에 입력된다. Feature dimension은 $D$로 표기하며 이는 linear projection단계에서 정하면 되는 hyperparameter이다. 논문에서는 projection된 결과를 patch embedding이라고 한다. Projection자체는 MLP레이어를 사용하는 것으로 lookup table방식의 Embedding과 헷갈리지 않도록 유의해야 한다.

Transformer에 들어가기 전에 첫 번째 token은 BERT에서 사용한 [class] token을 사용하게 된다. 이 토큰에 대응되는 마지막 출력값이 이미지에 대한 summarization처럼 여겨지고 Transformer output에서 첫 번째 token에 MLP를 붙여서 원하는 task를 수행한다. 기존 Transformer처럼 순서정보를 알려주는 positional embedding을 사용하며 여기서는 learnable positional embedding을 사용하였다고 한다. (다른 방식에 대한 실험도 있으나 결괒거으로 성능차이가 없었다고 한다.)

Transformer encoder 구조는 original Transformer encoder와 완전하게 동일하다. Inductive bias측면에서 CNN과는 다르게 주변부의 정보를 활용하는 inductive bias가 주어지지 않으므로 ViT는 image-specific inductive bias가 훨씬 낮은 구조이다. 따라서 ViT는 이미지의 2차원 공간적인 구조를 모두 학습단계에서 학습해야한다.

논문에서는 hybrid architecture도 제시하는데 여기서 hybrid는 이미지를 patch로 나누어서 바로 입력하는 것이 아니라 CNN을 사용해 feature map을 만들고 이를 활용하는 구조를 의미한다. 즉, CNN이 갖는 inductive bias와 Transformer를 동시에 활용하는 구조인 것이다.

Fine-Tuning and Higher Resolution

언어모델처럼 ViT도 큰 데이터셋에 pre-training을 하고 downstream task에 fine-tune하는 것을 염두에 두고 만들어졌다. 따라서 downstream task를 수행할 때는 prediction head를 떼고 여기에 $D \times K$의 feedforward layer를 붙여서 사용한다. 예를 들어 downstream task는 10개 class에 대한 분류문제를 수행하는 것이라면 기존 head에서는 훨씬 큰 label에 대한 분류작업을 수행하는 MLP layer가 있었을 것이므로 이를 버리고 $D \times K$의 MLP를 넣고 fine-tune하는 방식을 사용하는 것이다. 보통 fine-tune단계에서 pre-training때보다 고해상도 이미지를 사용하는 것이 유리하다고 한다. (Touvron et al., 2019, Kolesnikov et al., 2020) 고해상도 이미지를 사용할 때는 patch size는 유지하였고 따라서 sequence 길이는 길어지게 된다. 다만 이렇게 할 경우 기존에 학습한 position embedding을 활용하지 못한다는 문제가 생기게 된다. 이를 해결하기 위해 이미 학습된 position embedding에 대해 2D interpolation을 사용하였다고 한다. 짐작컨데 학습 때 패치를 14 x 14로 사용하였고 downstream에서 18 x 18개의 patch size가 되었다면 학습하지 못한 15 ~ 18에 해당하는 patch는 이전 14개를 사용해 interpolation한 것으로 생각한다. 저자는 ViT에서 2차원 정보에 대한 inductive bias가 주어지는 부분은 resolution adjustment와 patch extraction뿐임을 강조한다.

Experiments

실험섹션에서 눈여겨 볼 점은 ViT로 SOTA를 달성했다는 것과 ViT는 학습데이터가 충분히 제공될 때 SOTA를 달성할 수 있으며 이는 앞서 언급한 inductive bias에 대한 가정을 뒷받침한다는 것이다.

/vit/table1.png
Table 1

사용된 ViT는 BERT와 비슷하게 파라미터수에 따라 위와 같이 나누어진다. 실험에서 모델을 지칭할 때 ViT-L/16과 같이 표기되는데 이 때 L은 Table 1에서의 Large모델을 사용하였다는 것이며 /16은 하나의 patch가 가지는 해상도가 16 x 16임을 의미한다.

/vit/table2.png
Table 2

Table 2에서 확인할 수 있듯 benchmark에서 기존 성능을 상회하고 있음을 볼 수 있다. 개인적으로는 이미 이미지 분류에 있어서는 워낙에 높은 수준의 benchmark가 형성되어있기 때문에 새로운 architecture에서 대등한 성능을 달성한 것 자체가 충분히 의미가 있다고 생각한다. 게다가 저자의 주장대로 inductive bias가 낮은 구조이기 때문에 scalability측면에서도 더 가능성이 크다고 생각한다.

앞서 언급한 hybrid architecture는 CNN에 의한 local feature가 ViT학습에 상당히 도움이 될 것이라는 예상과는 다르게 실제 실험에서는 큰 차이가 없었으며 model이 커질수록 이런 경향은 강하게 나타났다고 한다. 세부적인 내용은 본문과 appendix의 자료를 참고하면 된다.

이 논문에서는 추가적으로 BERT와 같은 masked patch prediction을 사용한 self-supervision 학습도 수행하였다. Patch를 가리고 이를 맞추도록 pre-train하는 방식이나 masked된 patch의 원래 pixel을 맞추는 것이 아닌 3-bit mean color를 맞추는 방식으로 한다는게 특이한 점이다. 이와 관련된 부분은 Appendix B.1.2에서 설명하고 있다. 참고로 BERT방식으로 ViT를 학습하는 논의는 이후 나오는 논문인 BEIT: BERT Pre-Training of Image Transformers에서 본격적으로 다루어진다.

Conclusion

이 논문은 그동안 vision분야에서 다방면으로 다루어진 CNN기반의 접근이 아닌 NLP에서 활용되던 Transformer를 상요해 vision task에서 SOTA를 달성하였다는데 의의가 있다. 특히 이미지를 patch로 나누어 접근하는 것은 이후 연구들에도 차용되어 활용되는 구조로 vision분야에서 Transformer를 사용하는 하나의 기준을 마련하였다는 점에서 크게 기여하였다고 생각한다.

Bibliography

Dosovitskiy, Alexey, et al. “An image is worth 16x16 words: Transformers for image recognition at scale.” arXiv preprint arXiv:2010.11929 (2020).