https://arxiv.org/abs/2010.11929
An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
While the Transformer architecture has become the de-facto standard for natural language processing tasks, its applications to computer vision remain limited. In vision, attention is either applied in conjunction with convolutional networks, or used to rep
arxiv.org
0. Abstract
- Transformer는 자연어 처리에서는 표준이 되었지만, 비전 분야에서는 여전히 제한적으로 사용되고 있다.
- 기존 비전 모델은 주로 CNN과 함께 attention을 사용하거나 부분적으로만 Transformer를 도입해왔다.
- 본 논문은 CNN 없이도 순수 Transformer만으로 이미지 분류에서 우수한 성능을 낼 수 있음을 보여준다.
- ViT는 대규모 데이터로 사전 학습하고 중소 규모 벤치마크(ex. ImageNet, CIFAR-100, VTAB)에 전이하면 기존 SOTA Conv 네트워크보다 적은 계산 자원으로 더 우수한 성능을 낸다.
1. Introduction
- 배경: NLP에서는 Transformer가 기존 방식들을 뛰어넘는 성과를 보여주었고, 대규모 사전 학습 + fine-tuning 방식이 주류가 됨.
- 문제 의식: 반면, CV는 여전히 CNN 기반 구조가 주류이며, Transformer는 널리 확산되지 못하고 있음.
- 제안 방식 (ViT):
- 이미지를 패치로 나눠 각 패치를 토큰처럼 취급.
- 패치를 선형 임베딩 후 Transformer에 입력.
- 별도 convolution 없이 순수 Transformer 구조만으로 이미지 분류 수행.
- 초기 실험 결과:
- 중간 규모 데이터셋에서는 ResNet 대비 성능이 낮음.
- 이는 CNN의 귀납적 편향(국소성, 변환 불변성 등)이 없기 때문.
- 주요 발견:
- 대규모 데이터(14M~300M 이미지)를 활용한 사전 학습 시, ViT가 ResNet보다 높은 성능을 보임.
- ImageNet-21k 및 JFT-300M 사전 학습 모델은 다양한 벤치마크에서 SOTA 달성.
2. Related Work
- Transformer의 기원과 NLP 성공:
- Transformer는 원래 기계 번역을 위해 제안되었고, 이후 NLP에서 pre-training 후 fine-tuning 방식의 표준이 되었음.
- BERT는 denoising self-supervised learning, GPT 계열은 language modeling 기반 사전 학습을 사용함.
- 이미지에 self-attention 적용 시 문제점: 모든 픽셀이 서로를 참조하게 되면 계산량이 픽셀 수의 제곱으로 증가하여 현실적인 고해상도 이미지에 직접 적용하기 어려움.
- 기존의 해결책들 (Transformer 기반 이미지 처리):
- Local Attention: 각 쿼리 픽셀의 local neighborhood에만 Attention 적용. Local Multi-Head Dot-Product Self-Attention으로 구현되며, 일부 구조에서는 합성곱 자체를 대체함.
- Sparse Attention: 전역 어텐션을 sparse pattern나 축 방향(axial)으로 제한하여 연산량을 줄임.
- CNN + 어텐션: CNN이 추출한 feature map 위에 self-attention을 얹거나, 후처리 단계에 적용. (DETR, Non-local NN 등).
- ViT의 차별점:
- Convolution 없이 순수 Transformer만 사용.
- 이미지를 패치 단위로 분할하고, 각 패치를 토큰처럼 처리.
- 기존 연구와 달리 2×2 같은 극소 패치가 아닌 중간 해상도 패치도 다룸.
- 복잡한 구조 없이도 대규모 사전 학습으로 SOTA CNN 성능 달성 가능함을 입증.
- iGPT와 비교:
- iGPT는 이미지를 픽셀 단위로 시퀀스화하고 비지도 학습 기반의 생성 모델로 학습.
- 분류 성능을 위해선 별도 탐색(linear probing) 또는 미세 조정이 필요하며, 정확도는 ImageNet 기준 최대 72% 수준.
- ViT는 지도 학습 기반의 이미지 분류 모델이며 더 높은 정확도를 기록.
3. Method
ViT는 기존 Transformer 구조를 거의 그대로 사용한다. 구조를 단순하게 유지함으로써, NLP에서 검증된 대규모 Transformer 구현을 그대로 재활용할 수 있다.
3.1 Vision Transformer (ViT)
모델의 개요는 Figure 1에 도식화되어 있다. 표준 Transformer는 1차원 토큰 임베딩 시퀀스를 입력으로 받는다. 이를 2차원 이미지에 적용하기 위해, 이미지를 다음과 같이 변환한다:
- 이미지 \( x \in \mathbb{R}^{H \times W \times C} \)는 패치 시퀀스 \( x_p \in \mathbb{R}^{N \times (P^2 \cdot C)} \)로 변형
- \( x \in \mathbb{R}^{H \times W \times C} \) : 입력 이미지
- \( H, W \) : 원본 이미지 해상도
- \( C \) : 채널 수
- \( P \times P \) : 패치 크기
- \( N = HW / P^2 \) : 패치 개수 = Transformer 입력 시퀀스 길이
Transformer는 모든 층에서 고정된 은닉 차원 D를 사용하므로, 각 패치를 평탄화한 후 학습 가능한 선형 프로젝션을 통해 D차원으로 매핑한다.
$$
\mathbf{z}_0 = [\mathbf{x}_{\text{class}}; \mathbf{x}_p^1 \mathbf{E}; \mathbf{x}_p^2 \mathbf{E}; \cdots ; \mathbf{x}_p^N \mathbf{E}] + \mathbf{E}_{\text{pos}}, \quad
\mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}, \quad
\mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}
\tag{1}
$$
이 결과를 patch embedding이라 부른다.
BERT의 [class] 토큰과 유사하게, 패치 임베딩 시퀀스 앞에 학습 가능한 클래스 토큰 \( x_{\text{class}} \)을 추가한다:
- \( z_0^0 = x_{\text{class}} \)
Transformer 인코더의 마지막 층에서 이 토큰의 출력 상태 \( z_L^0 \)는 이미지 전체의 표현 y로 사용된다.
$$
\mathbf{y} = \text{LN}(\mathbf{z}_L^0)
\tag{4}
$$
pre-training과 fine-tuning 모두에서 \( z_L^0 \)에는 classification head가 연결된다.
- 사전 학습 시: 은닉층 1개를 가진 MLP
- 미세 조정 시: 단일 선형 계층
패치 임베딩에는 위치 정보를 유지하기 위해 패치 임베딩에 position embedding을 더한다. 이때 고급 2D-aware position embedding은 성능 이득이 거의 없어, 표준 학습 가능한 1D position embedding을 사용했다.
위 과정을 거친 임베딩 시퀀스는 Transformer encoder의 입력이 된다.
Transformer 인코더는 multi-head self-attention(MSA)과 MLP 블록이 교차 반복되는 구조다. 모든 블록 앞에는 Layer Normalization (LN)이 적용되며, 각 블록 뒤에는 Residual Connection이 사용된다.
$$
\mathbf{z'}_\ell = \text{MSA}(\text{LN}(\mathbf{z}_{\ell - 1})) + \mathbf{z}_{\ell - 1}, \quad \ell = 1, \dots, L
\tag{2}
$$
$$
\mathbf{z}_\ell = \text{MLP}(\text{LN}(\mathbf{z'}_\ell)) + \mathbf{z'}_\ell, \quad \ell = 1, \dots, L
\tag{3}
$$
Inductive bias.
ViT는 CNN에 비해 inductive bias가 훨씬 적다. CNN은 locality, two-dimensional neighborhood structure, translation equivariance 와 같은 inductive bias가 모델 구조에 자연스럽게 포함되어 있으나, ViT는 이런 성질이 거의 없고 대부분의 공간 정보는 학습을 통해 습득해야 한다.
Hybrid Architecture.
입력 시퀀스는 CNN의 feature map으로부터 생성할 수 있는 하이브리드 구조를 사용할 수 있다. 원시 이미지 대신, CNN feature map에서 추출한 패치를 입력으로 사용하며, 이들 패치에는 동일한 방식으로 임베딩과 위치 정보가 부여된다. 특수한 경우에는 패치의 공간 크기를 1×1로 설정하여 feature map을 평탄화하고, 이를 Transformer의 입력으로 사용할 수 있다.
3.2 Fine-Tuning and Higher Resolution
ViT는 대규모 데이터셋에서 pre-train한 후, downstream task에 맞게 fine-tune하며, 이때 pre-trained prediction head는 제거하고 D×K 크기의 zero-initialized feedforward layer를 연결한다. 여기서 KK는 downstream task의 클래스 수를 의미한다. fine-tuning은 보통 pre-training보다 높은 resolution에서 수행되며, patch size는 그대로 유지되므로 입력 이미지의 해상도가 높아질수록 시퀀스 길이가 길어진다. ViT는 시퀀스 길이가 달라져도 처리할 수 있지만, 기존 position embeddings는 무의미해질 수 있기 때문에 원래 이미지상의 위치를 기준으로 2D interpolation을 통해 조정한다. 이 resolution 조정과 patch extraction이 ViT에서 2D 구조에 대한 inductive bias를 수동으로 삽입하는 유일한 지점이다.
4. Experiments
ResNet, ViT, hybrid 모델의 표현 학습 성능을 다양한 데이터 크기에서 평가한 결과, ViT는 적은 pre-training 비용으로도 높은 성능을 보였으며, self-supervision 기반 실험에서도 향후 가능성을 확인할 수 있었다.
4.1 Setup
Datasets.
모델의 확장성과 전이 성능을 평가하기 위해 ImageNet, ImageNet-21k, JFT와 같은 대규모 데이터셋을 사용하고, pre-training 데이터셋은 downstream task의 test set과의 중복을 제거한 후 여러 benchmark에 전이하여 평가한다. 전이 대상에는 원래의 validation label과 정제된 ReaL label이 포함된 ImageNet, CIFAR-10/100, Oxford Pets와 Flowers, 그리고 19개 task로 구성된 VTAB이 포함되며, VTAB은 Natural, Specialized, Structured 세 가지 task 그룹에서 low-data transfer 성능을 평가한다.
- Natural 그룹:
- 자연 이미지 기반의 task
- 예시: Pets, CIFAR 등
- Specialized 그룹:
- 의료 영상 및 위성 이미지 기반의 task
- Structured 그룹:
- localization처럼 geometric understanding을 요구하는 task
Model Variants.
ViT 구성은 BERT의 Base, Large 설정을 따르며, Huge 모델을 추가해 사용한다. (표기법: ViT-L/16처럼 모델 크기와 patch size를 함께 나타낸다.) patch size가 작을수록 시퀀스 길이가 길어지고 계산 비용이 증가한다.
baseline CNN은 ResNet을 기반으로 하되 normalization과 convolution 방식을 수정한 ResNet (BiT)를 사용한다.
hybrid 모델은 중간 feature map을 patch size 1 pixel로 ViT에 입력하며, 시퀀스 길이를 조절하기 위해 ResNet의 stage 4 출력을 사용하거나 stage 3을 확장해 사용하는 두 가지 구조를 실험한다.
- 방식 (i):
- 일반적인 ResNet50의 stage 4 출력을 사용하는 방식이다.
- 방식 (ii):
- stage 4를 제거하고,
- 해당 계층 수만큼을 stage 3에 추가하여
- 전체 계층 수는 유지하면서 stage 3의 출력을 사용하는 방식이다.
이 두 방식 중 두 번째는 시퀀스 길이를 4배 증가시키며, 그에 따라 계산 비용이 더 큰 ViT 모델이 된다.
Training & Fine-tuning.
- Pre-training
- 옵티마이저: Adam (β₁ = 0.9, β₂ = 0.999)
- batch size: 4096
- weight decay: 0.1
- ResNet조차도 이 설정에서 Adam이 SGD보다 성능이 더 좋았음
- 학습률 스케줄: linear warmup + linear decay
- Fine-tuning
- 옵티마이저: SGD with momentum
- batch size: 512
- ImageNet fine-tuning 세부 설정
- ViT-L/16: resolution 512
- ViT-H/14: resolution 518
- 평균화: Polyak & Juditsky averaging (계수 0.9999)
Metrics.
downstream 데이터셋에 대해 fine-tuning accuracy 또는 few-shot accuracy를 사용해 성능을 평가한다.
fine-tuning은 각 데이터셋에 대해 모델을 재학습한 결과이며, few-shot은 frozen feature를 regularized least-squares regression으로 \( \{-1, 1\}^K \) 의 target vector에 매핑해 계산하고, closed-form 해를 구할 수 있다.
fine-tuning 성능을 주로 보고하지만, 연산 비용이 클 경우 빠른 평가를 위해 few-shot accuracy를 사용하기도 한다.
4.2 Comparison to State of the Art
논문에서는 먼저 ViT의 대형 모델인 ViT-H/14와 ViT-L/16을 기존 문헌의 최첨단 CNN들과 비교하였다. 비교 대상은 두 가지로, 첫째는 대형 ResNet을 이용한 지도 전이 학습 방식인 Big Transfer (BiT), 둘째는 ImageNet 및 JFT-300M에서 라벨 없이 반지도 학습된 EfficientNet 기반의 Noisy Student였다. 현재 기준으로 ImageNet에서는 Noisy Student, 다른 데이터셋에서는 BiT-L이 SOTA 모델로 간주되고 있다.
모든 모델은 TPUv3 하드웨어에서 학습되었으며, 사전 학습 시 소요된 TPUv3-core-day 수가 함께 보고되었다. 이 수치는 TPUv3 코어 수(칩당 2개)와 학습 일수를 곱한 값이다.
Table 2에 따르면, JFT-300M으로 사전 학습한 ViT-L/16은 같은 데이터셋으로 학습된 BiT-L보다 모든 작업에서 더 높은 정확도를 기록하면서도, 필요한 계산 자원은 훨씬 적었다. 더 큰 모델인 ViT-H/14는 특히 ImageNet, CIFAR-100, VTAB과 같은 더 어려운 벤치마크에서 성능을 더욱 향상시켰다. 그럼에도 불구하고 이 모델은 이전의 최신 모델들보다 적은 연산량으로 학습되었다는 점이 주목할 만하다.
다만 저자들은 이러한 사전 학습 효율성은 단지 아키텍처만의 문제가 아니라, 학습 스케줄, 옵티마이저, weight decay 등 다른 요소들의 영향을 함께 받는다고 언급한다. 이에 대해서는 4.4절에서 아키텍처별 성능과 계산량 간의 관계를 통제된 조건 하에서 정리해 제시하고 있다.
또한, 공개 데이터셋인 ImageNet-21k로 사전 학습된 ViT-L/16도 대부분의 데이터셋에서 좋은 성능을 보였으며, 자원 소모도 적어, TPUv3 8코어 기준 약 30일이면 학습이 가능하다는 점도 실용적이다.
Figure 2에서는 VTAB의 작업들을 Natural, Specialized, Structured 세 가지 그룹으로 나누어 비교하였다. 여기서는 기존 기법인 BiT 외에도, VIVI(ResNet 기반으로 ImageNet과 YouTube에서 공동 학습된 모델), S4L(ImageNet 기반의 지도 및 반지도 학습 모델)과의 성능 비교가 이루어졌다. 그 결과, ViT-H/14는 Natural 및 Structured 작업에서 BiT-R152x4와 다른 기존 방법들보다 더 우수한 성능을 보였으며, Specialized 작업에서는 상위 두 모델 간의 성능이 유사하였다.
4.3 Pre-training data Requirements
실험 배경 및 목적
- ViT는 ResNet보다 시각적 inductive bias가 적기 때문에, 모델의 성능이 학습 데이터 크기에 얼마나 민감한지를 실험을 통해 분석하고자 했다.
- 이를 위해 두 가지 실험이 진행되었다:
- 다양한 크기의 정형화된 데이터셋을 사용해 사전 학습 성능을 비교.
- JFT-300M의 일부 서브셋을 무작위로 추출해, 정규화 없이 모델 특성만 평가.
실험 1: 데이터셋 크기에 따른 사전 학습 효과 비교
사용한 데이터셋: ImageNet (약 130만 개), ImageNet-21k (약 1400만 개), JFT-300M (3억 개 이상)
실험 설정
- 작은 데이터셋에서는 성능 향상을 위해 다음과 같은 기본 정규화 기법 세 가지를 조정함:
- weight decay
- dropout
- label smoothing
- 학습된 모델들은 ImageNet에 파인튜닝되어 Figure 3에서 성능을 측정 (다른 데이터셋 결과는 Table 5 참고).
주요 결과
- ImageNet 사전 학습: ViT-Large가 ViT-Base보다 성능이 떨어짐 (과대 모델이 데이터에 비해 과적합되기 쉬움).
- ImageNet-21k 사전 학습: ViT-L과 ViT-B가 비슷한 수준의 성능을 보임.
- JFT-300M 사전 학습: ViT-L이 ViT-B를 확연히 능가하며, 대형 모델의 이점이 명확하게 드러남.
- 비교 모델 (BiT): ImageNet 기준에서는 BiT가 ViT보다 성능이 좋았으나 데이터셋 규모가 커질수록 ViT가 BiT를 추월함.
실험 2: JFT-300M의 무작위 서브셋에서 모델 학습
데이터 서브셋 크기: 9M, 30M, 90M, 300M (전체)
실험 설정
- 정규화 기법은 사용하지 않음.
- 모든 실험에서 동일한 하이퍼파라미터를 유지함.
- 모델의 intrinsic property만 평가하고자 함.
- 학습 중 early stopping 적용, 최고 검증 정확도 기준으로 결과 보고.
- 계산 비용 절감을 위해 full fine-tuning 대신 few-shot linear 평가 방식 사용.
주요 결과 (Figure 4)
- 작은 데이터셋(9M): ViT가 ResNet보다 심하게 과적합함. ex. ViT-B/32는 ResNet-50보다 연산 속도는 빠르지만, 성능은 현저히 낮음.
- 중간 이상 데이터셋(30M, 90M): ViT 성능이 점점 개선되며 90M 이상의 경우에는 ResNet보다 우수한 성능을 보임.
- ResNet152x2와 ViT-L/16 조합에서도 동일한 경향이 나타남.
- 해석:
- 작은 데이터셋에서는 CNN의 inductive bias(예: 국소성, translation equivariance)가 효과적.
- 데이터가 충분히 크면, ViT는 귀납적 가정을 강제하지 않고도 더 나은 표현을 학습할 수 있음.
4.4 Scaling Study
- 실험 목표: JFT-300M 기반 전이 학습 환경에서 다양한 모델의 성능 대 사전 학습 비용을 비교하는 스케일링 실험 수행.
- 모델 구성:
- ResNet 7종, ViT 6종, ResNet+ViT 하이브리드 5종 포함.
- 각 모델은 7 epoch 또는 14 epoch 동안 동일한 데이터셋으로 사전 학습됨.
- 핵심 결과 요약:
- ViT는 ResNet보다 훨씬 적은 연산으로 같은 성능을 달성 (약 2~4배 연산 효율 높음).
- Hybrid 모델은 연산 자원이 적을 때만 ViT보다 소폭 우수하나, 모델이 커지면 차이 사라짐.
- ViT의 성능은 실험 범위 내에서 포화되지 않았으며, 더 큰 모델로 확장할 여지가 충분함.
- 결론: ViT는 전이 성능과 계산 효율에서 매우 유리한 구조이며, 소형 모델에 hybrid 접근이 잠깐 유리할 수 있지만, 결국 순수 ViT가 더 나은 확장성을 가진다.
4.5 Inspecting Vision Transformer
ViT가 이미지 데이터를 어떻게 처리하는지를 이해하기 위해, 모델의 internal representations를 분석해 보자.
- 입력 이미지는 flatten 패치 단위로 분할되어 linear projection됨 → 저차원 벡터로 변환됨
- Figure 7의 왼쪽은 학습된 임베딩 필터의 Principal Components. 각 패치 내부의 세밀한 구조를 저차원에서 표현하기에 적합한 basis function와 유사한 형태를 보임
- 투영된 패치 벡터에 학습된 position embedding이 추가됨
- Figure 7의 중앙은 모델이 위치 임베딩의 유사도를 통해 이미지 내 거리 정보를 학습하고 있다는 점을 보여줌
- position embedding은 다음과 같은 특성을 보임:
- 가까운 패치일수록 임베딩이 유사
- 행/열 구조 인식 (같은 행이나 열에 있는 패치들은 유사한 임베딩 공유)
- 큰 이미지일수록 sinusoidal structure 나타남
- 이로 인해, 수작업으로 설계한 2D 구조(topological) 임베딩은 추가적 이득이 없음
ViT의 self-attention 메커니즘은 네트워크의 가장 낮은 계층에서도 이미지 전역 정보를 통합할 수 있도록 해 준다. 모델이 이 전역적 통합 능력을 실제로 어느 정도 활용하는지 분석하였는데, 구체적으로, attention weight를 기반으로 모델이 정보를 통합하는 평균 거리를 계산하였으며, 이 값은 CNN에서의 receptive field size와 유사한 개념이다.
- 관찰된 경향:
- 일부 attention head는 입력의 넓은 범위(전역 패치들)를 한 번에 연결해 정보를 처리함
→ CNN보다 훨씬 빠르게 전역 정보를 통합할 수 있음 - 다른 head는 인접한 패치들만을 대상으로 정보를 처리함
→ CNN의 초기 계층처럼 지역적인 특징만 학습하는 경향 - 하이브리드 모델(ResNet + Transformer)은 Transformer 초반 계층에서 국소 attention이 덜 나타남
→ 이는 ResNet 계층이 이미 지역 정보를 충분히 추출해줬기 때문 - 네트워크 깊이가 깊어질수록, attention이 더 멀리 떨어진 위치로 확장됨
→ 점점 전역적인 정보 통합이 가능해짐
- 일부 attention head는 입력의 넓은 범위(전역 패치들)를 한 번에 연결해 정보를 처리함
- 결과적으로, ViT는 semantically relevant regions에 주목하는 경향을 보임 (Figure 6 참고)
4.6 Self-Supervision
Transformer는 NLP에서의 성능뿐 아니라 대규모 self-supervised pre-training의 덕분이라는 점에 착안해, 이와 유사한 방식으로 BERT의 masked language modeling을 모방한 masked patch prediction 실험을 수행하였다. 그 결과, ViT-B/16은 ImageNet에서 79.9% 정확도를 기록했으며, 이는 scratch 대비 2% 향상이지만 pre-training보다는 4% 낮았다. contrastive pre-training은 향후 연구 과제로 남겨두었다.
5. Conclusion
이 연구는 이미지를 패치 단위의 시퀀스로 변환한 뒤 NLP에서 사용되는 표준 Transformer 인코더를 그대로 적용하는 방식을 통해 ViT를 설계하였다. 기존 시각 분야의 self-attention 연구들과 달리, 초기 패치 분할을 제외하면 이미지 전용 inductive bias를 도입하지 않았다. 이러한 구조는 supervised pre-training과 결합할 경우, 여러 이미지 분류 벤치마크에서 기존 최첨단 기법에 비슷하거나 더 나은 성능을 보였으며, 비용도 비교적 낮다.
초기 결과는 고무적이지만, ViT를 detection, segmentation 같은 다른 컴퓨터 비전 과제에 확장하는 문제, 그리고 self-supervised pre-training의 성능 격차를 줄이는 것이 주요 과제로 남아 있다.
초기 self-supervised pre-training 실험에서는 scratch 대비 성능 향상이 있었지만, 대규모 지도 사전 학습에는 여전히 크게 못 미치는 수준이었다. 또한 ViT의 모델 크기를 더욱 scale up할 경우 성능이 추가적으로 향상될 여지가 있는 것으로 나타났다.