[논문 리뷰] An Image is Worth 16x16 Words: Transformers for image Recognition at scale(VIT)

최근들어 필자가 가장 많은 관심을 가지고 있는 vision transformer가 처음 등장하는 논문이다. 물론 이 논문 이전에도 transformer를 vision 분야에 적용한 시도들이 있었지만, 실제로 vision분야에서 transformer가 널리 알려지게 된 것은 이 논문부터라고 해도 과언이 아니다. 한번 vision transformer를 뿌셔보자!!

---

1. Introduction

Transformer in NLP

Transformers는 NLP에서 먼저 언급된 개념이고 지금까지 dominant하게 쓰이는 approach이다. 기존에 NLP에서는 Sequence model(ex. RNN, LSTM)이 많이 사용되었는데 그러한 모델들은 순차적으로 token들을 넣어주어야 하기 때문에 gpu처리가 비효율적인 단점이 있었다. 따라서 sequence model을 전혀 사용하지 않고 오직 attention만을 사용하여 architecture를 만들었고, 이를 "Transformer"한다. Transformer는 gpu에 최적화되었을 뿐만 아니라 성능도 이전을 뛰어넘으면서 NLP에 큰 파장을 일으켰다.

 

https://aistudy9314.tistory.com/63

[논문리뷰] Attention is All you need

 

Transformer to Vision

하지만 VIsion 분야에서는 VIT가 등장하기 전까지 convolutional architecture가 대부분을 차지하고 있었는데, 몇몇 연구자들은 위와 같이 NLP에서 큰 성공을 거둔 Transformer 구조를 vision에도 적용하고자 여러 시도들을 하였다. 예를 들어, CNN architecture와 self-attention을 결합한다거나 아예 self-attention로 대체하는 방법들이 있다. 후자의 경우, 이론적으로는 efficient하지만 specialized attention patterns을 사용하여 hardware accelerators에 잘 맞지 않는 문제가 있었다고 한다(필자도 해당 논문을 읽지 않았기 때문에 자세히는 모르지만 일반적인 self-attention과 다른 방식으로 사용한 듯 하다).

 

논문 저자들을 NLP에서 사용된 Transformer가 위와 같은 문제를 일으키지 않고 성공적인 scaling을 이루어냈기 때문에, 최대한 standard transformer를 그대로 이미지에 적용하려고 노력을 하였다. 

 

2. Method

 

이제 Vision Transformer(VIT)가 어떤 식으로 동작하는지 알아보자!

위에서도 말했듯이, standard transformer를 사용하기 위해서는 input을 1차원으로 만들 필요가 있다. 하지만 NLP에서 사용되는 문장 또는 단어와 비교해서 이미지는 매우 큰 차원으로 이루어져있다. 따라서 그대로 self-attention을 적용하게 되면 메모리적으로나 computation적으로나 너무 비싸다는 문제가 발생한다. 

 

저자들은 이러한 문제를 해결하기 위해 이미지를 patch 단위로 나누어서 이를 token처럼 사용한다. 48x48x3의 크기를 가지는 이미지 $x$를 생각해보자. 논문에서 default로 사용하는 patch size 16으로 이를 나누면, 9x(16x16x3), 즉 9개의 16x16x3 patch를 input으로 하고, 이를 flatten 한 후에 D차원으로 embedding하는 과정을 거친다(9xD).

 

여기에 learnable embedding vector인 class token을 만들고, 이를 다른 token들에 concat해준다(10xD). Transformer모델에 특성상 transformer layer를 거친 output이 input과 동일한 shape를 가지는데 이렇게 나온 Transformer Encoder의 최종 output에서 class token에 해당하는 vector만을 classification header의 input으로 사용한다. 

**추후에 나오는 논문들에서는 class token만 쓰는 것이 아니라 모든 token의 average vector를 사용하기도 한다.

 

class token까지 concat한 vectors(10xD)에 추가로 positional information을 주기 위해 "positional embedding"을 더해준다. positional encoding도 NLP transformer논문에 등장한 개념으로 Sequence 모델을 사용하지 않기 때문에 positional information을 따로 학습할 수 없어 이를 별도로 추가해 주기 위한 방법이다. 이러한 방법으로는sin과 cos같은 주기함수를 사용하는 법과 이를 learnable하게 만들어 학습시키는 방법, 또 최근들어서 token의 위치에 따라 relative하게 만드는 법등이 있다. vit에서는 이 중에서 "learnable embedding"을 사용하여 positional information을 더해준다. 

**논문에서는 1D와 2D-aware position embedding 모두 시도하였지만, 별다른 성능의 향상이 없어서 그냥 1D positional embedding을 사용하였다고 한다.

**별도의 이야기로 positional encoding을 더해주는것이 좋은가 concat해주는게 좋은가와 같은 논쟁도 있다. 결론적으로 성능은 concat이 더 좋을 수 있지만 이는 추가적인 memory와 computational cost가 들기 때문에 trade-off를 가지고 있다고 생각된다.

 

이제 이 vector를 NLP에서 사용된 transformer 구조에 그대로 넣어준다. 이는 마찬가지로 Multi-head self-attention(MSA)와 MLP block, 그리고 Layer Norm형태를 가지고, 이러한 encoder block을 L개 통과하여 output을 내보낸다. 위에서 말했듯이 class token에 해당하는 vector를 사용해 하나의 MLP($D\rightarrow n\_class$)를 거쳐서 최종 결과를 낸다.

 

Inductive bias

위와 같이 self-attention으로만 이루어진 vision transformer는 CNN architecture에 비해서 더 작은 image-specific inductive bias를 가진다. 대표적인 inductive bias로는 첫번째로 어떠한 픽셀은 가까이 있는 neighborhood 픽셀들에 영향을 많이 받고 거리가 멀어질 수록 그 영향도가 떨어진다는 "locality"가 있고, 두번째로는 object가 xy축으로 이동하거나 회전을 주더라도 같은 object로 인식하는 "translation invariant"가 있다. 

 

inductive bias가 적다는 것은 그만큼 contraint 없이 이미지 전체에서 정보를 얻을 수 있다는 장점이 있지만 optimal parameter를 찾기 위한 space 또한 커져버리기 때문에 데이터가 충분하게 많지 않으면 학습이 잘 안되는 문제가 발생한다. 때문에 vit는large datasets(ex. 14M-300M images)로 vit를 pretrained 시키고 이를 사용하여 specific task with fewer datapoints에 transfer learning을 한다.

 

3. Experiments

이 파트에서는 여러 데이터셋과 모델을 비교하면서 실험을 한다. 여러 실험이 많지만 결과만 간략하게만 보도록 하겠다.

 

Comparision to state-of-art

 

최근 state-of-art 모델인 BIT와 Noisy student 모델과 성능 비교를 한 표이다. Vit-Huge 모델을 썼을 때 대부분의 task에서 가장 좋은 성능을 얻었고, 학습 또한 더 빠른 것을 볼 수 있다.

**TPUv3-core-days: 하루 기준 사용된 TPU core 개수

 

Pretraining data requirements

Vit를 pretraining할 때 사용하는 데이터셋의 크기에 따라서 결과가 어떻게 달라지는지에 대한 실험이다.

 

 

작은 데이터셋(ex. imagenet)의 경우, VIT가 BIT보다 under-perform하고 데이터셋의 크기가 커지면서 BIT보다 성능이 높아지는 것을 볼 수 있다. 오른쪽 이미지의 경우, JFT 데이터셋에서 samples의 개수를 달리하면서 성능을 재었는데 비슷하게 더 많은 samples를 사용할 수록 VIT의 성능이 올라간다. 

 

Scaling study

모델들의 scale을 맞춘 후에 성능을 비교하는 실험이다. 논문에서는 FLOP을 scale의 지표로 사용한다.

 

 

여기서 알 수 있는 것은 VIT가 ResNet(BIT)와 비교해서 performance/compute trade off가 좋다는 점과 Hybrid가 작은 computation scale을 가질 때 VIT의 성능을 outperform 한다는 점이다.

**Hybrid는 CNN의 feature를 transformer block의 input으로 사용한 모델이다.