An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (ViT 논문 리뷰)

AGI foundation model에 관련된 스터디를 진행하면서, AI 전공자라면 이제는 기본적으로 알아야 하는 모델들에 대한 논문들을 함께 읽기 시작했다.

나도 많은 모델들을 알고만 있었지, 어떤 방식으로 작동하고 어떤 약점이 있고 어떤 강점이 있는 지에 대해서는 paper review를 자세하게 진행해보지 않았기 때문에 기초부터 다진다는 생각에 논문을 하나씩 읽어 나가고 있다.

 

이번 시간에는 NLP 분야에서 정말 많이 사용되는 transformer를 vision 분야에서 성공시킨 ViT에 대한 논문이다.

NLP에서 transformer가 나오기 시작하면서 LLM에 대한 발전이 시작됐다고 생각한다. 한 분야에서 정말 강한 모델을 만들 수 있는 구조가 나오면서 그 당시, vision 분야를 연구하던 사람들도 당연히 이 아키텍처가 vision 분야에도 적용될 것이라 생각했고, 그렇게 해서 탄생하게 된 모델인 것 같다.

 

[논문 링크]

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

 

참고로, transformer와 관련된 정말 유명한 논문, "Attention is All you need"를 읽고 나서 ViT를 살펴보는 것이 더 도움이 될 것 같다.

 

Abstract

 

• transformer architecture가 많이 사용되고 있지만 vision 분야에서는 X
• vision 분야에서 attention은 CNNs을 보조하는 느낌
  • replace components
  • conjunction with CNNs
• 따라서, CNNs에 의존적이지 않고 sequences of image patch를 통한 image classification을 수행할 수 있는 transformer 구조 제안
• 현재 (이 당시) SOTA인 CNNs보다 좋은 성과를 얻을 수 있었음.

Introduction

• Self-attention 기반 구조인 transformers는 NLP 분야에서 주로 사용되는 모델
  • large text corpus에서 pre-training → 원하는 task로 fine-tuning
  • 모델의 크기를 100B params까지 늘릴 수 있었다.
    • 그럼에도 saturate 되지 않음
  • 하지만 vision 분야에서는 각광 받지 못한 모델
• CV 분야에 적용하려는 여러 시도가 있었으나 효과적인 모델 제시가 어려웠음

그래서, 가장 작은 tuning을 통해 standard transformer에 직접 image를 사용할 수 있는 모델을 개발

• 이미지를 patch로 나누고 sequence of linear embedding으로 transformer에 입력
  • NLP 분야에서 문장을 토큰화하고 transformer에 입력하는 것과 비슷한 원리
• supervised learning으로 image classification 훈련
• 데이터가 작다면 기존 ResNet의 성능을 넘지 못하지만, 데이터가 충분할 경우?
  • NLP에서 fine-tuning을 통해 새로운 task를 수행하는 것처럼, 조금의 데이터로도 task를 전환해서 사용이 가능
• ImageNet-21K / JFT-300M 데이터셋으로 사전 학습 → SOTA에 준하는 성능

2. Related work

Transformer in NLP

• BERT : self-supervised pre-training
• GPT

image에서 단순하게 self-attention ? ⇒ pixel 마다의 attention을 계산해야 하므로 cost가 너무 높다.

따라서 여러 방법들이 시도 됨.

• 각 query pixel에 대한 local neighborhood에서만 self-attention
  • local multi-head dot-product self attention은 convolution을 완벽히 대체
• sparse transformer
• etc

논문에서 제시한 ViT 모델은 input으로부터 2x2 크기의 patch들을 추출하여 full self-attention을 수행하는 모델과 유사

• 차이점이 있다면 더 성능이 좋다는 점
• 2 x 2 픽셀의 작은 patch → 중해상도 이미지까지 처리 가능

3. Method

3.1 Vision Transformer (ViT)

pure transformer를 사용 → scalable NLP transformer architecture + 효율적인 implementation

ViT는 다음과 같은 과정으로 input을 처리한다.

1. $x \in \mathbb{R}^{H\times W\times C}$ (2D) → $x \in \mathbb{R}^{P^2 \times C}$ (1D)
   1. H : Height / W : Width / C : Channel / P : Patch의 크기
   2. $N = \frac{H \times W}{P^2}$ : patch의 개수이자 input sequence의 길이
2. Patch를 flatten + trainable linear projection ⇒ D dimension으로 mapping, 여기서 각각의 patch를 D-dimension으로 mapping한 하나의 sequence를 $\mathbf{z_k}$라고 한다.
   1. 이를 patch embedding이라 한다.
   2. standard transformer는 1D sequence를 입력으로 받기 때문에 이미지를 2D → 1D로 처리해주는 과정을 거쳐야 한다.
   3. transformers는 각 layer에서 일정한 크기 (D)의 latent vector를 사용
3. BERT의 [class] 토큰처럼 embedding된 patch sequence에 학습 가능한 embedding 1개를 추가한다. $(z^0_0 = \mathbf{x} \space class)$
   1. 이 후, transformer의 encoder를 거쳐 나온 $z^0_L$ (즉, label) 활용
   2. Pre-training과 finetuning동안 classification head에 $z^0_L$이 부착됨
4. Position embedding이 patch embedding에 추가 → 위치 정보 보존
   1. learnable 1D position embedding을 사용한다. 왜냐면, 2D position embedding을 사용한 연구로부터 눈에 띄는 성능 향상을 보지 못했기 때문.
5. Position embedding + patch embedding으로 얻어진 vector sequence를 encoder에 입력
   1. encoder : multi-head self-attention + MLP block으로 이루어짐
6. LayerNorm (LN)이 모든 block 앞에 적용되며 Residual connection이 모든 block 뒤에 적용

• (1) : $\mathbf{z_0}$은 입력으로 들어가는 embedding을 의미한다.
  • learnable한 $\mathbf{x}_{class}$와 각 patch에 대한 정보, 그리고 positional embedding이 함께 들어간다.
• (2) : 이렇게 구성된 $\mathbf{z}$는 LayerNorm을 거쳐 multi-self attention을 수행
• (3) : attention이 수행된 $\mathbf{z} \prime$은 MLP를 거치게 된다.
• (4) : MLP를 거친 \mathbf{z}의 가장 앞쪽에 있는 token을 바탕으로 y (label)을 반환

---

multi-head self-attention (Appendix A)

• sequence인 $\mathbf{z}$는 N x D dimesion이다.
• 우리는 sequence 안에 있는 모든 value에 대해서 weighted sum을 계산한다.
• Attention weights ($A_{ij}$)는 sequence의 2개 요소 사이의 pairwise similarity를 계산한다.

---

[Inductive bias]

ViT는 CNNs보다 image-specific inductive bias가 적다.

• CNNs은 locality, 2D neighborhood structure, translation equivariance 등이 각 layer에 내장되어 있다.
• ViT는 MLP layer에서는 local & translationally equivariant / self-attention에서는 global
  • ViT는 2D neighborhood structure를 최대한 적게 사용하려고 함.
  • 처음 시작할 때, position embedding은 2D position 정보를 담지 않는다, 따라서 patch간의 모든 spatial relations은 처음부터 학습되어야 한다.
    • 이 말은, 우리가 flatten을 했을 때, 실제로는 상하좌우로 가까이 있는 pixel들이 flatten을 통해 멀어지게 되면서 2D 차원에서의 공간적인 정보를 활용할 수 없다는 것을 의미한다.
    • 이런 점들을 보완하기 위해 spatial relations을 처음부터 학습하도록 하여 해당 정보들을 반영할 수 있도록 모델을 훈련시키는 것 같음.

[Hybrid Architecture]

image patch를 사용하는 것 대신에 CNNs에서 나온 output인 feature map을 사용해도 된다.

이는 (1)에서 patch embedding을 projection하는 $\mathbf{E}$가 image 자체에 적용되는 것이 아니라 CNNs에서 나온 feature map에 적용되는 것을 의미한다.

이런 경우는 특이한 경우이며, 1 x 1 patch를 사용해서 patch별로 crop하는 것 없이 feature map을 단순하게 flatten해서 transformer에 넣어서 사용한다.

3.2 Fine-tuning and higher resolution

큰 데이터로 ViT를 pre-training → downstream task에 맞춰서 fine-tuning

• fine-tuning을 하기 위해서 pre-trained prediction head를 제거하고 초기화된 feedforward layer를 부착한다. ($feedforward\space layer \in \mathbb{R}^{D \times K}$, K는 class 수)
• fine-tuning시에, pre-training할 때보다 더 높은 resolution으로 진행하면 더 좋은 성능을 얻을 수 있다.
  • high-resolution image를 입력할 때, patch 크기는 동일하게 유지 → sequence 길이가 길어짐
  • 이 경우, 처리하는 것 자체는 문제가 되지 않지만 사전 훈련된 positional embedding은 의미가 없을 가능성이 높아짐
    • 따라서, 원래 이미지에서의 위치에 따라 pre-training된 positional embedding에 대해서 2D interpolation을 수행하여 더 길어진 sequence에서도 위치 정보를 사용할 수 있도록 진행
    • 이 말의 뜻은 앞서 flatten을 통해 2D 위치 정보를 활용하지 못한다는 점 때문에 처음부터 학습하도록 만든다고 했는데, 만약에 input image의 크기가 달라지면 sequence의 길이가 달라지기 때문에 달라진 길이만큼 2D 위치 정보를 활용 못함. 따라서 기존 positional embedding에서 interpolation을 통해 위치 정보를 활용할 수 있도록 만든다는 의미

여기서, resolution을 변경하는 것과 patch extraction하는 것은 2D structure에 대한 inductive bias가 수동적으로 들어감을 의미.

→ 내가 이해한 바로는, CNN에서 추출한 feature map은 inductive bias가 모델을 통해 주입되는 반면에, 이 부분에서는 inductive bias가 우리가 설정한 세팅에 따라 달라지기 때문에 그런 것 같음.

 

4. Experiments

ResNet / ViT / Hybrid architecture, 이렇게 3개의 representation learning capability를 평가

pre-training의 computational cost를 고려하면 ViT는 performance가 좋고 대부분의 벤치 마크 데이터에서 SOTA에 준하는 성능을 얻을 수 있었다.

추후 self-supervised ViT에 대한 작은 실험을 진행

4.1 Setup

[Datasets]

Model scalability 측정 → 모델이 데이터 양이나 데이터의 complexity가 증가함에 따라서 그 성능이 향상되는 능력을 의미

• ILSVRC-2012 ImageNet dataset
  • 1K class + 1.3M images
• ImageNet-21K
  • 21K class + 14M images
• JFT
  • 18K class + 303M high-resolution images

이 외에 downstream task에 필요한 데이터 셋 (CIFAR, Oxford 등) 사용

 

[Model Variants]

ViT architecture에 대한 자세한 variants

• BERT 모델을 기반으로 ViT 구성
  • Base / Large는 BERT-base / Large와 같은 variants를 구성
  • 그보다 더 큰 “Huge”는 ViT에서 따로 정의

Base / Large는 BERT-base / Large와 같은 variants를 구성

• 표기법은 ViT-’model’/’patch size’
  • ViT-L/16 : 16 x 16 input patch를 가지는 ViT-large model
  • 패치 크기의 제곱과 시퀀스 길이는 반비례 → patch size가 작을수록 계산량 증가
• 비교하고자 하는 CNNs은 ResNet을 사용
  • batch norm → group norm, standard convolution 사용
    • 이를 통해 transfer를 개선 (BiT 논문)
• Hybrid의 경우, CNNs으로부터 나온 intermediate feature map을 ViT에 patch size에 맞게 입력
• transfer에 용이한 hyper-parameter 설정 (Appendix B 참조)

[Metric]

downstream dataset에서 few-shot / fine-tuning accuracy로 평가

• fine-tuning
  • 각 모델이 해당 데이터셋에서 fine-tuning한 후의 성능 측정
  • main 측정 방법
• few-shot
  • 연산량이 부담될 떄 사용

4.2 Comparison to State-Of-The-Art (SOTA)

가장 큰 모델인 ViT-H/14 & ViT-L/16에 대해서 최신 CNNs과 비교

• ViT가 BiT나 efficientNet보다 좋은 성능을 보여주며 특히 더 어려운 task에서 더 큰 성능 향상폭을 보여줌.
• 또한 이전 기술보다 pre-training에 계산량이 줄어듬.
  • 일부 parameter에 의해 영향을 받을 수 있긴 함.

• VTAB task를 분류하여 이전 SOTA와 비교

4.3 Pre-training data requirements

ViT는 dataset 크기가 가장 큰 JFT-300M에서 가장 좋은 performance를 보여준다.

기존 CNN 기반 모델보다 inductive bias가 작기 때문에, dataset의 크기가 얼마 정도 되어야 할까에 대한 실험 진행

1. 점차 증가하는 크기의 dataset에서 ViT 모델 pre-training (ImageNet → ImageNet-21k → JFT-300M)

2. 9M, 30M, 90M, JFT-300M에 대한 dataset에서 훈련 및 평가

전반적으로 작은 데이터 셋에서는 convolution inductive bias가 유용하지만, 더 큰 데이터 셋에 대해서는 데이터에서 직접 관련된 패턴을 학습하는 것이 더 유용하다는 것을 확인할 수 있음. (future work로 ViT의 Few-shot에 대한 추가 분석 제시)

4.4 Scaling study

JFT-300M에서 transfer performance 평가

자세한 setting은 paper 참조

4.5 Inspecting Vision Transformer

ViT가 어떻게 image data를 처리하는 지 이해하기 위해서 internal representation을 분석

• ViT 1st layer에서 flattened patch를 low-dimension으로 projection
  • Figure 7의 가장 왼쪽
  • 이 구성 요소들은 각 patch 내의 미세한 구조를 low-dimension으로 표현하기 위한 basis function과 비슷하다.
    • 모든 feature들을 다 capture할 수 있는 가장 기본적인 pattern들이 다 모여있다. → 2차원에서 vector를 표현하기 위해 서로 다른 2개의 vector만 필요한 것처럼. (여기서 2개의 vector가 basis function)
• projection 이 후에 positional embedding이 추가
  • Figure 7의 가운데 그림은 모델이 이미지 내의 거리를 positional embedding의 유사성으로 encoding
  • 이해한 바 : 자세하게 보면 맨 왼쪽 위에 있는 그림은 맨 왼쪽 위에 highlight가 가장 강하다. 이와 같이 2D structure에 대한 정보를 반영하고 있다는 것, 따라서 이 정보 덕분에 비슷한 위치에 있는 patch들간의 유사한 embedding을 얻을 수 있음.
• Grid가 큰 경우에 종종 sinusoidal structure가 나타난다.
• self-attention은 가장 낮은 layer에서도 ViT가 전체 이미지에 걸쳐 정보를 통합할 수 있게 만드는데, 어느 정도 활용하는 지에 대해 조사
  • Figure 7의 오른쪽 그림

4.6 Self-supervision

NLP에서 transformer가 효과적으로 작동할 수 있도록 만든 masking modeling 방법을 차용하여 ViT에서도 patch를 mask하고 예측하는 훈련법을 사용

→ 처음부터 훈련하는 것보다 2% 향상 / supervised pre-training에 비해 여전히 낮은 성능

따라서 이와 관련된 작업은 future work로 제시

 

5. Conclusion

• transformer 구조를 image에 적용
  • 초기 단계를 제외하고 transformer 구조 그대로 사용
• 데이터 규모가 클수록 훨씬 잘 작동됨
  • 그럼에도 cost가 낮게 pre-training이 가능

Future work

• ViT를 classification 외에 다른 작업에 적용
• Self-supervised pre-training 방법에 대한 고안
• ViT의 scaling을 키우는 것이 performance를 향상 시키는가

마치며

 

이 논문 이전까지 computer vision 분야에서 가장 많이 사용되는 모델은 CNN이었다. CNN 모델에서 사용하는 convolution 대신, transformer 구조를 가져와 image patch를 sequential하게 처리하는 방식으로 훈련된 모델 구조는 꽤 혁신적이었을 것으로 보인다. 물론 지금은 더 큰 모델들과 더 많은 모델들이 생겨나고 있지만, ViT를 기점으로 이후에 많은 CV 모델들이 발전을 할 수 있지 않았을까? 생각한다.

 

그리고 늘 느끼는 거지만, 내가 모든 연구가 끝난 결과물만 봐서 그런지 생각보다 엄청 어려운 접근들을 사용하지는 않는 것 같다. 기존에서 어떻게 참신하게 modification을 주느냐가 정말 중요한 것 같다.