[X:AI] ViT 논문 리뷰

논문 원본 : https://arxiv.org/abs/2010.11929

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

 

 

1. Introduction

• Self-attention-based architectures, 특히 Transformer는 NLP 분야에서 핵심 모델로 자리잡음

• 한편 CV 분야에서는 CNN이 여전히 주류
• 본 연구는 NLP에서의 Transformer 성공을 바탕으로, 최소한의 수정만으로 이미지를 처리할 수 있는 표준 Transformer 모델 실험

• 이를 위해 이미지를 여러 개의 패치로 나누고, 각 패치를 선형 임베딩하여 Transforemer에 입력
• 이 패치들은 NLP의 단어 토큰처럼 취급
• 모델은 이미지 분류 task를 위한 지도 학습 방식으로 학습

• ImageNet과 같은 중간 규모 데이터셋에서 강력한 정규화 없이 학습한 경우, 이 모델의 정확도는 동급 크기의 ResNet보다 몇 % 낮은 수준에 머무름
• 이러한 결과는 Transformer가 CNN이 가진 Locality(지역성)나 Translation Equivariance(이동 불변성) 같은 Inductive Bias(귀납적 편향)가 없기 때문에, 데이터가 충분하지 않으면 일반화 성능이 낮다는 것을 의미 
  • 쉽게 말해 Transformer는 이미지 구조에 대한 선천적인 가정이 없기 때문에, 데이터가 적으면 일반화가 어려움
  • Inductive Bias : 데이터가 부족해도 일반화를 잘 할 수 있도록 도와주는, 모델 구조 자체에 내재된 '가정' 또는 '지식'
  • Locality : 이미지에서 가까운 픽셀끼리는 서로 관련성이 높다는 가정. CNN은 작은 커널로 주변 픽셀만을 보며 연산하기 때문에 해당 구조를 잘 반영
  • Translation Equivariance : 물체가 이미지 내에서 위치만 바뀌어도 같은 물체로 인식해야 한다는 성질. CNN은 커널이 전체 이미지 위를 움직이면서 같은 연산을 하므로 이 특성을 갖음

• 하지만 데이터셋 크기를 1400만~3억 장 수준의 대규모 데이터로 학습할 경우, Inductive Bias보다 데이터의 규모가 더 중요한 요소로 작용
• ViT는 충분한 규모로 사전 학습한 후, 소규모 데이터셋으로 전이 학습을 진행했을 때 매우 우수한 성능을 보여줌
  • ImageNet: 88.55%
  • ImageNet-ReaL: 90.72%
  • CIFAR-100: 94.55%
  • VTAB 19개 과제 평균: 77.63%

 

 

3. Method

• 본 연구는 ViT를 설계할 때 가능한 Transformer 구조를 그대로 따름
• 이렇게 단순하게 설계하면, 기존의 대규모 Transformer 구조 및 최적화된 구현을 거의 그대로 사용할 수 있다는 장점 있음

   3.1 Vision Transformer (ViT)

• 1) 이미지를 패치로 분할
  • 이미지 x ∈ ℝ^H×W×C 를 P×P 크기의 패치로 나눔
  • 전체 패치 수는 N = HW / P² 
  • 각 패치는 flatten해서 1D 벡터로 바꿈
• 2) 패치 임베딩
  • 각 패치를 D 차원 벡터로 변환하기 위해 학습 가능한 linear projection을 사용
  • 해당 벡터들을 패치 임베딩이라고 부름
• 3) [CLS] 토큰 추가
  • BERT처럼, 전체 이미지를 대표할 수 있도록 [class] token을 시퀀스 맨 앞에 붙임
  • 해당 토큰은 Transformer를 통과한 후 분류에 사용
• 4) Positional Encoding
  • 패치들의 순서를 알 수 있도록 1D Postional Embedding을 더함
  • 2D-aware Embedding도 실험했지만 성능 차이가 없어서 사용하지 않음
• 5) Transformer Encoder 입력
  • 입력 시퀀스는 여러 층의 Transformer Encoder에 입력
  • 각 layer는 다음을 포함
    • multiheaded self-attention (MSA)
    • MLP (GELU 사용)
    • 각 block 앞에 LayerNorm, 뒤에는 residual connection 사용
• 6) 출력
  • 최종적으로 [class] token의 출력 벡터에 분류기를 붙여 예측

https://github.com/lucidrains/vit-pytorch/blob/main/images/vit.gif

 

 

 

• Pytorch에서 제공하는 VIT 시각적 설명
  • https://discuss.pytorch.kr/t/vision-transformer-a-visual-guide-to-vision-transformers/4158

Vision Transformer에 대한 시각적 설명 (A Visual Guide to Vision Transformers)

 

 

   3.2 Fine-Tuning and Higher Resolution

• ViT는 보통 대규모 데이터셋으로 사전 학습한 뒤, 작은 다운스트림 데이터셋에 대해 fine-tuning 진행
• fine-tuning 시 사전학습된 분류기 head를 제거하고, DxK 크기의 zero-initalized 선형층을 새로 붙임 (K : 새 class 수)

• 또한, 사전학습 때보다 더 높은 해상도 이미지로 fine-tune 하는 것이 도움이 되는 경우가 많음
• 이때는 패치 크기는 그대로 유지하지만 이미지 크기가 커지므로 시퀀스 길이가 더 길어짐

• ViT는 시퀀스 길이가 달라도 처리 가능하지만 사전학습된 위치 임베딩이 해상도가 다르면 의미를 잃음
• 그래서 위치 임베딩은 이미지 상 위치 기준으로 2D interpolation을 통해 맞춰줌

• 이미지를 패치로 자르거나, 해상도 조정 후 위치 임베딩 보간을 하는 것만이 ViT에서 2D 이미지 구조에 대한 명시적인 inductive bias를 주입하는 유일한 순간

 

 

4. Experiments

 

   4.2 Comparison to State of The Art

• ViT-L/16 (JFT-300M) → BiT-L보다 적은 연산으로 모든 task에서 더 좋은 성능
• ViT-H/14 → ImageNet, CIFAR-100, VTAB 등 더 어려운 과제에서 특히 우수
• ViT는 기존 SOTA보다 계산량 대비 성능 우수 (2~4배 효율적)

 

 

 

   4.3 Pretraining Data Requirments

• 작은 데이터셋에서는 ViT가 ResNet보다 성능이 낮음 → 과적합
• 대규모 데이터셋일수록 ViT의 성능이 급상승
• 즉, CNN의 inductive bias는 소규모 데이터에 유리, ViT는 대규모 데이터에 강함

 

 

   4.4 Scaling Study

• ViT는 ResNet보다 같은 성능을 2~4배 적은 연산량으로 달성
• 작은 모델에서는 hybrid가 더 좋지만, 커지면 ViT 단독이 더 뛰어남
• ViT는 더 크게 만들수록 성능이 계속 증가 → scaling 여지 많음

 

 

   4.5 Inspecting Vision Trasnformer

• ViT는 위치 임베딩이 2D 공간 구조를 학습함 
• Self-attention은 낮은 레이어부터도 전역 정보 통합 가능
  • 일부 attention head는 국소적(=CNN 역할), 일부는 전역적(=Transformer 강점)
• ViT는 의미 있는 시각 영역에 attend함 → 시각적 인식 가능

 

   4.6 Self-Supervision

• BERT처럼 마스킹된 패치 예측 실험
• ViT-B/16: ImageNet에서 79.9% 정확도 → from scratch보다 2% 향상
• 하지만 여전히 supervised pre-training보다 4% 낮음
• Contrastive 학습 등은 미래 연구 방향

 

5. Conclusion

• 전략 : 이미지를 패치 시퀀스로 바꾼 뒤, NLP에서 쓰는 표준 Transformer에 그대로 넣음

• 성과
  • 단순하지만 대규모 데이터로 사전학습하면, SOTA 성능을 따라잡거나 능가함
  • 그럼에도 불구하고 사전학습 비용은 비교적 낮음

• 남은 과제
  • detection·segmentation 같은 다른 비전 과제로 확장 필요
  • Self-supervised 학습 개선 필요 -> 현재 방식도 향상은 있지만, supervised pre-training과는 격차 존재
  • ViT를 더 크게(scale-up) 만들수록 성능이 더 좋아질 가능성 있음