[논문] An image is worth 16x16 words (Vision transformer)

What is Attention?

 기존의 기계번역은 인코더가 input을 차례대로 읽어 vector를 만들어내면 이를 이용해 디코더가 translation을 생성하는 과정으로 이루어져 있었다. 그런데 인코딩 과정에서 문장의 길이에 관계없이 고정된 길이의 벡터를 생성하기 때문에 문장구조가 제대로 표현되지 않으며, 트레이닝 세트보다 테스트 세트의 문장이 더 길다면 성능이 떨어지는 문제가 생긴다. 논문에서는 translation을 진행할 때마다 초점을 두어야 할 부분(subset)을 알려주는 방식으로 이 문제를 해결하고자 하였다. 각 단어가 다른 모든 단어와 얼마나 연관되어 있는지(alignment score)를 계산해서 context score를 계산해낸다. 이러한 방식을 attention(해당 단어가 집중해야 할 부분)으로 칭하게 되었다.

(‘Neural Machine Translation by Jointly Learning to Align and Translate’, 2014, Bahdanau et al)

 

 

Vision Transformer(ViT) 개요

 Attention을 이용한 Transformer(2017)는 NLP에서 매우 좋은 성능을 내지만 이미지에서는 예외였다. Transformer가 CNN의 특성에 적합한 inductive bias*가 부족하기 때문이라고 추측되고 있다. 그런데 이 논문에서는 데이터셋을 늘리면 이러한 제약을 해소할 수 있다는 것을 발견했고, 큰 데이터셋에서 학습시킴으로서 Transformer의 비전 태스크에서의 성능을 높이고자 하였다. ViT는 기존 Transformer 모델의 구조를 거의 그대로 사용하는데, 그로 인해 NLP에 적용된 여러 응용 Transformer application을 쉽게 적용할 수 있게 된다.

*inductive bias: 학습 과정에서 경험하지 못한 데이터를 처리하는 데 필요한 가정

 

 

ViT 모델 구조

 ViT에서는 이미지를 여러 개의 조각으로 나눠 패치화된 sequence로 바꿔 Transformer 모델에 입력시킨다. 이때 class 정보를 담을 class token(0th vector)과 패치 간의 위치적 정보를 학습하는 position embedding이 추가로 인코더에 입력된다.

 ViT는 CNN이 없을 때 생기는 문제점을 큰 데이터셋을 활용함으로써 해결한다. 먼저 저화질의 큰 데이터셋에 대해 pre-training 시킨 후, 훈련된 네트워크의 마지막 레이어를 제거하고 분류하고자 하는 클래스 개수만큼의 fully connected layer를 연결해 fine tuning한다. 이 때 고화질 이미지를 사용하면 패치 숫자가 늘어나 pre-training에 사용한 position embedding은 맞지 않기 때문에 새로운 데이터셋의 고화질 크기에 맞게 수정한다.

 

 

 

Code

 

1. Transformer

 Attention과 Feedforward 모듈을 이용하여 트랜스포머 모듈을 구현한다.

 

2. Vision Transformer

 먼저 인풋 이미지를 시퀀스로 변환하는데, 논문에서는 이미지를 14x14 그리드로 나누기 때문에 총 196개의 768 크기 벡터가 생성된다. 이 때 class token이 추가되므로 input size는 197x768이 된다. 여기에 positional embedding을 더해(concat) transformer에 입력시킨다. Transformer를 통과한 마지막 MLP head만 출력하여 class 결과를 얻어내는 구조이다.

 

 

 

코드 참고)

github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit_pytorch.py