논문 리뷰/CV

ViT 논문 리뷰

manfromearth1 2025. 8. 1. 01:36

AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

Summary & Key Contribution

본 논문은 기존의 CNN 구조 없이, 순수한 Transformer 구조만으로 이미지 분류를 수행할 수 있음을 최초로 보여준다.

핵심 아이디어는 이미지를 고정 패치로 분할하고, 이 패치를 토큰처럼 처리하여 NLP에서 사용되는 Transformer Encoder에 그대로 입력하는 것이다.

이러한 접근은 사전 훈련을 대규모 이미지 데이터셋에서 수행하고, 이후 소규모 데이터셋에 파인튜닝함으로써 기존 SOTA인 CNN기반 모델을 능가하는 성능을 입증한다.

또한, CNN 구조 모델의 locality나 translation equivariance 같은 inductive bias 없이도 충분한 데이터와 computation만 존재하면 학습 및 일반화가 가능하다는 것을 증명하였다.

1. Preliminary

Inductive Bias

우선 Bias라는 용어는 Bias and Variance에서 많이 들어봤듯이, 데이터에 관한 편향성을 의미한다. 즉 Bias가 높은 경우 데이터로부터 타겟과 연관성을 찾아내기 어려운, 과소적합의 문제가 발생한다.

여기서 Bias를 학습 알고리즘의 잘못된 가정에 의한 오차라고 부르기도 하는데, 이를 거꾸로 생각해볼 수 있다. 만약 그 가정 자체를 올바르게 한다면, 성능을 높일 수 있을 것이다. Inductive Bias란 모델이 학습 데이터 바깥의 새로운 데이터를 일반화(generalization)할 때, 올바른 예측을 하기 위해 사전에 주어진 구조적 가정이나 제한이고, 즉 학습을 잘 하기 위해 '설계된 편향'인 것이다.

Linear Regression: 독립변수와 종속변수 간의 관계가 선형적이라 가정하고, 오차는 정규분포를 따른다고 가정한다.
Decision Tree: 특정 feature의 값에 따라 데이터를 구분할 수 있다.

CNN 또한 inductive bias를 가지고 있다.

Locality: 가까운 픽셀들끼리 더 높은 상관관계를 가진다고 가정한다. 즉, 지역적인 정보(작은 receptive field)만 먼저 보고 점차 큰 패턴을 인식하도록 설계되어 있다.
- CNN은 작은 Kernel을 통해서 국소적인 픽셀 정보들을 이용하여 정보를 추출하기 때문에, 공간적으로 가까운 정보일수록 더욱 연관이 있으며 중요하다는 귀납적 편향을 내포한다.
Translation Equivariance: "이미지 상의 특정 패턴은 위치에 관계없이 동일한 의미를 가진다"는 가정이다. 특징이 이미지 내 어디에 있든 같은 방식으로 인식된다. 예를 들어, 고양이 귀가 왼쪽 위에 있든 오른쪽 아래에 있든, 동일한 convolution filter로 감지할 수 있다.
- 동일한 필터를 이미지 전체에 공유(weight sharing)하며 적용하기 때문이다.

Transformer with Inductive Bias

Transformer는 CNN 아키텍처에 비해서 inductive bias가 낮다. 그러한 이유는 self-attention에서 나오는데,

CNN은 공간적으로 가까운 정보는 서로 연관이 있다고 가정하지만(filter), transformer은 모든 위치 쌍에 동일한 self-attention을 적용하기 때문에 가까운 것이 더욱 중요하다는 가정이 약하다.
Transformer은 같은 사진에 위치 변화가 일어나는 경우 완전히 새로운 것으로 간주한다.

의문이 하나 들 수 있는 것은, CNN이 inductive bias 덕분에 적은 데이터 셋으로도 좋은 성능을 낼 수 있었다는 것에서 나온다. 왜 inductive bias를 포기하면서 ViT를 다루는지에 대해서 추후에 설명하도록 하겠다.

2. Structure

이 장에서는 ViT의 모델 구조에 관한 간략한 설명과 finetuning 과정에 대해서 다룬다.

Structure

기본적으로 Transformer은 1차원 토큰 임베딩 시퀀스를 입력으로 받는다. 따라서 이를 2D 이미지에 적용하기 위해서는, 입력 이미지를 평탄화된 2D 패치들의 시퀀스로 변형한다.

$(H, W)$ : 원본 이미지의 해상도
$C$ : 채널 수
$(P, P)$: 각 이미지 패치의 해상도
$N = HW/P^2$ : 생성된 패치 수(Transformer 입력 시퀀스의 길이)

그러나 Transformer은 고정된 인풋 차원 $D$ 를 받기 때문에,선형 변환을 거쳐야 한다. 여기서 주목할 점은, BERT의 [class] 토큰처럼 학습 가능한 임베딩을 두는데, 이는 해당 이미지 $y$에 대한 표현이 된다. 토큰은 이후 이미지 분류 레이어와 연결된다.

또한, 위치에 대한 정보를 추가하기 위해 1D position embedding을 추가한다. - (1)

(2): 이전 층의 출력을 LayerNorm → MSA(Multi-head Self-Attention)에 넣고 다시 residual connection으로 더함

Layer Normalization: 레이어 정규화로 학습 안정성 유도
여러 개의 self-attention head로 global context 모델링
Residual connection: 기존 입력을 더해 gradient 흐름 보존

(3): self-attention의 출력을 정규화한 후 MLP에 통과, 다시 residual connction

(4): 최종 인코더 층의 출력 층 [CLS] 토큰을 LayerNorm 후 최종 출력으로 사용

Fine-Tuning

일반적으로 ViT를 대규모 데이터셋에서 사전학습 후, 더 작은 다운스트림 테스크에서 파인튜닝한다.

Pretrain 작업에서는 대규모 데이터셋에서 학습되기 때문에, 분류할 테스크가 다운스트림 테스크와 맞지 않을 수 있기 때문에 마지막 레이어 $W$는 버리고, 새로운 분류기 레이어를 통해 학습한다.

또한 파인튜닝 시에 해상도를 높이는 경우가 성능 향상에 도움이 되는 것이 경험적으로 밝혀졌는데, 이 경우 $P$가 그대로이면 생성되는 이미지 패치 수가 늘어나 입력 시퀀스의 길이가 늘어난다. 물론 Transformer은 attention 기반 연산으로 입력 시퀀스 길이에는 문제가 없으나, 포지셔널 인코딩에는 왜곡 생길 수 있다.

따라서 2D interpolation을 통해서 포지셔널 임베딩을 조절한다.

이렇게 해상도를 조정하거나 패치를 추출하는 과정이, ViT 구조에 2D 이미지 구조에 대한 귀납적 편향을 수동으로 주입하는 부분이 된다.

3. Evaluation of ViT

다음과 같은 비교 모델을 통해 ViT의 우수성을 입증하였다.

모델명	설명
BiT-L (Big Transfer)	ResNet152x4, supervised pretraining on JFT-300M
Noisy Student	EfficientNet-L2, semi-supervised (ImageNet + JFT300M unlabeled)
ViT-L/16, ViT-H/14	Vision Transformer, pretraining on JFT-300M 또는 ImageNet-21k

Summary

ViT-L/16 (JFT-300M 사전학습)는 BiT-L보다 모든 task에서 더 나은 성능
ViT-H/14는 모든 모델 중 최고 성능
연산량(TPU core-days) 도 BiT나 Noisy Student보다 훨씬 적음
ImageNet-21k처럼 작은 공개 데이터셋을 써도, 꽤 괜찮은 성능 (85.3%)