4 minute read

ViT

Introduction

NLP 분야에서 Transformer에 영감을 받아 해당 논문은 Tranformer 를 Vision 분야에 직접적으로 적용을 한 첫 논문이다.

해당 논문에서는 Image 에 Sequnencial 한 특성을 부여하도록 하기 위해 Patch라는 단위로 나눈 다음 이 Patch의 Embedding 값을 Transformer 의 Input으로 넣게 된다. Image Patch의 경우에는 NLP 분야에서의 Token 으로 취급된다.(한 단어) 이후 Supervised learning 으로 학습 진행.


그림 1.

해당 Transformer 는 ImageNet 과 같은 mid-sized data에 대해서는 ResNet보다 낮은 성능을 보였으나, Pretrained 된 후 transformer 를 통과한 경우 CNN based 모델보다 월등한 성능을 보였다.

ImageNet data에서 성능이 잘 안나온 이유에 대해서 논문의 저자는 얘기하길 Transformer는 Inductive한 biases가 없기 때문에 Locality와 Translation Equivariance 같은 CNN의 특성이 없다고 한다.

Inductive bias

Inductive bias란, 학습하는 모델이 지금까지 만나보지 못했던 상황에서 정확한 예측을 하기 위해 사용하는 추가적인 가정을 의미한다. 즉, 다시말해 보다 강력한 가정을 하여 좀 더 강한 Filter 로 문제 해결을 하겠다는 것이다.

하지만 이러한 Inductive bias는 모델의 robust 하게 만들지는 않는다. 그 이유는 모델이 내세운 가정 자체가 좀 더 강력한 가정이기 때문에 오히려 낯선 데이터에 대해서 이리저리 튀는 예측을 할 수 있는 문제가 있기 때문이다.

이러한 Inductive bias 는 CNN에서도 잘 나타나는데, 기본적으로 CNN에서의 가정은 feature 가 Locally한 information을 담고 있을 것이란 가정이다. 모델 설계 자체에서 Vision 분야에 공간 정보를 더 뽑자라는 가정이 들어 간 것이다.

하지만 CNN 의 경우에서도 전체에서 Global한 영역의 처리가 어렵기 때문에 Receptive field를 넓히려는 연구가 지속되고 있다. 즉, 다시말해 Inductive bias 가 전체적인 Robustness 보다 우월하다고는 할 수 없다

Transformer의 경우에는 그러나 Attention 구조이기 때문에 전체적으로 Global한 정보는 얻을 수 있겠지만, 추가적인 가정이 부족할 수 있다. 따라서 논문의 저자는 이러한 Transformer 의 specifically locality 의 Inductive bias 의 부족을 한 가지 문제점으로 지적하였다.


그림 2.

또한 Translation Equivariance 의 경우 해당 그림과 같이 동일하지만 특정한 지역적 특성 혹은 Variance 를 갖고있는 동치 그림들에 대해이야기 하는 것인데, Transformer의 경우 전체 정보를 처리하기 때문에 이러한 동치에 대해서도 어떠한 강한 가정을 통해 문제를 풀기 어렵다.

따라서 기존 CNN 모델보다 Robust한 Transformer 기반 모델을 많은 양의 데이터로 학습할 때엔 Robust 하게 잡은 가정들도 상대적으로 Inductive 해지기에 강건하면서도 강력한 모델을 만들 수 있을 것이다.

실제로도 큰 데이터 셋에 대해서는 CNN기반 모델보다 매우 좋은 성능을 보이게 된다.

Method


그림 4.

해당 아키텍쳐가 VIT의 전체 구조인데, 수식적으로 살펴보자.

기본적인 Transformer의 경우 Input으로 토큰 임베딩이 완료된 1D vector를 받는다. 2D image 를 이러한 Input 으로 처리하기 위해 본 논문을 Image를 Reshape한 후 각 Patch를 Flatten하여 1D vector로 만들게 된다.

  • $x \in R^{H \times W \times C} $ : 원본 이미지 $H,W$는 image 의 해상도이며, $C$는 image 의 channel
  • $x_p \in R^{N \times (P^2 \cdot C)}$ : 각 patch 의 shape , $C$는 동일하게 Channel 이며 $N$은 Patch의 수, $P^2$ 이란 각 Patch의 크기 (ex : $4 \times 4$) , $N = HW/P^2$ . 즉, 각 Patch $x_{pi}$ 는 $R^{P\times P \times C} $ 공간에 있는 Tensor 이다.

즉, 각각의 Input 이 들어가게 되는데, 해당 input의 Patch 는 Flatten 된 후 Transformer의 Input으로 들어가게 되는 것이다. Transformer의 경우 D차원의 Latent vector를 사용하게 되는데, train 가능한 linear projection layer를 거쳐 이러한 Patch의 Flatten을 진행하게 된다.

또한 BERT와 흡사하 [class] Token을 Patch Sequence의 앞단에 붙이게 된다. 해당 임베딩은 학습이 가능하며 이미지의 출력 단에서 encoder 에서 y의 역할을 하게 된다. 이후 Equation 을 보며 좀 더 이야기 하겠다.

VIT의 경우 원본 Image 가 2D 이기 때문에 2D positional embedding 을 활용하는 것이 더 합당해보이나 본 논문의 저자는 실험 결과 유의미한 차이를 발견하지 못했기 때문에 1D positional embedding 을 사용하였다.

Equation

  • $z_0 = (x_{class} ; x^1p E ;x^2_pE;,\dots ; x^n_p E) + E{pos}$
    • $E \in ~ R^{(p^2\cdot C ) \times D}, E_{pos} \in ~ R^{(N+1) \times D}$
      • Transformer 의 Input으로 들어가는 Sequence
      • 각 $E$는 embedding 을 말한다.
  • $ z_{l}^{\prime} = MSA(LN(z_{l-1})) + z_{l-1} $, $l = 1,\dots,L$
    • 해당 과정은 Multi-head Self Attention 과정이다.
    • 이전 입력값에 Layer Normalization 을 해준 이후 이전 값을 Skip connection 해준다.
  • \(z_l = MLP(LN(z_{l}^{\prime})) + z_{l-1}^{\prime }\) , $l = 1,\dots , L$
    • Transformer Encoder를 통과한 이후 Layer Normalization 을 한 후 MLP Layer를 거치게 된다.
  • $y = LN(z^0 _L)$
    • 이후 가장 첫단에서 Tokenization 해준 Token 을 output y 로 뽑게 된다.
    • 해당 과정은 Transformer의 Encoder구조를 사용했기 때문에 BERT 와 비슷하게 동작한다.
    • Masking 이 된 어떤 Sequence 를 추론하는 것이 아니기 때문에 Transformer의 Decoder 구조를 사용하지 않았고 따라서 Encoder구조에서 결과값 $y$를 뽑기 위해 사용

ViT 의 경우 Large dataset으로 pretrained를 시킨다. 이후 더 작은 데이터로 Downstream task 를 수행한다. Down-stream task를 진행할 때 prediction head 를 제거한 후 zero-initialized 된 $D \times K$ layer를 붙인다. ($K$는 down-stream class) 또한 high resolution image를 통해 학습하는 것이 좋다고 한다. 또한 Pre-trained 된 positional embedding 은 의미가 없으며, 모호한 길이의 Sequence 를 다룰 수 있다.

Experiments

실험 방식은 VIT base , Large, Huge 모델에 순차적으로 ImageNET-1k(1.3M), ImageNET-21k(14M), JFT-18k(303M) 데이터를 통해 학습을 진행하였다.

이어 Natural, Specialized, Storectured Image에 대해 Fine Tuning을 진행한 뒤 성능을 비교하였다.


그림 5.

Transfer Learning 결과 해당 그림과 같이 SOTA를 달성하였으며, 밑 차트처럼 원 학습 데이터에 대해서도 이전 SOTA인 BIT를 앞선 성능을 보인다.


그림 6.

또한 밑 Plot을 살펴보면 VIT에서의 Pre-train은 필수적인 것이라는 것을 판단할 수 있다. 기존에 적은 양의 데이터로 Pretrain을 진행한 경우 BIT에 비해 성능이 안나오는 것을 확인할 수 있으나, Pre-train 의 Data가 커질수록 이후 transfer learning에서 SOTA를 달성하는 것을 확인할 수 있다.


그림 7.

본 논문의 저자는 ViT가 아직 성능의 포화상태를 달성하지 않았다고 본 논문에서 밝히고 있다. 따라서 더 큰 데이터 셋으로의 training이 이루어진다면, 보다 더 높은 성능 향상을 기대할 수 있다.

Conclusion

VIT는 기존 CNN 이 갖고있던 Inductive bias의 영향에서 보다 더 자유롭고 global한 특성을 가질 수 있도록 설계한 모델이라고 할 수 있다. 국소적 특성에 대한 가정은 VIT 초기 Patch를 생성할 때에만 적용되고 이후에는 Self-attention을 통해 학습을 진행하여 Transformer Encoder 구조를 통해 SOTA를 달성하게 되었다.

아직까지 많은 분야에서 SOTA를 낸 것이 아니라 개선해야 할 부분이 많지만, 비교적 저렴한 pretrain Cost 와 아직까지 bottle neck 에 걸리지 않은 점을 미루어 보았을 때 충분히 발전 가능성이 풍부한 모델이라고 할 수 있다.