[Classification] ViT: Vision Transformer

 

 사족조차 잊어버린 바쁘다 바빠 현대 사회

 

* Transformer를 이해하고 읽으면 훠어얼씬 쉽습니다. 제 Transformer 포스팅을 참고해주세요^,^

https://hyelimkungkung.tistory.com/33?category=935163 

[NLP] Tansformer

 

 An Image Is Worth 16 X 16 Words:

Transformers For Image Recognition At Scale

 

초록 Abstract

Transformer 구조가 NLP에서 대세로 자리잡음에 따라 CV 분야에서도 이를 활용하고자 하는 움직임이 많아졌다. 하지만 여전히 합성곱 신경망에 연결하는 등 Transformer가 RNN을 성공적으로 제거했던 것처럼 CV에서는 CNN을 성공적으로 제거한 사례가 없었다. 하지만 본 논문에서는 CNN에 rely 하는 것은 필요하지 않을 뿐더러 순수하게 transformer만을 이용해서 바로 image patch의 sequence에 적용하면 좋은 성과를 낼 수 있음을 보여주었다. 한편 큰 데이터 셋에 pre-trained한 후에 mid-size나 작은 사이즈의 이미지 recognition에서 downstream task를 수행할 때 이 비젼 트렌스포머ViT는 더욱 좋은 성과를 보여주었다. 

 

1. Introduction

 자연어 처리 분야에서는 Transformer의 효율성과 확정성 덕분에 매우 큰 사이즈의 데이터 셋도 가능해졌다. 그러나 cv 분야에서는 여전히 합성곱 구조가 dominant 했다. 왜냐하면 NLP의 성과를 따라하려고 이제 CNN 같은 구조와 self-attention 을 combine 하거나 합성곱을 대체하는 등의 시도가 있었으나, 이론상으로는 효율적이지만 현실적으로 specialzed 한 어탠션 패턴 때문에 현대 하드웨어로는 확장시키기 어려운 한계가 있었다. 

 

 따라서 본 논문에서는 standard한 Transformer를 최대한 modify하지 않고 바로 이미지에 적용시키고자 했다. 이를 위해서 (1) 이미지를 패치로 쪼갠다 ( 한 token이 input sequence가 되었던 것처럼, 이미지에서도 token 비스무리한 걸 만들어야 했고, 그게 패치로 쪼개진 이미지다) (2) 그리고 이 패치를 linear embedding한  sequence를 Transformer의 input으로 이용했다. 

 

 중간 사이즈의 데이터 셋인 ImageNet을 강한 정규화strong regularization을 사용하지 않고 훈련을 하면, 비슷한 규모의 ResNet보다 몇 % 떨어진 정확도를 보인다. 다소 실망스러운 결과인데, 이는 예측된 결과라고 한다: Transformer는 CNN에 내제되어 있는 inductive bias 가 부족하기 때문이다. 이 말인 즉슨, CNN 구조를 사용함으로써 얻을 수 있는 translation equivariance와 지역성을 학습할 수 없기 때문에 충분한 데이터로 학습하지 않는다면 일반적으로 잘 쓸 수 없다는 뜻이다. 

 

(* inductive bias 에 대해서는 여기 참조)

 

 따라서 엄청 큰 데이터 셋(14M-300M이미지)로 학습하게 되면 상황이 긍정적으로 변한다.데이터 셋이 커지면 inductive bias를 이길? 그런 성능이 나오기 때문이다. 따라서 ViT는 충분히 큰 규모의 데이터로 pre-trained 시킨 후에 더 적은 데이터의 task로 transfer-learning 시키면 excellent한 결과가 나온다고 한다. 한마디로 양치기로 승부함 

 

2. Related Work 

 

 크게 model architecture와 관련 없다. 넘어가도 괜찮다. 

 

 Transformer같은 경우 큰 텍스트 코퍼스(말뭉치라고도 하던데) 를 이용해서 pre-trained 한 후 해결해야 하는 task를 수행하기 위해 fine-tuning을 한다. 대표적으로 BERT가 있는데, 중요하진 않지만 자기 지도 학습이라는 새로운 걸 알아버려서 좀 써볼까 한다! 관심이 없는 사람은 접은 글을 펴보지 않으면 된다!! 

 여하튼 BERT는 denoising이라는 pretext task를 통해서 자기지도 학습을 시행했다~ 라는 맥락이다. 

 

 ** self - supervised learning 자기 지도 학습

 

 머신러닝에 크게 두 가지 학습 방법이 있는데, 지도 학습과 supervised 비지도 학습 unsupervised이 그 둘이다. 

 지도 학습을 할 수 있다면 참 좋지만, 불행히도 질 좋은 label이 모두 갖춰진 데이터 셋은 흔치 않다! 그래서 대신 고안해 낸 것이 self-supervised learning이다. 아래와 같은 개요를 따른다

 

(1) 사용자가 지정한 task(pretext task)를 수행함

  단 이때 사용자가 직접 unlabeld 데이터에 label을 달아서 input으로 이용함 

  pretext task를 통해서 직접적으로 우리가 하고자 하는 목표와는 상관이 없더라도, 데이터 자체에 대한 이해도를 높이는 것을 목표로 한다. 

(2) weight는 freeze시키고, (1)에서 학습된 모델을 transfer-learning하여 downstream task를 수행함

 downstream task는 우리가 해야하는 task라고 생각 ( 이 경우, classification이 되겠다)

 

자세한 사항은 아래 참고 

https://hoya012.github.io/blog/Self-Supervised-Learning-Overview/

https://velog.io/@tobigs-gm1/Self-Supervised-Learning

 

 이미지에 self-attention을 적용시키기 위해 각 픽셀이 서로 다른 픽셀에 attend 하게 naive한 시도를 한 적이 있다. 그러나 이때 픽셀 수의 quadratic하게 cost가 필요했기에 실질적인 input  사이즈에는 적용시킬 수 없었다. 그래서 image 처리에 Transformer를 적용시키기 위해서 여러 시도들이 시행되어 왔고 관련 work들을 소개한다. (1) applied self-attention only in local neighborhoods for each query pixel instead of globally => completely replace convolutions (2) scalable approximation to global self-attention. 하지만 이런 attention 구조는 컴퓨터 비전 분야에서 좋은 결과가 나올 것이 자명했으나, 복잡한 기술을 요구했기 때문에 현실적으로는 사용이 어려웠던 듯하다. 

 

 한편 본 논문과 가장 관련 깊은 모델은 input 이미지에서 2 X 2 이미지 패치를 추출해서 full self-attention을 시도하는 것이었다. 이 말만 언뜻 들어도 ViT와 굉장히 유사한데, 큰 스케일의 pre-training을 통해서 SOTA를 달성했다는 것이 다른 점이라고 한다. 더욱이 앞서 언급한 모델은 작은 resolution 이미지에만 잘 작동하는 반면, ViT의 경우 medium한 resolution에도 잘 작동한다는 점에서 주목할 필요가 있다. 

 

Figure 1

 

Figure 1에 대해서 대략적인 모델 구조를 알 수 있다. 그럼 본격적으로 해보자! 

실은 별거 없다. Transformer만 잘 알면.. 어떤 변형이 있는지만 알면 될듯..

 

3. METHOD

 

 최대한 Transformer의 구조를 변형하지 않았다. 

 

3.1 Vision Transformer(ViT)

 

 자연어를 다루는 NLP에서는 input이 1차원의 token embeddings의 sequence였다. 하지만 이제 우리는? 이미지를 처리해야 한다. 그래서 이 2차원의 이미지를 처리하기 위해서 input의 차원을 변형해준다. 

 

 기존의 이미지가 $x \in \mathbb{R}^{H \times W \times C}$ 크기라면,

$x_{p} = \in \mathbb{R}^{N \times (p^{2} \cdot C)}$, 즉 패치들로 이루어진 sequence로 resize 해준다. 

이때 flattened 2D patches라고 원문에서는 언급한다. 

그림으로 설명해보겠다! 난 그림에 자신없다. (편의상 흑백이라고 생각하자 즉 C=1인 경우다)

 처음에 $x_{p}$가 그냥 패치 하나인줄 알고 뒤에 있는 내용을 줄줄이 이해못했는데, 그게 아니고 패치들의 묶음을 $x_{p}$로 부른다. 실은 차원만 잘 봤더라면 해결됐을 문제... 

* notation
(H,W) : 원본 이미지의 해상도(그냥 가로 세로 라고 생각)
C : 채널 수 (이미지 색상 채널)
(P,P) : flattened 2D 패치의 해상도(패치의 가로 세로)
N = HW/$P^{2}$ : P X P 사이즈로 flatten 함에 따라 생기는 패치의 수 
D : constant latent vector size  flatten한 패치를 D 차원으로 learnable linear projection을 통해 (Eq(1)에서 $E$)  mapping 함 (자세한 내용은 Equation1 참조), 이때 projection한 결과를 패치 임배딩이라고 부른다. 
 

 그래서 위에서 언급한대로 과정을 따르면 아래 수식에서 $z_{0}$인 input을 만들 수 있다.

 다만 아래 수식에서 $x_{class}, E_{pos}$는 아직 언급은 안 했으니 나중에 고려하도록 하자!  

 

 $z_{0} = [x_{class}; x^{1}_{p}E; x^{2}_{p}E; \dots ; x_{p}^{N}E] + E_{pos},$

 

$E \text{ } \in \mathbb{R}^{(P^{2}\cdot C) \times D} ,\  E_{pos} \text{ }\in \mathbb{R}^{(N+1) \times D} \dots (1)$

 

 $z_{0}$의 두번째 term인 $x^{1}_{p}E$를 보겠다. 이때 $x^{1}_{p}$는 패치 묶음에서 첫번째 패치를 의미한다. 

 이 첫번째 패치와 E를 곱해줌으로써 얻는 $x^{1}_{p}E$의 차원은? $\mathbb{R}^{1 X D}$이다!! 

 이렇게 길이가 D인 벡터로 바꿔줌으로써, 우리는 패치 임배딩한 결과를 얻게 된다. 나아가 앞에서 언급한 learnable은 E에 대해 하는 말이다. 

 

휴 한 시름 놨다... 

 

 그 다음으로 위에서 나중에 보기로 했던 $x_{class}$를 보겠다. 대관절 얘는 뭐냐??

 Transformer를 이용해서 NLP를 수행했던 모델, BERT에서는 [class]라고 하는 토큰이 있었다고 한다. 이걸 ViT에도 적용시키기 위해서 본 논문 팀이 데려왔는데, classification을 도와주는 녀석이다. 앞에서 임배딩한 패치의 머리에 저 learnable embbeding을 달아준다. $z_{0}$의 첫번째 term($z_{0}^{0}$)인 $x_{class}$이다. 

 Transformer encoder의 마지막 레이어의 output에서 이 $x_{class}$의 state가 image representation $y$로서 기능을 하게 된다.  Eq.4로 넘어가자!

 

 $$ y = LN(z_{L}^{0})$$

 

(LN은 layer normalization) L 번째 layer의 z의 0번째 term, 즉 $x_class$ 값이 되겠다. 

 

 pre training, 그리고 fine tuning할 때 classification head는 $z_{L}^{0}$에 attached 됐다고 한다. classification head는 pre training 때는 은닉층에서 MLP로 시행하고, fine tuning 때는 하나의 선형 layer로 시행했다고 한다. 

 

 이제 마지막으로 앞서 보기로 하고 넘겼던 $E_{pos}$를 보자. 

 이건 대충 눈치 채겠지만, 포지션 임배딩과 연관이 있다. 즉 위치 정보를 유지하기 위해서 필요한 term이다. 학습가능한 보통의 1차원 위치 임배딩을 사용했다고 한다. (2차원 임배딩을 써도 큰 성능 개선이 없었기 때문이다) 그래서 여기까지 거친, 즉 패치 임배딩에 포지션(위치) 임배딩까지 거친 이 녀석이 인코더의 인풋으로서 이용되었다. 

 

 자 그럼 Transformer의 인풋을 지금까지 알아봤다. 그럼 이제 작동하는 걸 수식으로 한 번 보자! 

 

 

 

 

 요렇게 생긴 구조를 수식으로 바꾸면 다음과 같다. (인풋과 아웃풋 모두 같이 놓고 보자) 

 

 $z_{0} = [x_{class}; x^{1}_{p}E; x^{2}_{p}E; \dots ; x_{p}^{N}E] + E_{pos}, E \text{ } \in \mathbb{R}^{(P^{2}\cdot C) \times D} ,\  E_{pos} \text{ }\in \mathbb{R}^{(N+1) \times D} \dots (1)$

 

$z_{l}^{'} = MSA(LN(z_{l-1})) + z_{l-1} \quad l = 1, \cdots N \dots (2)$

 

$z_{l} = MLP(LN(z_{l}^{'})) + z_{l}^{'} \quad l = 1, \cdots N \dots (3)$

 

$y = LN(z_{L}^{0}) \qquad\dots (4)$

 

 

MSA: Multihead Self-Attention

LN: Layer Normalization

MLP: Multi Layer Perceptron

 

이제 작동하는 것도 다 봐버렸다!

 

- Inductive bias

 Intro에서 언급을 했었다. Transformer에서는 CNN에 비해 이미지 specific한 inductive bias가 부족하다. 왜냐하면 CNN에서는 locality, two-dimensional neighborhood 구조, 그리고 translation equivariance가 전체 모델을 통해서 각 레이어에 backed 되어 있기 때문이다. 한편 ViT에서는 오직 다층 퍼셉트론에서만 지역적이고 translation equivariant한 특성이 나타난다. (self-attention의 경우 전역적이다) 

 그렇기 때문에 two dimensional neighborhood 구조는 아주 조금 쓰인다 (1) 모델을 초반에 이미지를 패치로 자를때 (2) fine tuning할 때 서로 다른 해상도의 이미지에 포지션 임배딩을 해줄 때 

 

- Hybrid Architecture 

 CNN과 함께 ViT를 쓰면 어떨까? 하는 생각에서 나온 하이브리드 구조다! 이때 input은 raw한 이미지 패치 대신에 CNN을 통해서 나온 특성맵 feature map으로부터 input sequence를 이용한다. 즉 위에서 봤던 Eq.1 에서의 E는 CNN 특성맵으로부터 추출된 패치에 multiplication 된다. 이때 spatial 패치 사이즈는 1 X 1이었다고 한다. 마찬가지로 input 임배딩과 포지션 임배딩도 실시해준다. 

 

 

3.2 Fine-Tuning and Higher Resolution

 

 ViT는 큰 데이터 셋에 pre-train을 시키고, 그 후에 작은 데이터 셋으로 fine-tune을 시행했다. 이때 downstream task라고 하는데, 그냥 쉽게 생각해서 지금 당장 해결해야 하는 과제라고 생각하면 된다. pre train을 할때 쓰는 데이터 셋은 우리가 원래 분류할 때 쓰고 싶은 데이터가 아니지 않은가. 그런 맥락이다. 

 fine tuning을 하기 위해서는, pre-trained prediction head를 없애고 0으로 초기화시킨 D X K 차원의 FF layer를 붙여주었다고 한다. 이때 K는 우리가 분류해야하는 class의 수다. 

 

저기 하이라이트 친 녀석을 바꿨다는 말이다!

 

그나저나 보통 fine-tuning할 때 더 높은 해상도의 사진을 쓰는 게 좋다고 하다. pre-training을 할 때 보다. 더 높은 해상도의 이미지를 feeding할 때, 우리는 패치 사이즈는 동일하게 유지하면 더 큰 sequence length의 결과가 나온다고 한다. ViT의 경우, 이제 이런 길이의 sequence를 처리하는 것은 메모리가 닿는 만큼 가능하지만, 이렇게 되면 미리 학습한 포지션 임배딩이 더 이상 의미가 없어진다. 그렇기 때문에 이때는 원래의 이미지에서 패치의 위치를 참고해서 pre-trained한 포지션 임베딩에 2D의 보간법을 시행했다고 한다. 흠 그래서 훈련 데이터로 썼던 녀석보다 더 좋은 해상도를 써도 잘 학습할 수 있었다는 말이군!

 

4. Experiments 

 

 

 ResNet, ViT, hybrid 이렇게 3가지의 모델의 learning capabilities를 비교해봤다. 

 그래서 각 모델이 가장 학습이 잘 될 때를 알아보기 위해서 다양한 사이즈의 데이터 셋으로 pre-train, 그리고 test 했다. 

 pre-trianing할 때 ViT의 경우 가장 low한 계산 cost가 들었다. 

 

4.1 Setup

 

- Datasets

 

 train: ImageNet(class: 1K, image:1.3M) , ImageNet-21K(class: 21k, image: 14M), JFT(class:18k, image:303M)

 task set: ImageNet validation set, cleaned-up ReaL lables, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102

 

 데이터 셋 전처리는 Klensnikov et al(2020)을 따름

 

+ 19-task VTAB classifcation suite 셋도 이요했다. 여러 task가 있고, 각 task 별로 1000장의 training example이 있다. 

task는 크게 세 가지로 나눈다: Natural, Specialized(medical, statellite), Structured(require geometric understanding ex. localization)

 

 - Model Variants 

 

Table 1

 

ViT-Huge 모델을 제외한 두 친구는 BERT의 것을 고대로 가져왔다고 한다. 

 

모델 variants 읽는 방법: ViT-L/16 이 있으면 L은 large model, 16은 input 패치 사이즈

(* 이때 패치 사이즈가 작을수록, cost가 비싸다는 걸 잊지말자!)

 

 또 비교를 위해서 사용하는 ResNet 같은 경우 Batch Normalization layer 대신에 Group Normalization을 이용했다고 한다. 

 한편 하이브리드에서는 intermediate feature map을 ViT로 feed 했고, 이때의 패치 사이즈는 1 "픽셀".

 

- Training &  Fine-tuning 

 

 Train 할 때: 세 모델 모두 Adam optimizer를 사용했다고 한다. 하이퍼 파리미터는: $\beta_{1} = 0.9, \beta_{2}=0.999$ 

 배치 사이즈는 4096(이렇게 큰 배치를 쓸 수 있다니 역시 GPU대박) 

 weigh decay는 좀 높은 0.1 (=> 더 빠른 transfer이 가능해짐) 

 

 Fine-tuning 할 때: momentum과 함께 SGD 사용, 배치 사이즈 512

 

 여기서 언급안 된 자세한 사항은 Appendix D.1 

 

 - Metrics

 

  성능 지표로 크게 두 가지를 이용한다 (1) fine-tuning accuracy (2) few-shot accuracy

 (1): fine-tuning이 끝난 후 각 데이터 셋에 대한 퍼포먼스를 이용한다

 (2): 흠 few-shot learning은 들어봤는데.. 이건 뭔지 잘 모르겠다 원문 표현대로는

 obtained by solving a regularized lest-squares regression problem that maps the (frozen) representation of a subset of training image to {$-1,1^{K}$} target vectors. 

 

 (1) 계산을 위주로 하되, 너무 expensive 해지면 (2)으로 갈아타! 

 

4.2 Comparison to State Of The Art

 

Table 2

 

누가 결과를 묻거든 고개를 들어 표를 보게 하라.. 

 

 보면 outperform하는 것을 볼 수 있다! 다만 이때 계속 강조하고 있는 건 pre-train할 때 드는 cost가 다른 모델에 비해 적다! 라는 점이다. (그런데도 30일이나 걸렸다 이거 실화냐;;)

 

Figure 2

 

VTAB에 대한 결과는 따로 Figure 2에서 확인할 수 있다. 보면 Specialized task를 제외하고선 대부분 ViT가 outperform 하고 있다. 

 

4.3 Pre-training data requirements 

 

 JFT-300M 데이터 셋으로 pre-train을 하면 매우 성과가 좋았는데, 그렇다면 이렇게 좋은 성과를 보기 위해서 최소한으로 필요한 데이터셋 크기는 어느 정도일까? 의 의문에서 시작된 섹션이다. 크게 2가지의 실험을 소개한다.

 

 (1) 데이터 셋 사이즈를 키우면서 ViT 모델을 pre-train 시켜봤더니:

 좀 더 큰 규모의 모델인 Vit-L을 작은 데이터 셋(ImageNet)으로 pre-train한 결과는, 좀 더 큰 규모의 데이터 셋으로 pre-train한 작은 규모의 모델인 ViT-Base 모델에 비해서도 성능이 떨어졌다. ImageNet-21k로 pre-train을 시켜봐도 Vit-L과 Vit-Base의 성능이 비슷했다고 한다. 오직 JFT-300M으로 학습한 것만이 규모가 큰 이점을 누릴 수 있었다고 한다! 

Figure 3

 

 (2) JFT-300M 데이터 셋에서 9M, 30M, 90M 의 크기로 subset를 만들어서 pre-trian 해봤더니:

 ViT overfit more than ResNets with comparable computational cost on smaller dataset. 이 (2)번 실험을 통해서 알 수 있는 것은 CNN의 inductive bias를 통해서 얻는 이점은 smaller dataset에서는 크지만, 더 큰 데이터 셋에서는 직접적으로 모델에서 패턴을 학습하는게 훨씬 효과적이다! 라는 사실을 암시한다. 

 

 4.4 Scaling Study 

 

Figure 5

 

알 수 있는 사실 

 (1) 성능과 연산 사이의 trade-off를 고려하면 ViT가 ResNet보다 dominate하다. 동일한 성능을 보이기 위해서 필요한 연산이 2-4배 적다. 

 (2) small computational budgets일 때, 하이브리드 모델이 ViT를 살짝 outperform 하는데, 규모가 큰 모델이 되면 그런 이점은 사라진다. 

 (3) ViT appear not to saturate within the range tried, motivating future scaling efforts.

 

 

4.5 Inspecting ViT

 

그림 및 주석으로 이해 바람..

 

4.6 Self-supervision

 

future work로 남겨버리심

 

5. Conclusion

 

 이미지 인식 분야에 Transformer를 직접적인 적용을 통해 ViT 모델을 구상해냈다. inductive bias를 처음에 이미지를 패치로 나눠줄 때 제외하고 사용하지 않았다는 점에서 이전의 연구와 차별점을 가진다. 대신, 이미지를 패치들의 sequence로 받아들였고, NLP에서 쓰인 것처럼 standard한 Transformer 인코더로 그 패치들을 연산?했다. 많은 이미지 분류 데이터 셋에서도 SOTA를 달성했으며, 상대적으로 pre-train cost가 싸다는 점에서 박수쳐줄 만하다.(라고 실제로 쓴 건 아니고 내가 쓴거다 ㅎㅎ) 

 다른 컴퓨터 비전의 영역에서 적용해볼 것과, self-supervised pre-training에 대한 연구가 futrue에 진행되어야 한다는 말로, 본 논문을 마무리한다. 

 

 

 

[reference]

참고하고 싶어도 없었음 ㅠㅠ 난 레퍼런스 없이 못 살아 

https://arxiv.org/pdf/2010.11929.pdf

 

사족

 그나저나 2021년 논문이라니(정확하게 2021년에 나온 건 아니지만).. 드디어 현대를 살고 있는 느낌이 든다

흑흑.. 일주일에 세 개나 블로그 올리다니 나 방학 하고 있는 거 맞아..??

 

 다들 이 긴 글 읽느라 수고하셨습니다. 저도 방금 이 글 스크롤 보고 깜짝 놀랐어요.

 Appendix에 대한 이야기는 다루지 않아서 마음이 아프지만 여유가 생기는 날 돌아오겠습니다.

 읽어주셔서 감사해요.. 우리 인생 파이팅 👏