[개념 정리] 비전 트랜스포머 / Vision Transformer(ViT) (2/2)

 

지난 포스팅에 이어서 비전트랜스포머(Vistion Transformer/ ViT)에 대해서 알아보겠습니다.

 

[개념 정리] 비전 트랜스포머 / Vision Transformer(ViT) (1/2)

 

지난 시간에는 인코더 부분에 대해서 알아보았습니다. 포지셔널 임베딩이 추가된 패치 임베딩이 Transformer의 인코더에 들어가서 Q, K, V로 나뉘고, Multi Head Self Attention + Skip connection, 이후 Multi Layer Perceptron + Skip connection이 추가된 다수의 레이어를 지나 최종적으로 인코더의 출력 벡터가 나오게 됩니다. 이제 그 다음 단계에 대해서 알아보겠습니다.

 

1. Classification Head (분류 헤드) 

- 위 그림에서처럼, 인코더 부분이 끝나고 도출된 벡터는 최종 MLP Head를 만나게 됩니다. ViT 논문에서는 Classification task에 ViT를 사용했는데요. 그래서 해당 논문에서 제시하는 마지막 Head는 Classification head 입니다.

 

- Classification은 다음과 같이 구성되어 있습니다.

 

1) Class 토큰 : ViT는 이미지 패치들을 처리하기 전에, 특별한 [Class] 토큰을 입력 시퀀스의 맨 앞에 추가합니다.(위 그림에서는 * 표시와 함께 Extra learnable [class] embedding 이라고 표시된 부분입니다.) 이 토큰은 이미지 전체에 대한 대표 정보를 학습하는데 사용됩니다. Transformer 인코더를 통과한 후, [Class] 토큰은 이미지 전체에 대한 종합적인 정보를 담게 됩니다. (각 패치들과 상호작용을 하면서 이미지 전반적인 특성과 Context를 학습하게 됩니다.

 

2) Linear Layer : [Class] 토큰의 출력은 위에서 보이는 MLP Head에 들어가게 됩니다. 이는 Linear 형태를 띄고 있는 Layer이며 Fully Connected 되어 있습니다.  이 레이어는 모델이 학습하는 과정에서 최적화되며, 각 클래스에 대한 Logit 값을 출력하게 됩니다. 

 

3) Activation function for classification : Binary class의 경우 Sigmoid, Multiple Class의 경우 Softmax 함수가 Activiation Function이 되어 최종적으로 입력된 Logit 값을 각 클래스에 대한 확률 값으로 변환하고 가장 높은 확률을 가진 클래스를 해당 이미지의 클래스로 선정하게 됩니다.