[논문정리📃] An Image is Worth 16X16 Words : Transformers for Image Recognition at Scale

An Image is Worth 16X16 Words : Transformers for Image Recognition at Scale

-ViT-

오늘은 Vision Transformer에 대해서 리뷰해보도록 하겠습니다. Vision Transformer, 줄여서 ViT는 원래 자연어 처리(NLP)에서 사용하던 모델을 Vision 쪽에 적용한 모델입니다.

transformer 과 attention 에 대한 개념이 나오는데, [논문정리📃]Attention is all you need 를 먼저 읽어보면 Vision transformer에 대해 이해하기 수월 할 것입니다. :)

Attention is all you need 논문 링크도 걸어두겠습니다!

🌕 Abstract

NLP 에서 100B 의 파라미터가 넘는 사이즈를 학습할 수 있게 된 반면, 컴퓨터 비전쪽에서는 아직 CNN 구조가 주를 이루고 있습니다. 이런 NLP의 성공에 영향을 받아 이미지에 직접적으로 Transformer을 적용하고자 하였고, transformer을 적용하여 ViT는 엄청난 크기의 데이터셋으로 학습하여 SOTA를 달성하는데 성공합니다.

프로세스를 간단히 살펴보면,

자연어에서 입력 시퀀스를 토큰으로 쪼개어 각각을 입력값으로 해서 진행하는 것 처럼 vision 쪽에서도 이미지를 패치 단위로 쪼개어서 개별적으로 입력을 넣어 처리해 진행하게 되는데

한편, transformer을 컴퓨터 비전에 적용하기에는 inductive bias를 갖고 있습니다. inductive bias는 학습에선 없었던 새로운 데이터가 입력으로 들어올 때, 해당 데이터에 대한 판단을 내리기 위해 학습과정에서 습득된 Bias를 가지고 판단을 하게 됩니다. 이것을 inductive bias라고 해석하면 될 것 같습니다.

CNN은 locality 한 정보를 추출하고 convolution연산은 translation equivariance와 locality하다는 특징을 갖고 있습니다.

따라서 이러한 문제를 갖고 있는 CNN은 새로운 분포의 데이터를 입력 받을 때, 해당 데이터에 대한 판단을 내리기 위해 학습과정에서 습득된 Bias인 “지역적 특성” 가지고 판단을 하게 되는 것입니다.

equivariance 란, 함수 입력이 바뀌면 출력 또한 바뀐다는 뜻이며, translation equivariance는 입력 위치가 변하면 출력도 마찬가지로 위치가 변한채로 나온다는 뜻입니다.

locality 란, 이미지를 구성하는 특징들은 이미지 전체가 아닌 일부 지역에 근접한 픽셀들로만 구성되고, 근접한 픽셀들끼리만 종속성을 가진다는 성질입니다.

결국, transformer은 이미지를 패치로 쪼개고 일렬로 세워서 학습(지역적 특징 무시)하는데, 각 픽셀들의 위치와 주변 값들이 중요한 CNN에 있어서, 이 패치들의 입력이 1차원의 값으로 병렬 처리가 된다면 원하는 출력이 나올 수 없게 된다는 뜻인 것 같습니다.

하지만!, 이 논문에서는 엄청난 양의 데이터를 사용해 이 inductive bias를 없앨 수 있었다고 말합니다. 결국.. 데이터가 많은게 한수 였나봅니다.🐱

🌕 Method

위의 이미지와 같이 이미지를 일정한 값의 patch로 분할하고 이 patch 들을 embedding을 시켜서 transformer의 입력값으로 사용하게 됩니다. 개별의 patch를 NLP의 token 처럼 간주하는 것이죠.

프로세스를 자세히 살펴보겠습니다.

먼저 48x48의 이미지를 input image로 한다면, (16x16) * 3 의 형태로 총 9개의 patch 형태로 쪼갭니다. 그리고 각각의 patch들을 linear 하게 prjection 시켜서 patch embedding을 진행합니다. 따라서 768차원의 하나의 벡터가 되는 것입니다.

16x16x3 = 768 = D

그 다음에는 각각의 embedding된 패치에 PE(position embedding) 를 더해준 값을 최종 transformer의 입력 벡터로 사용하게 됩니다. 이때, 맨 앞에는 CLASS TOKEN 이라는 것이 붙는데요, 자연어 처리의 BERT 모델에서 나온 개념을 사용한 것입니다. class token은 학습을 하면서 각 토큰들에 대한 정보값을 저장하게 됩니다. 따라서 이 0번째 패치는 각 패치들의 위치값과 정보들을 가지고 있어 실제로 classification으로 식별 할 때, 이 class token으로 진행한다고 합니다. 나머지 패치들을 학습용으로, 정보를 비교할때는 맨 앞의 패치(class token)으로!

여기서 잠깐!🖐 기존의 transformer 과 vision에서 사용하는 transformer의 차이에 대해 짚고 넘어가겠습니다.

둘을 비교했을 때! Normalization 의 위치가 다른것을 확인할 수 있습니다. Layer normalization의 위치가 중요하다는 것이 후속 연구들에서 밝혀지게 되면서 Attention과 MLP이전에 정규화를 수행하여 더 깊은 레이어에서도 학습이 잘 되도록 했다고 합니다.

🌍 Fine-tuning and higher resolution

ViT는 그대로 사용하기 보다, 먼저 큰 데이터셋으로 pre-train 을 시킨 모델에 fine-tune 하여 사용한다고 합니다. 따라서, transformer encoder 출력값에 하나의 hidden layer를 가진 MLP를 사용하여 pre-train 하게 됩니다. fine-tunning 시에는 랜덤 초기화된 하나의 linear layer를 사용하며 pre-train에서 학습된 positional embedding은 fine-tunning시에 다른 해상도로 이미지에 대한 position embedding을 조정해야 합니다.

학습 초기에 position embedding은 patch의 2D 위치에 대한 아무런 정보를 제공하지 않고, patch 사이에 공간적인 관계는 처음부터 학습되어야 합니다. 논문에서는 고해상도 이미지로 fine tunning 할 때는, 학습된 pre-trained position embedding에 2D interpolation하는 방법(수동으로 bias를 넣는 방법)을 사용합니다.

🌍 hybrid architecture

CNN과 transformer를 함께 사용하는 hybrid archtecture도 소개하고 있습니다. CNN의 feature map으로부터 patch를 추출하여 patch embedding을 적용하는 프로세스를 거치게 되는데, CNN으로 localty한 정보를 추출하고, 이 정보들의 patch의 sequence를 embedding하여 transformer로 전달하는 것입니다. VIT에 비해 좀 더 적은 dataset으로 성능이 saturate하다는 장점이 있습니다. 하지만 큰 데이터셋을 사용한다면 VIT보다 성능이 뒤떨어 집니다.

아래는 그 결과를 보여주는 그래프 입니다.