[논문 리뷰] AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE (Vision Transformer : ViT)

 

포스팅 개요

 

NLP에서 뛰어난 성능을 보인 transformer 구조를 비전 분야에 적용한 Vision Transformer(ViT)에 대해 리뷰해보려고 한다.

사실 OCR 프로젝트를 하기 위해 TrOCR을 썼었는데 이때 처음 ViT에 대해서 알게 되었다. TrOCR의 인코더로 ViT가 이미지의 텍스트를 인지하는 역할을 했다. Hugging Face에서 pre-train model을 사용하여 fine-tunning을 진행했었는데 정확하게 모델의 구조와 디테일한 부분은 몰랐기 때문에 이 논문을 리뷰하면서 ViT를 구현하는 것까지 해보려고 한다.

 

 

ViT에서 가장 주목할 포인트!

1. ViT는 CNN에 대한 의존이 필요하지 않고, 이미지 패치들에 순수 트랜스포머가 사용되고 이미지 분류에 잘 작동한다.
2. 많은 양의 학습 데이터가 있다면 CNN 구조의 모델보다 성능이 더 뛰어나다.
3. ViT의 한계는 CNN보다 inductive bias가 부족해 중간 사이즈의 데이터셋(예 ImageNet)에서는 input sequnce를 잘 추론하지 못하고 떨어진 성능을 보인다.

 

 

모델 구조

 

 

전체 모델 개요

1. 이미지를 patch (P x P)로 분할
2. Linear Projection of flattened patches
3. Patch에 Position embeddings 추가 (feat. CLS token)
4. Standard Transformer encoder에 input 시퀀스 데이터 넣기
5. Multi Layer Perceptron (MLP) Head - 출력단

 

 

1. 이미지를 Patch (PxP)로 분할

일반적인 Transformer 모델의 경우 1차원의 토큰 임베딩 시퀀스가 들어간다.

하지만 이미지의 경우 2차원(채널까지 3차원)이기 때문에 차원을(1차원으로) 맞춰줘야 한다.

 

원본 이미지의 해상도 (H, W), 이미지 채널 수 (C)

 

패치 해상도 (P), reshape된 패치 수 N= HW/P^2

 

 

보통 패치의 크기는 16 x 16 사이즈이고, 32 x 32로도 만들 수 있다.

이때 작은 사이즈의 패치일수록 시퀀스가 길어질 것이고, 더 높은 해상도를 유지할 수 있다.

 

 

2. Linear Projection of flattened patches

다음의 수식을 통해서 1차원으로 flatten 한 패치들을 D차원으로 매핑하고, Projection 된 output을 patch embedding이라고 한다.

E는 learned embedding matrix이고, Linear Projection에 대한 식은 아래와 같다.

 

Linear Projection

 

이때 D는 transformer의 모든 hidden layer의 사이즈와 같다.

 

 

3. Patch에 Position embeddings 추가 (feat. CLS token)

Transformer 모델의 중요한 특징 중 하나는 시계열 데이터를 순서대로 입력으로 넣어주지 않고, 시퀀스 데이터에 위치 데이터를 추가해 병렬적으로 학습을 시킬 수 있는 장점이 있다.

 

위치 정보를 입력할 때 Positional Encoding처럼 사인, 코사인 함수를 이용하여 단어 벡터들과 더하는 방법이 아닌, 위치 정보를 위한 Embedding layer를 하나 더 사용하여 문장의 길이에 따른 Position Embedding 벡터를 더해준다.

 

 

• 위의 식에서 position embedding은 위치 정보(E_pos)를 담고 있으며 Patch embedding에 추가가 된다.
• 이때 E_pos가 R(NxD)이 아니고 R(N+1 x D)인 이유는 맨 앞의 [Class] token을 concat 하기 때문이다.
• [Class] token은 BERT의 [Class] token과 유사하게 임베딩된 patch의 맨 앞에 하나의 학습 가능한 [Class] Token이다. 이것의 역할은 Transformer의 여러 encoder 층을 거쳐 최종 output으로 나왔을 때, 이미지에 대한 1차원 representation vector로써 사용된다.
  • representation vector로 사용될 수 있는 이유는?
    • BERT에서 Embedding vector는 Token embd + Segment embd + Position embd로 구성되는데 (워드 임베딩 + 첫 번째 문장인지 두번째 문장인지 + 단어의 위치 정보) [Class] token은 모든 input에서 항상 맨 앞에 있기 때문에 Segment embd, Position embd이 동일하다. 
    • [Class] token은 단어로써 의미가 없는 빈칸에 가깝기 때문에 Transformer에서 다른 word token에 대한 정보가 인코딩 된다. 이것의 의미는 전체 Sentence에 대한 정보를 넘긴다는 의미가 된다.
    • 패치 임베딩에서의 모든 정보를 함께 모아 전체 이미지를 나타내는 역할을 하는 의미에서 CLS token은 representation vector로 사용될 수 있다.

•  
• Standard 1-D positional embedding을 사용 (2-D positional embedding과 차이가 크게 없었음)

 

4. Standard Transformer encoder에 input 시퀀스 데이터 넣기

MSA (Multi Head Attention) block - (2)

MLP (Multi Layer Perceptron) block - (3)

LN (Layer Norm) block - (4)

 

encoder의 output에서 시퀀스의 첫번째 요소인 z^0_L을 MLP head에 넣어 class label을 예측한다.

Transformer Encoder 구조

 

 

5. Multi Layer Perceptron (MLP) Head - 출력단

• 일반적인 CNN의 Image Classfier과 동일하다.
• 특별히 봐야 할 부분은 [Class] Token만을 사용한다는 것이다. [Class] Token이 Transformer의 여러 층을 거치고 최종 y가 나왔을 때 이 y는 이미지에 대한 1차원 representation vector로써의 역할을 수행하기 때문에 classification의 input으로 사용된다.

 

Inductive bias

이 논문에서 "Inductive bias"라는 단어가 반복적으로 나오는데 난 이 논문 리뷰를 하면서 처음 알았다. 

이 단어 자체만으로는 뜻이 와닿지 않아서 풀어서 설명하자면 주어지지 않은 입력에 대해 예측하기 위해서 즉, 일반화의 성능을 높이기 위해서 만약의 상황에 대한 추가적인 가정(Additional Assumptions)이라고 보면 된다.

 

학습 시에는 경험하지 않은 input이 주어져도 output의 적절한 귀납적 추론이 가능하도록 하기 위한 일련의 가정의 집합이라고 할 수 있다. 데이터의 특성에 맞게 적절한 Inductive Bias를 가지는 알고리즘을 사용해야 높은 성능을 낼 수 있다.

 

Transformer 아키텍처를 기반으로 두는 ViT 모델을 CNN과 RNN의 특징과 비교하면 아래와 같다.

 

• CNN
  • Locality와 Transitional Invariance 한 특성이 있다.
  • Convolution Filter가 입력을 Window Sliding 하면서 같은 weight를 공유한다. 이것의 의미는 global이 아닌 local feature들과 연산하면서 각 Entities 간의 관계가 서로 가까운 요소들에 주목할 수 있다.
  • Transitional Invariance는 input의 위치가 달라져도 output이 동일한 값을 갖는 것을 말한다. CNN에는 Max pooling이 있기 때문에 k x k filter 사이즈만큼의 값들을 1개의 max 값으로 치환한다. 즉, k x k 범위 내에의 translation에 대해서는 invariance 하다.
• RNN
  • Sequentiality와 Temporal Invariance 한 특성이 있다.
  • Sequential이란 입력이 시계열의 특징을 갖는다고 가정하며, Temporal Invariance는 동일한 순서로 입력이 들어오면 출력도 동일하다는 것을 말한다.
• Transformer와 비교해 보면?
  • Transformer는 Positional Embedding과 Self-Attention을 이용해 모든 정보를 활용한다.
    즉, Transformer는 CNN에 비해 Inductive Bias가 부족하다고 볼 수 있다.
  • Global 한 정보에 대해서는 좋은 결과를 낼 수 있지만 Local feature에 대한 정보는 CNN에 비해 추론 능력이 떨어진다.

 

 

ViT Training & Fine-Tuning

ViT 모델은 대규모의 데이터셋을 개인이 학습시키기는 어렵기 때문에 일반적으로 large-scale dataset에 대해 ViT를 pre-train 하고 downstream task에 대해 fine-tuning을 수행한다.

 

이를 위해 pre-trained Multi Layer Perceptron (MLP) head를 제거하고 0으로 초기화된 D×K feedforward layer를 추가한다.

이때 높은 화질로 fine-tunning 하는 것이 도움 된다. 그 이유는? 고화질에 Patch size가 같게 유지되면 effective sequence length가 커진다.

 

Conclusion

• Computer vision에 self-attention을 사용하는 방법과 달리, image를 squence of patches로서 NLP에서 사용하는 standard Transformer encoder로 처리한다.
• 대규모의 데이터 셋과 fine tunning을 적용했을 때 뛰어난 성능을 보인다.

 

논문 구현 코드 github

https://github.com/minglet/paper_review/blob/main/vit.ipynb

 

 

Reference

https://simonezz.tistory.com/102

https://velog.io/@mink7878/Vision-Transformer-An-Image-is-Worth-16x16-Words-Transformers-for-Image-Recognition-at-Scale-논문-리뷰

https://medium.com/rate-labs/transformer-based-optical-character-recognition-with-pre-trained-models-논문-학습-3435f3d5afec

https://wikidocs.net/137253

https://re-code-cord.tistory.com/entry/Inductive-Bias란-무엇일까