ViT: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scal

논문: https://arxiv.org/abs/2010.11929v2

 

NLP에서 널리 쓰이는 transformer를 vision task에 direct하게 적용하여 좋은 성과를 얻은 논문

0. Abstraction

 transformer 구조는 NLP 분야에서는 de-facto standard이지만, CV 분야에서는 convolution과 함께 사용되거나 convolution을 대체해서 사용하는 등 convolution 구조를 유지한 채 함께 사용되는 정도로만 활용되었었다. 하지만 본 논문에서는 convolution을 완전히 배제하고 transformer 구조에 바로 image를 집어넣는 방식을 보인다.

 특히, 대용량의 dataset으로 pre-train한 뒤 상대적으로 작은 dataset(ImageNet, CIFAR, VTAB 등)으로 transfer(fine-tunning) 했을 때 적은 계산량을 요구함과 동시에 성능이 좋았다.

 

1. Introduction

 Transformer의 self-attention 구조를 computer vision 분야에 접목하려는 시도는 많이 있었다. CNN과 같이 사용하는 것, CNN을 완전히 대체하는 것 등이 있었으나 변형된 attention pattern으로 인해 이론적으론 좋을지라도 scale up하기에는 h/w적인 문제가 있었다. 따라서 지금까지 SOTA는 여전히 ResNet base의 구조였다.

 

 본 논문에서는 transformer의 구조를 최대한 변경하지 않고 그대로 image에 적용하려는 실험을 했다. image를 작은 정사각 patch로 자르고 patch를 linear projection을 통해 patch embedding vector로 변경하고 이 emebedding을 하나의 token으로 하여 patch(token) sequence를 만들어 input으로 사용했다.

 

Normalizaion의 순서가 attention, mlp 앞으로 온 것만 차이가 나는데 이는 transformer 후속 논문에서 앞에 쓰는 것이 transformer를 깊게 만들 때 더 효과적이라는 것이 밝혀졌기 때문이다.

 

 ImageNet과 같은 비교적 적은 양의 dataset으로 학습할 경우 ResNet 보다 성능이 조금 낮게 나왔는데 이는 CNN이 기본적으로 갖고 있는 "inductive bias (translation equivariance, locality)" 특성이 transformer에게는 없기 때문이다. 하지만, large dataset(14M~300M)으로 학습하면 inductive bias이 없음에도 불구하고 좋은 성능을 내는 것을 확인할 수 있었다.

 

Inductive bias

 CNN은 image processing 분야에서 image가 갖는 특성을 잘 뽑아낼 수 있도록 설계되어 있다(biased). 일종의 image processing에 대한 가정이 적용되어 있다는 것인데 그 두가지가 locality와 translation equivalance이다.

 

Locality

 Image의 feature는 특정 부분에 모여있다. 예를 들어 사람의 코라는 object는 (대충) image의 중앙 어딘가의 pixel들로 표현이 될 것이다. CNN 구조는 2차원에서 인접한 특정 부분에 대해 filter가 적용되기 때문에 자연스레 지역적인 feature를 뽑아내는 구조로 동작한다.

 

Translation equivariance

 Translation이란 image의 모든 pixel을 같은 방향으로 같은 거리만큼 이동하는 것을 말한다(즉 image의 slide).

translation equivariance는 input의 translation이 output에도 그대로 영향을 미친다는 의미이다. translation 후 처리를 하나, 처리를 하고 translation을 하나 결과가 같다고 할 수 있다.

 

2. Related Works

 Pixel level로 나이브하게 image self-attention을 하게되면 qudratic cost가 들면서 scale을 키우기가 힘들다. 이에 따라 다양한 approximation 방법들(SA를 local neighborhood에만, global로 하지만 sparse하게, block size를 다양하게 적용 등)이 실행 됐고, 좋은 성능을 보였지만 여전히 복잡한 engineering을 필요로해 효율적으로 구현할 수 없었다.

 특히 https://arxiv.org/abs/1911.03584 논문이 ViT와 굉장히 비슷한 구조인데, ViT는 large scale dataset을 학습했을 때 inductive bias 없이도 좋은 성능을 낸다는 점을 보였다는 것과, 더 큰 patch size로 medium-resolution image에도 적용했다는 점이 더 좋다고 말한다.

 CNN을 large dataset에서 학습/평가하고 transfer 하는 연구들 처럼 transformer로도 했다.

 

3. Method

 

3.1 Vision Transformer (ViT)

 일반적인 transformer 처럼 token embedding의 1D sequence로 input을 받기 위해 (Height x Width x Channel) dimension인 image를 변경해야한다. 먼저 patch size로 image를 쪼개 N(H/P * W/P)개의 patch를 만들어 N x (P x P x C) 차원으로 변경한다. (P x P x C)의 각 patch를 D size의 latent vector로 linear projection 한다. 이를 patch embedding이라고 부름.

 그리고 learnable parameter인 [class] token x_class를 첫번째 token으로 넣고, 이 token이 encoder를 거친 후에 오직 이 token만이 MLP layer에 들어가서 최종 output y를 뽑게 된다. 

 그 다음 learnable parameter인 position embedding이 있다. patch embedding에 element wise로 더해지며 각 token(position) 별로 있다(E_pos (N+1 x D)). 각 embedding을 N개 1D로 나열한 형태로 학습하였다. D/2 size로 N개의 X embedding, Y embedding을 각각 학습해 조합해서 사용하는 2D positional embedding도 사용해봤으나 큰 차이가 없었다.

z_0의 0은 0번 encoder block의 input이라는 의미

 위 그림에서 i, j 번째 그림은 i, j 번째 positional embedding vector와, 모든 다른 positional embedding vector와의 cosin similarity를 나타낸다. 3, 3번째 그림을 보면 3,3 근처에서 제일 밝고, 같은 행, 열의 vector와 비슷한 유형을 보여준다. 즉, 학습된 positional embedding은 가깝거나 행, 열이 같은 positional embedding과 유사한 값을 갖는다.

 

 Transformer는 일반적인 transformer오 마찬가지로 하나의 block은 Multi head self attention layer와 MLP layer로 구성되어 있다.

 

Inductive bias

 앞서 얘기한대로 ViT는 image에 대한 inductive bias가 CNN에 비해 적다. MSA layer는 token 위치에 상관 없이 연산되기 때문에 local, translation equivariant한 특성이 없으며, MLP layer가 1D의 local, translation equivariant한 특성이 있을 뿐이다. 2D neighborhood의 정보를 갖는 부분은 오로지 아래 두가지경우 뿐이다. 

1. Image를 patch로 자르는 것
2. Fine tunning할 때 다른 resolution을 위해 positional embedding을 조절하는 것

 뿐만 아니라 positional embedding의 초기값을 설정할 때도 position에 대한 어떠한 정보도 반영되지 않기 때문에 처음부터 position에 대한 관계를 학습해야한다.

 

Hybrid Architecture

 CNN은 inductive bias로 (학습이 많이 안 됐을 경우) transformer에 비해 feature를 뽑는 능력이 좋기때문에, raw image를 patch로 자르기보다 CNN에서 뽑은 feature map을 cutting하여 patch로 만드는 hybrid 방법도 사용할 수 있다. 뒤에 나오겠지만 결론적으로 모델이 작을 때는 효과가 있었으나 model이 커짐에 따라 그 효과는 사라졌다.  

 

3.2 Fine-tunning and Higher resolution

 일반적으로 많은 양의 dataset으로 pre-train을 한뒤 적은 양의 dataset으로 fine-tuniing을 하게 된다. fine-tunning을 할 때는 pre-train에서 학습된 MLP head는 버리고 0으로 초기화 된 1개의 layer(D x K)만 넣어서 학습하는 게 고해상도 이미지에서 효과적이었다.

 

  ViT는 sequence의 길이에 상관없이 처리할 수 있다. 하지만 positional embedding은 특정한 길이의 sequence에 대해 학습된다. 고해상도 image의 경우 patch size가 동일해 sequence가 길어지게 되는데, 이때 짧은 sequence로 학습된 positional embedding을 그대로 사용하면 position의 정보가 별로 의미가 없게 된다. 따라서 학습된 sequence의 2D position 정보에 기반해 늘어난 sequence에 맞게 2D interpolation을 해서 positional embedding이 적절히 반영되게 한다.

 

4. Experiments

 

4.1 Setup

Datasets

ImageNet(1.3M), ImageNet-21k(14M), JFT(303M) 세 가지 dataset으로 pre-train했으며 ImageNet, ImageNet ReaL, CIFA-10, CIFA-100, Oxford-IIIT Pets, Oxford Flowers-102, VTAB  dataset들에 대해 fine-tunning해서 성능을 비교했다.

 

Model Variants

 ViT는 model size별로 base, large, huge 세 가지로 나눴고 patch size 별로 16x16, 32x32로도 나눴다.

CNN 모델로는 ResNet에 group normalization, standardized convolution을 적용한 ResNet(BiT)을 사용했다.

Hybrid 모델로는 ResNet50의 4번 stage output을 input으로 사용한 base 모델과, 4번 stage에 3번 stage의 layer랑 같은 수를 넣고 얻은 output (4배 더 긴 sequence)을 input으로 사용한 large 모델 두 가지를 사용했다.

 

Metrics

 Downstream dataset에 대한 평가를 했으며 대부분 fine tunning한 뒤 그 모델로 측정했고, 일부 fine-tunning이 너무 cost가 높은 경우 빠른 평가를 위해 few-shot accuracy를 측정했다.

 

4.2 Comparison to State-Of-The-Art

 JFP로 학습된 ViT-H/14, ViT-L/16,  ImageNet-21k로 학습된 ViT-L/16과 기존 SOTA인 BiT-L(JFP), Noisy Strudent(ImageNet, JFP)와의 비교이다.

 

ViT-L/16-JFT가 Bit-L보다 훨씬 계산량이 적으면서도 모든 task에서 성능이 우수했다. VitT-H는 여전히 BiT-L의 1/4밖에 계산량을 쓰지 않으면서도 더 우수한 성능을 보인다.

 

4.3 Pre-training Data Requirements

 CNN에 비해 적은 inductive bias를 갖고 있음에도 large dataset에서는 좋은 성능을 보였다. dataset의 중요도를 알기 위해 두 가지 실험을 진행했다.

 

서로 다른 dataset을 사용해서 pretrain

•  가장 작은 ImageNet에 pre-train한 경우에는 ViT-L가 각각 patch size에 해당하는 ViT-B보다 성능이 안나왔다. ImageNet-21k에서는 비슷했고, 가장 큰 dataset인 JFP에서는 Large model의 성능이 더 좋아졌다.
• 회색 영역은 다양한 사이즈의 BiT 결과의 영역인데 작은 dataset에서는 ViT에 비해 우수하지만 JFP에서는 ViT-H가 더 뛰어난 것을 확인할 수 있다.

 

동일한 dataset에서 학습에 사용하는 양을 다르게

JFP dataset에서 random sampling을 하여 학습 데이터의 양을 조절해봤다.

• 적은 dataset에서는 ViT-B/32, VitT-L/16은 overfit되어 성능이 좋지 않지만 30M부터는 ResNet50보다 좋아진다. 이는 CNN의 inductive bias가 smaller dataset에서 더 효과가 좋을 것이라는 직관에 대한 근거가 된다.
• 반면 dataset이 커지면 inductive bias가 효과가 있을지라도 data로부터 그 특성을 충분히 찾아낼 수 있다는 것이 보인다.

 

4.4 Scaling Study

 pre-training cost 대비 transfer learning (fine-tunning) 했을 때의 성능이다.

• 같은 computational cost대비 ViT가 ResNet보다 뛰어 났으며 같은 성능을 위해서는 ResNet이 2~4배의 cost가 들었다.
• Hybrid의 경우 small dataset(low computation)에서 ViT보다 뛰어남을 보였으나 dataset이 커짐에 따라 그 차이는 없어졌다. CNN의 filter가 inductive bias로 local feature를 잘 뽑아내기 때문에 dataset의 크기에 상관 없이 도움이 될 것으로 예상할 수 있지만 실제로는 많은 data로부터 transformer의 inductive bias만으로도 feauture를 잘 뽑아낼 수 있음을 나타낸다.
• ViT는 실험한 정도의 scale에서는 saturate하지 않았다. 즉, 더 큰 scale에서는 더 좋은 성능을 낼 가능성이 있다.

4.5 Inspection Vision Transformer

 ViT가 어떻게 image를 처리하는지 보기 위해 internal representation을 보았다.

 

Patch embedding layer

 첫 번째 layer는 단순히 patch를 linear projection하여 patch embedding을 만든다. 아래는 이 embedding projection의 학습된 결과를 보여준다(PxPxC x D projection은 D개의 filter가 있다고 볼 수 있다.).

 CNN의 학습된 low level filter와 유사한 모습을 볼 수 있다. 즉, image low level feature(edge, color 등)를 잘 뽑아낼 수 있는 형태로 학습이 되어 있다는 의미이다. 2D geometrically feature를 뽑은 것이 아님에도 많은 dataset으로 학습하여 2D의 feature를 잘 뽑아낼 수 있게 되었다는게 포인트.

 

Positional embedding layer

 Model은 가까운 거리는 positional embedding의 similatiry가 유사하도록 positional embedding을 학습한다. 아래 그림에서 i, j 에 위치한 사진의 i, j pixel 근처의 pixel이 밝다는 것이 이를 나타낸다. 또한 row, column의 pixel이 밝다는 것은 같은 row, column의 positional embedding 역시 similarity가 높다는 의미이다. 이렇게 학습된 positional embedding이 2D topology를 잘 나타내기 때문에 manual하게 positional embedding을 사용하는 것보다 더 효율 적이라는 의미이다.

 추가로 다양한 방식으로 positional embedding을 학습하도록 abblation을 해봤지만 차이는 미미했다. 이 이유는 patch 단위로 position이 들어가기 때문에 224 image 16 patch size라면 겨우 14x14개의 position 정보만 알면 되기 때문에 인코딩 방식은 별로 안중요하다고 함.

 

Self-attention

 어떤 head가 attention한 patch들의 attention weigh을 기준으로 각 attention한 patch들의 weighted average distance를 구해봤다. 가까운 patch들의 attention weight이 높으면 distance가 짧고 멀리 있는 patch들의 attention을 많이 하면 distance가 길어지는 방식. 이는 CNN의 receptive field size와 유사하다.

 lower layer에서는 local한 정보를 보는 attention head, global한 정보를 보는 attention head가 골고루 있는데 반해 layer가 깊어짐에 따라 점점 global한 정보를 보는 attention head만 남게된다. 최종적인 classification을 위해 점차 image 전체의 semantic 정보를 보게 된다는 것이다.

 또한 재미있는 점은 hybrid model의 경우 lower layer의 아주 distance가 짧은 head들이 관측되지 않는다는 점이다. 이는 CNN의 convolution layer에서 이런 low level(highly local) feature들을 이미 뽑아서 transformer에 전달하기 때문이라고 볼 수 있다.

 

4.6 Self-Suprevision

 NLP에서 large scale self-supervised pre-training을 통해 좋은 효과를 냈듯이, patch를 masking하여 예측하도록 self-supervised learning도 적용해 보았다. 결과는 pre-training 없이 처음부터 학습한 것보다는 좋았지만 논문에서 계속 얘기한 방법인 supervised pre train을 한 성능보다는 좋지 않았다.

 

5. Conclusion

 Transformer 구조를 그래도 image recognition task에 적용했다. computer vision의 inductive bias를 고려해서 self-attention 구조를 차용한 이전 연구들과는 다르게 이 논문에서는 image-specific inductive bias를 architecture에 반영하지 않았다(맨 처음에 patch로 자르는 거 빼고). 그저 patch를 embedding하여 input으로 사용했을 뿐이다. 이런 간단하고 scalable(다른 복잡한 engineering이 필요 없기에)한 방법으로 large scale dataset에서 pre-train하고 다른 dataset으로 fine tunning하는 방법으로 놀라운 효과를 봤다. SOTA를 뛰어넘을 뿐만 아니라 상대적으로 적은 계산량만을 사용했다.