[5주차 / 천원준 / 논문리뷰] AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

ViT (Vision transformer) 논문은 NLP 분야에서 사용되던 Transformer 구조를 Vision task에 적용한 논문입니다.

먼저 Transformer 논문을 읽고 보시는 걸 추천드립니다.

 

Abstract

NLP task에서는 Transformer가 가장 스탠더드 한 모델로 많이 사용되었지만 computer vision task에서는 주로 convolution network와 함께 사용하거나 convolutional network의 일정 부분들을 대체하는 방식으로 사용되었습니다. 하지만 transformer가 image patches에 직접 적용되어도 image classification을 매우 잘할 수 있다는 것을 보여주었습니다.

 

많은 데이터를 pre train 하고 중간 사이즈와 작은 사이즈의 image recognition benchmark들에 전이 학습을 진행할 때 ViT는 SOTA convolutional network보다 계산량은 적으면서 성능은 좋았습니다. 

 

Method

가장 기본적인 Transformer 모델과 최대한 유사하게 구성하였다.

 

Model Overview

 

위 사진은 기본적인 모델의 오버뷰이다. 우선 이미지를 패치들로 나누고 linear projection을 진행한다.

각 패치마다 position embedding을 하고 이를 인풋으로 transformer encoder에 넣는다.

이는 기존 transformer가 1차원 토큰 임베딩 시퀀스를 인풋으로 받기 때문에 2차원 이미지를 비슷한 형태로 만들어준 것이다. 

 

이를 수식으로 나타내면 다음과 같다. 

• x∈R^(H×W×C)를 xp∈R^{N×(P2⋅C)}로 변환
• (H, W): 원본 이미지의 해상도
• C: 채널의 수
• (P, P): 각 이미지 패치의 해상도
• N = HW/P^2 : 패치의 수
• D: 모든 레이어에 사용하는 constant latent vector size / D 차원으로 mapping 한다.

BERT의 [class] token과 비슷하게 learnable embedding을 sequence of embedded patches (z00 = Xclass)에 추가하는데 transformer encoder (z0L)의 아웃풋 state가 image representation y이다.

Fine tuning과 pre training때 classification head는 z0L에 부착된다. Classification head는 하나의 hidden layer를 가진 MLP로 구현되고 fine tuning시 single linear layer로 구현된다.

위치 정보를 위해 patch embeddings에 1D position embeddings가 추가된다. 

 

여기서 나온 최종 embedding sequence가 Encoder의 인풋으로 들어간다. 

Transformer encoder는 multi headed self attention layer들과 MLP blocks로 구성되어 있다. Layer normalization이 모든 block 전에 적용되고 residual connection이 모든 블록 후에 사용된다.

MLP는 GELU non linearity를 사용하는 2개의 layer로 구성된다.

 

ViT 수식

 

ViT는 다른 CNN 구조들 보다 공간에 대한 inductive bias가 적다. 

Image patch 대신 CNN의 feature map을 사용해 input sequence를 만들 수 있다. Feature map의 spatial dimension을 flatten해서 spatial size를 1x1로 만들어 사용한다. 

 

일반적으로 pre training을 large dataset으로 진행하고 fine tuning을 작은 downstream tasks로 진행한다.

이 과정에서 pre trained prediction head는 제거되고 0으로 초기화된 D x K의 feedforwad layer가 부착된다. 이때 K는 downstream classes의 수이다.

Pre training때 보다 높은 해상도로 fine tuning을 진행하는게 효과적이다. 고해상도 이미지를 사용할 때 patch size는 동일하게 가져가  sequence length가 증가한다. Fine tuning을 진행할 때 pre trained position embedding은 의미가 사라진다. 따라서 원본 이미지에서의 위치를 기준으로 pre trained position embedding의 2D interpolation을 진행한다. 

 

Conclusion

• Transformer만 사용했지만 성능이 SOTA CNN 모델보다 좋다
• 상대적으로 Pre train 비용이 적다
• Pre training data가 많아야 성능이 좋아진다

References

https://arxiv.org/abs/2010.11929

https://kmhana.tistory.com/27

https://simonezz.tistory.com/102