ViT(Vision Transformer) 모델 설명

Introduction

트랜스포머는 NLP분야에서 정말 획기적인 발전을 이루고 뛰어난 성과를 이루어졌다. 하지만 이미지 인식에서는 아직도 CNN이 주로 사용되고 있다. ViT는 2021년 ICLR에서 An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale 논문에서 처음 발표되었으며, CNN없이도 트랜스포머를 통해 뛰어난 성능을 보이는 Computer Vision 모델을 만들었다. 해당 모델은 현재까지도 후속 연구가 진행되고 있으며, ViT를 백본으로 사용되는 모델 또한 많으며, 인공지능 역사에서 굉장히 중요한 모델 중 하나인 CLIP에서도 ViT를 활용하고 있다.

Background

ViT를 이해하기 위해서는 inductive bias가 무엇인지에 대해 이해를 해야한다.

Inductive Bias란 한국어로 번역하면 귀납 편향인데, 머신 러닝 모델이 학습을 할때 선천적으로 가지는 가정을 의미합니다. 보통 모델들은 generalization problem 문제가 있는데, 이는 brittle(불안정, input이 달라지면 결과가 바뀐다)하다는 것과 spurious(데이터의 본질을 파악하는게 아닌 결과과 bias를 학습하는 것)을 의미합니다. 이를 해결하기 위해서 처음 보는 상황에서 정확한 예측을 위해 추가적인 가정을 하는 겁니다. 다시 한번 요약하자면 학습되지 않은 새로운 데이터에도 적절한 추론을 할 수 있도록 일반화를 도와주는 가정입니다.

위 표를 보면 각 머신러닝 모델에 따른 Inductive Bias를 나타냅니다.

• Fully Connected : FC는 입력과 출력 사이의 뉴런이 모두 연결되어 특별한 가정이 없어, Inducted bias가 거의 없습니다.
• Convolutional (CNN) : CNN은 필터를 통해 local receptive field만을 연결합니다. 또한 동일한 필터가 이미지 전체에 적용됩니다. 해당 특성은 인접한 픽셀끼리 강한 상관 관계가 있다고 가정(Locality)을 합니다. 또한 사물의 위치가 바뀌어도 출력값의 위치만 바뀌면서 값은 유지합니다.(Translation Equivariance)
• Recurrent (RNN) : RNN은 데이터가 순차적으로 처리합니다. 이를 통해 입력이 시간 축에 따라 연속적으로 들어온다고 가정합니다.(Sequential)

참고로 ViT는 Computer Vision 모델이기 때문에 CNN과 비교를 해야한다. Transformer는 Self-Attention만을 사용하기에 CNN과 같이 local receptive field를 보지 않습니다. 따라서 CNN보다 Inductive Bias가 더 적다고 볼 수 있습니다. 하지만 이러한 지역적인 패턴을 충분히 데이터셋을 사용하면 CNN보다 더 좋은 성능을 보인다고 합니다.

METHOD

다음 figure는 ViT의 전체적인 과정을 설명하는 그림입니다. 전체적인 과정은 총 5가지 과정을 이미지가 처리됩니다.

1. 입력 이미지를 패치로 변환

기존의 이미지가 $H \times W \times C$ (높이 × 너비 × 채널) 일때 이를 $P \times P$의 크기로 분할 합니다. 하나의 패치는 $P \times P \times C$의 크기를 갖게 되고, 총 $N = \frac{HW}{P^2}$개의 패치가 생성됩니다. (논문 상에서는 조금 더 복잡하게 표현했지만 내용의 핵심만 간단하게 서술하였습니다.) 해당 패치는 트랜스포머의 입력인 토큰으로서 사용됩니다.

2. Patch Embedding

\[z_p = E x_p\]

이제 trainable linear projection을 통해 해당 패치를 flatten 시켜 D차원으로 변환한 후 patch embedding을 진행 한다. 여기서 $x_p$는 flatten 시킨 D차원의 벡터이고, $E$는 학습 가능한 projection matrix이다. 그래서 $z_p$는 변환된 D차원 패치 벡터가 된다.

3. [CLS] 토큰 추가

\[Z_0 = [z_{class}; z_1; z_2; …; z_N] + E_{pos}\]

BERT에서 처럼 클래스 토큰을 추가해줘야한다. 여기서 BERT에서 [CLS]와 같은 역할을 하는 $z_{class}$를 첫번째 토큰으로 하여 최종 입력 벡터 $Z_0$를 구성한다. 여기서 해당 토큰은 처음 입력될 때는 label과 같은 역할을 위한게 아니라, 최종 출력에서 전체 이미지를 대표하는 벡터의 역할을 한다. 즉 처음에 입력될 때는 특별한 의미를 갖지 않지만 학습이 진행되면서 self-attention을 통해 수집한 모든 패치들의 정보를 종합하고 이미지의 전체 representation을 나타내는 역할을 하게된다.

그리고 뒤에 보면 $E_{pos}$가 있는데 이는 위치정보를 갖고 있는 positional embedding을 나타낸다.

4. Transformer Encoder

\[\begin{gather} z_{\ell}{\prime} = \mathrm{MSA}\bigl(\mathrm{LN}(z_{\ell-1})\bigr) + z_{\ell-1} \\ z_{\ell} = \mathrm{MLP}\bigl(\mathrm{LN}(z_{\ell}{\prime})\bigr) + z_{\ell}{\prime} \\ y = \mathrm{LN}\bigl(z_{0}^{L}\bigr) \end{gather}\]

• LN : layer Nomalization
• MSA : Multi-head Self Attention
• MLP : 두개의 MLP로 구성된 feed forward network
• y : 최종 출력

ViT에서는 해당 $Z_0$를 그대로 트랜스포머 인코더에 집어넣는다. 해당 수식을 보면 트랜스포머 인코더의 수행과정을 보여준다. 위 2개의 수식은 $\ell$번째 인코더에서 연산을 수행하는 과정을 나타내고 있고, 3번째 수식은 마지막 출력에 대해 Layer Nomalization을 적용하는 것이다. 그리고 해당 수식을 통해 residual connection 또한 적용되어 있음을 확인해볼 수 있다.

5. Classification Head

트랜스포머의 마지막 출력 벡터의 [CLS]토큰을 MLP Head에 넣어 최종 클래스를 예측한다. 이때 모델을 사전학습 한 후 특정 데이터셋에 대해 fine-tuning을 진행할 수도 있고, pre-train한 모델은 유지한채로 마지막 MLP 분류기만 학습할 수 도 있다.

위 과정을 애니메이션화한 이미지는 다음과 같다. (출처 : lucidrains ViT 레포지토리)

Inductive bias

위에서 언급됐든 ViT는 CNN보다 Inductive bias가 약하다. CNN는 Locality와 Translation Equivariance의 특성을 가진다. 하지만 ViT에서는 Self-Attention으로 전역(global) 연산을 수행한다. 따라서 MLP에서만 Locality와 Translation Equivariance를 가지게 되며, 패치간의 공간적인 정보를 positional embedding에 의존해 학습하게 된다.

정리하자면 ViT는 이미지의 공간적 정보를 직접 반영하지 않는 한계가 있다. 패치간의 관계를 완전하게 학습해야하기 때문에 작은 데이터셋에서는 일반화 능력이 떨어질 수 있다. 즉 ViT는 positional embedding과 self-attention에 의존해 관계를 학습해야하므로 대량의 데이터셋이 필요하다는 한계가 있다.

Hybrid Architecture (CNN + ViT)

이러한 한계를 극복하기 위해 CNN과 ViT를 결합한 아키텍처도 제안되었다. 기존의 ViT와 다르게 먼저 이미지를 CNN에 통과시켜 feature map을 구하고, 해당 feature map을 패치단위로 분할하여 1x1패치로 취급하여(flatten 시켜) 트랜스포머의 입력으로 주는 방식으로도 사용 가능하다. 이러한 방식은 기존 ViT보다 적은 데이터에서 빠르게 학습이 가능해진다.

ViT 분석

우선 ViT의 실험부분은 생략하도록 하겠다. 그래도 실험부분의 내용을 요약하자면, ViT는 데이터가 충분히 학습된다면 CNN보다 더 좋은 성능을 보이며 다양한 벤치마크에서도 뛰어난 성능을 보였다. 또한 기존 CNN 모델들은 모델의 크기를 키울 수 록 성능이 줄어드는 현상이 있었지만 ViT는 모델 크기를 더 키울 수록 성능이 지속적으로 향상되었다. 즉 ViT는 스케일링의 법칙을 따른다는 것이다. 그리고 연산량을 보면 CNN은 이미지 크기가 클수록 FLOPs(연산량)이 급격히 증가하지만 ViT는 상대적으로 증가폭이 적어 동일한 연산량으로 더 높은 성능을 낼 수있다고 한다.

다음 figure는 ViT에서 주목하는 Attention Map의 예제이다. Attention Map은 이미지에서 중요한 영역을 강조하며, Global하게 이미지를 보고 중요한 부분을 파악할 수 있게 된다. 따라서 CNN보다 더 자연스럽게 중요한 객체에 집중할 수 있게 된다.

다음 figure에서는 ViT가 어떻게 이미지를 처리하는지에 대한 실험결과이다.

제일먼저 가장 왼쪽에 나와 있는 것은 RGB Embedding Filters로 ViT가 이미지를 RGB값을 선형변환(linear embedding)의 principal components를 분석하였다. 여기서 RGB Embedding Filters는 CNN의 Conv 필터와 유사한 역할을 하지만 ViT는 전체 이미지를 고려하기에 방식이 다르다. 이때 ViT의 필터가 초기 입력 단계에서 어떠한 이미지의 패턴을 파악하는지를 확인해 볼 수 있다.

두번째 결과에서는 각 패치간 position embedding similarity를 통해 가까운 위치의 패치들의 position embedding이 유사한지를 확인한다. 해당 이미지를 보면 유사도다 보니 자기 자신과의 유사도는 1이기때문에 본인 위치와의 유사도는 밝게 표시되는것을 볼 수 있고, 공간적으로 가까운 패치들이 유사도가 높음을 알 수 있다.

마지막은 ViT의 layer별 평균 attention distance를 확인한 것이다. X축은 네트워크의 깊이, Y축은 Attention distance를 의미하는것인데 CNN과 달리 처음부터 self attention을 통해 모든 패치와 연결되어 초반 레이어에서도 Receptive Field에 제한이 없는 것을 볼 수 있다. 하지만 초반 레이어에서는 global한 정보를 보되, local한 정보 위주로 보다가 점점 global하게 본다는 것을 볼 수 있다.