오랜만에 인사드립니다! 현재 Vision Language Model에 공부를 시작하게 되었는데요. VLM의 중요한 내용인 Vision Transformer에 대해 이야기해보겠습니다.

오늘 소개드릴 논문은 Dosovitskiy, Alexey, et al. “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.” ICLR 2021. 입니다. 이 논문은 Transformer를 비전 도메인에 직접 적용한 기념비적인 논문으로, 이후 CLIP, BLIP, LLaVA 등 수많은 멀티모달 연구의 백본으로 자리잡았습니다. Introduction, Related Work, Method, Experiments, Conclusion 순서로 이야기하겠습니다.

Introduction

자연어 처리(NLP) 분야에서 Transformer는 사실상 표준 모델입니다. 대규모 코퍼스에서 사전학습 후 작은 태스크에 파인튜닝 하는 패러다임이 정착했고, 1,000억 개 파라미터를 넘는 모델까지 등장했음에도 성능이 포화되는 조짐이 보이지 않습니다.

ViT Model Overview

[그림 1] ViT Model Overview

반면에 컴퓨터 비전에서는 여전히 CNN 계열이 지배적 입니다. NLP의 성공에 영감을 받아 self-attention을 비전에 적용하려는 시도들이 있었습니다. CNN과 self-attention을 결합하거나, convolution을 완전히 대체하는 방식이죠. 그러나 이러한 접근들은 특수한(specialized) attention 패턴 을 사용하기 때문에 현대 하드웨어 가속기에서 효율적으로 확장되지 못했습니다.

저자들은 표준 Transformer를 최소한의 수정만으로 이미지에 직접 적용 하는 실험을 시도합니다. 핵심 아이디어는 다음과 같습니다.

  • 이미지를 고정 크기 패치로 분할

  • 각 패치를 NLP의 단어(토큰)처럼 취급

  • 패치 임베딩 시퀀스를 표준 Transformer에 입력

  • 지도학습 방식으로 이미지 분류 태스크에 학습

흥미로운 점은 이 방식이 중간 규모 데이터셋(ImageNet)에서는 ResNet보다 성능이 떨어진다 는 것입니다. CNN이 본질적으로 가지는 translation equivariance와 locality 같은 inductive bias가 부족 하기 때문입니다.

그러나 대규모 데이터셋(1,400만 ~ 3억 장) 에서 사전학습하면 상황이 완전히 달라집니다. 저자들은 이를 “large scale training trumps inductive bias” 라고 표현하는데요, 이 한 문장이 ViT 논문 전체의 핵심 메시지라고 봅니다.

따라서 다음과 같이 논문의 기여를 이야기할 수 있습니다.

  • 표준 Transformer를 이미지에 직접 적용하는 Vision Transformer(ViT) 제안

  • ImageNet 88.55%, CIFAR-100 94.55%, VTAB 77.63% 등 다수 벤치마크에서 SOTA 달성

  • BiT-L 대비 약 4배 적은 사전학습 비용 으로 우수한 성능 달성

  • “데이터 규모가 inductive bias를 능가한다”는 점을 실증적으로 입증

Related Work

기존 비전 분야에서 self-attention을 적용하려는 시도들을 정리해보면 크게 네 가지로 나뉩니다.

먼저 self-attention을 이미지에 단순 적용(naive application) 하려면 모든 픽셀이 다른 모든 픽셀에 attend해야 합니다. 이 경우 픽셀 수에 대한 이차(quadratic) 비용 이 발생하여 224x224 이미지만 해도 약 25억 번의 유사도 계산이 필요해집니다. 따라서 여러 근사(approximation) 방식들이 시도되어 왔습니다.

  1. Local neighborhood attention: Parmar et al. (2018)은 각 query 픽셀에 대해 전역이 아닌 지역 이웃에서만 self-attention을 적용했습니다.

  2. Sparse Transformer: Child et al. (2019)은 전역 self-attention에 확장 가능한 근사를 도입했습니다.

  3. Block / Axial attention: Ho et al. (2019), Wang et al. (2020a)는 개별 축(axis)을 따라서만 attention을 적용하는 극단적 방식을 제안했습니다.

이러한 접근들은 좋은 결과를 보였지만 하드웨어 가속기에서 효율적으로 구현하려면 복잡한 엔지니어링 이 필요했습니다.

ViT와 가장 유사한 선행 연구는 Cordonnier et al. (2020) 입니다. 이 모델은 입력 이미지에서 2x2 패치를 추출하고 그 위에 full self-attention을 적용했습니다. ViT와 매우 유사하지만 두 가지 차별점이 있습니다.

  • ViT는 대규모 사전학습 을 통해 vanilla Transformer가 SOTA CNN을 능가함을 입증

  • Cordonnier et al.의 2x2 패치는 저해상도 이미지에만 적용 가능한 반면, ViT는 중간 해상도 이미지 까지 처리 가능

즉, ViT의 기여는 아이디어의 독창성 이 아니라 “scale로 증명한 것” 과 “실용적 해상도로 확장한 것” 에 있다고 볼 수 있습니다.

Method

Vision Transformer(ViT)

먼저 전체 구조부터 살펴보겠습니다. ViT의 데이터 흐름은 다음 8단계로 정리할 수 있습니다.

  1. 이미지 입력

  2. 패치 분할

  3. 패치 임베딩 (선형 투영)

  4. [class] 토큰 추가

  5. Position embedding 더하기

  6. Transformer 인코더 L회 통과

  7. [class] 토큰 추출 + LN

  8. Classification head

부분 별로 천천히 알아보겠습니다.

이미지를 패치 시퀀스로 변환

표준 Transformer는 1D 토큰 임베딩 시퀀스를 입력으로 받습니다. 따라서 2D 이미지를 다루기 위해 다음과 같이 변환합니다.

\[x \in \mathbb{R}^{H \times W \times C} \rightarrow x_p \in \mathbb{R}^{N \times (P^2 \cdot C)}\]

  • \(H, W\) : 원본 이미지 해상도

  • \(C\) : 채널 수 (RGB의 경우 3)

  • \((P, P)\) : 각 패치의 해상도

  • \(N = HW/P^2\) : 총 패치 수 (= Transformer의 유효 시퀀스 길이)

ViT-Base/16 (224x224 입력, 16x16 패치) 기준으로 보면, 패치 수는 \(N = 224^2 / 16^2 = 196\) 이 되고, 각 패치 차원은 \(16 \times 16 \times 3 = 768\) 이 됩니다.

이렇게 평탄화된 패치를 학습 가능한 선형 투영 \(E\) 를 통해 D차원에 매핑합니다. 이 결과를 패치 임베딩(patch embedding) 이라고 부릅니다.

[class] 토큰과 Position Embedding

BERT의 [class] 토큰과 유사하게, 패치 임베딩 시퀀스 맨 앞에 학습 가능한 임베딩을 추가 합니다. 이 토큰의 마지막 레이어 출력 \(z_L^0\) 이 이미지 전체 표현 으로 사용되어 분류 head로 전달됩니다.

위치 정보 유지를 위해 패치 임베딩에 position embedding 을 더합니다. 표준 1D learnable position embedding을 사용하며, 더 정교한 2D-aware 방식을 써도 유의미한 성능 향상이 없었다고 합니다. 이는 모델이 1D 위치 정보만으로도 2D 공간 구조를 자동으로 학습 하기 때문입니다(이후 실험에서 확인됩니다).

이를 종합하면, 입력 임베딩은 다음 식으로 표현됩니다.

\[z_0 = [x_{class}; \ x_p^1 E; \ x_p^2 E; \ \cdots; \ x_p^N E] + E_{pos}\] \[E \in \mathbb{R}^{(P^2 \cdot C) \times D}, \quad E_{pos} \in \mathbb{R}^{(N+1) \times D}\]

[class] 토큰 1개 + 패치 N개이기 때문에 총 \(N+1\) 개 토큰이 됩니다.

Transformer Encoder

Transformer encoder는 MSA(Multi-headed Self-Attention) 블록 과 MLP 블록 이 교대로 반복되는 구조입니다. Layer Normalization은 각 블록 앞(pre-norm) 에 배치하고, residual connection은 각 블록 뒤에 연결합니다. MLP는 2개 레이어와 GELU 비선형 함수로 구성됩니다.

  • MSA 블록 (\(\ell = 1, \ldots, L\)):
\[z'_\ell = \text{MSA}(\text{LN}(z_{\ell-1})) + z_{\ell-1}\]
  • MLP 블록 (\(\ell = 1, \ldots, L\)):
\[z_\ell = \text{MLP}(\text{LN}(z'_\ell)) + z'_\ell\]
  • 최종 이미지 표현:
\[y = \text{LN}(z_L^0)\]

마지막 레이어 출력 중 0번째 위치(즉, [class] 토큰)에 LN을 적용한 결과가 최종 이미지 표현이 됩니다. 즉, \(N+1\) 개의 토큰 중 [class] 토큰의 출력만 분류에 사용되고 나머지는 버려집니다.

Inductive Bias

이 부분이 ViT의 가장 중요한 설계 철학입니다. 저자들은 ViT가 CNN 대비 image-specific inductive bias가 훨씬 적다 고 강조합니다.

  • CNN: locality, 2D neighborhood 구조, translation equivariance가 모든 레이어에 내재

  • ViT: MLP 레이어만 local하고, self-attention 레이어는 global

ViT에서 2D 구조가 사용되는 지점은 단 두 곳 뿐입니다.

  1. 모델 시작부에서 이미지를 패치로 자를 때

  2. 파인튜닝 시 해상도가 다른 이미지의 position embedding을 조정할 때

그 외에는 position embedding이 초기에 2D 위치 정보를 전혀 담지 않으며, 모든 공간 관계를 처음부터 학습 해야 합니다. 이 점이 Introduction의 “scale trumps inductive bias”와 직결됩니다.

Hybrid Architecture와 Fine-tuning

raw 이미지 패치 대신 CNN의 feature map으로 입력 시퀀스를 구성 하는 hybrid 모델도 가능합니다. 이 경우 패치 임베딩 투영 \(E\) 가 CNN feature map에서 추출한 패치에 적용됩니다.

또한 ViT는 대규모 데이터셋에서 사전학습 후 다운스트림 태스크에 파인튜닝 하는 방식을 따릅니다. 파인튜닝 시 사전학습된 prediction head를 제거하고 0으로 초기화된 \(D \times K\) feedforward layer를 부착합니다(\(K\) 는 다운스트림 클래스 수).

흥미로운 점은 사전학습보다 높은 해상도에서 파인튜닝 하는 것이 유리하다는 것인데요, 고해상도 입력 시 패치 크기는 유지하고 시퀀스 길이를 늘립니다. 다만 사전학습된 position embedding이 더 이상 유효하지 않기 때문에, 2D interpolation 으로 재조정합니다.

Experiments and Results

저자들은 ViT, ResNet(BiT), Hybrid 세 모델의 표현 학습 능력을 평가합니다. 실험은 크게 SOTA 비교, 데이터 요구량 분석, Scaling Study, 내부 표현 분석으로 구성됩니다.

Setup

세 가지 데이터셋에서 사전학습을 수행했습니다.

데이터셋 클래스 수 이미지 수
ImageNet (ILSVRC-2012) 1,000 1.3M
ImageNet-21k 21,000 14M
JFT 18,000 303M

ViT는 BERT 설정을 기반으로 세 가지 변형을 사용합니다.

모델 Layers Hidden D Heads Params
ViT-Base 12 768 12 86M
ViT-Large 24 1,024 16 307M
ViT-Huge 32 1,280 16 632M

표기는 ViT-L/16 처럼 모델 크기와 패치 크기를 함께 적습니다(L = Large, 16 = 16x16 패치). 시퀀스 길이는 패치 크기의 제곱에 반비례하므로 작은 패치일수록 계산 비용이 높습니다.

Comparison to State of the Art

Comparison with state of the art

Comparison with state of the art on popular image classification benchmarks

가장 큰 ViT-H/14와 ViT-L/16을 BiT(Big Transfer), Noisy Student와 비교한 결과입니다. 핵심 관찰은 다음과 같습니다.

  • ViT-L/16(JFT)이 동일 데이터셋의 BiT-L을 모든 태스크에서 능가 하면서도 연산은 훨씬 적게 사용

  • ViT-H/14 가 ImageNet 88.55%, CIFAR-100 94.55%, VTAB 77.63%로 SOTA 달성

  • BiT-L: 9.9k TPUv3-core-days, Noisy Student: 12.3k vs ViT-H/14: 2.5k

  • 공개 데이터셋 ImageNet-21k만 사용한 ViT-L/16도 0.23k core-days로 충분히 경쟁력 있는 성능

Pre-training Data Requirements

여기가 ViT 논문에서 가장 중요한 실험입니다. “inductive bias가 적은 ViT에게 데이터셋 크기가 얼마나 중요한가?” 라는 질문에 답하는 부분이죠.

Transfer to ImageNet

[그림 2] Transfer to ImageNet (Figure 3 in the paper)

세 가지 데이터셋(ImageNet, ImageNet-21k, JFT-300M)에서 사전학습한 결과를 보면 흥미로운 패턴이 나타납니다.

  • ImageNet(1.3M): ViT-Large가 ViT-Base보다 오히려 떨어짐 (정규화에도 불구하고)

  • ImageNet-21k(14M): ViT-Base와 ViT-Large가 비슷한 수준

  • JFT-300M(303M): 큰 모델일수록 뚜렷한 이득

또한 BiT CNN은 ImageNet에서 ViT보다 우세하지만, 큰 데이터셋으로 갈수록 ViT가 역전 합니다.

JFT의 부분집합(9M, 30M, 90M, 300M)에서 정규화 없이 동일 하이퍼파라미터로 실험한 결과도 일관됩니다. ViT-B/32는 ResNet50보다 9M에서는 훨씬 나쁘지만, 90M+ 에서는 더 나은 성능을 보입니다. 작은 데이터셋에서는 convolutional inductive bias가 유용하지만, 큰 데이터셋에서는 데이터에서 직접 패턴을 학습하는 것이 충분하며 오히려 유리하다 는 결론입니다.

이것이 바로 “scale trumps inductive bias” 의 실증입니다.

Scaling Study

JFT-300M에서 데이터 크기가 병목이 아닌 상황에서 연산 비용 대비 성능 을 평가한 실험입니다. 세 가지 핵심 관찰을 얻었습니다.

  • ViT가 성능/연산 trade-off에서 ResNet 압도: 동일 성능에 약 2~4배 적은 연산

  • Hybrid는 작은 모델에서만 약간 유리: 큰 모델로 가면 차이가 사라짐 (저자들도 “다소 놀랍다”고 표현)

  • ViT는 실험 범위 내에서 포화되지 않음: 추가 scaling의 가능성을 시사

Inspecting Vision Transformer

저자들은 ViT가 어떻게 이미지를 처리하는지 세 가지 관점에서 분석합니다.

  1. 초기 선형 임베딩: 첫 레이어 투영 필터의 주성분이 패치 내 미세 구조의 기저 함수와 유사 → CNN의 초기 필터와 비슷한 역할

  2. Position Embedding: 가까운 패치끼리 유사한 임베딩, 같은 행/열에서 row-column 구조 출현 → 1D learnable이 2D 구조를 자동 학습

  3. Attention Distance: CNN의 receptive field에 대응되는 개념. 일부 head는 최하위 레이어부터 이미지 대부분에 attend (전역 정보 통합)하고, 다른 head는 작은 거리를 유지(지역 특징). Depth가 깊어질수록 attention distance 증가

이 중 attention distance 분석이 특히 인상적이었는데요, ViT가 처음부터 전역 정보를 활용 할 수 있다는 점이 CNN과의 핵심 차별점이라는 것을 시각적으로 확인할 수 있습니다.

Self-supervision

마지막으로 저자들은 BERT의 masked language modeling을 모방한 masked patch prediction 을 예비 실험으로 진행했습니다. ViT-B/16 기준 ImageNet 79.9% 정확도로 scratch 대비 +2% 향상을 보였지만, supervised pre-training 대비 -4% 차이가 있어 아직 격차가 크다는 점을 인정합니다.

Conclusion

저자들은 다음과 같이 논문을 마무리합니다.

  • 접근의 차별점: 패치 추출 단계 외에는 image-specific inductive bias를 도입하지 않음. 이미지를 패치 시퀀스로 해석하고 표준 Transformer encoder로 처리.

  • 전략의 성격: 단순(simple)하면서도 확장 가능(scalable)한 전략이 대규모 사전학습과 결합될 때 놀랍게 잘 작동

  • 성과: 많은 이미지 분류 데이터셋에서 SOTA 달성/능가하면서도 상대적으로 낮은 사전학습 비용

이어서 세 가지 향후 과제를 제시합니다.

  1. Detection, Segmentation 등 다른 비전 태스크로의 확장: DETR(Carion et al., 2020) 결과와 결합하면 유망

  2. Self-supervised pre-training 탐색: 현재의 격차(supervised 대비 -4%)를 메우는 것이 중요

  3. 추가 Scaling: ViT는 포화되지 않았으므로 더 큰 scaling이 추가 향상 가능

흥미로운 점은 이 세 가지 방향이 이후 비전 연구의 로드맵 역할 을 했다는 것입니다. Detection은 Swin Transformer와 DETR 계열로, Self-supervision은 MAE, DINO, MoCo v3로, 그리고 결정적으로 Self-supervision의 contrastive 방향이 CLIP으로 이어져 VLM의 시대 를 열었습니다.

마무리

ViT를 처음부터 다시 정독하면서 느낀 점은 “좋은 논문은 메시지가 단순하다” 는 것입니다. “Scale trumps inductive bias”라는 한 문장으로 요약되는 이 논문이 이후 비전과 멀티모달 연구의 흐름을 통째로 바꿔놓았으니까요.

오디오 처리를 공부하던 입장에서 보면 ViT의 “픽셀 → 패치” 토큰화는 오디오의 “샘플 → 프레임” 토큰화와 개념적으로 매우 유사합니다. 이런 모달리티 간 연결고리를 발견할 때마다 Transformer라는 아키텍처가 얼마나 보편적인지 새삼 느낍니다.

다음 포스팅에서는 ViT의 자연스러운 다음 단계인 CLIP 을 다뤄볼 예정입니다. ViT Conclusion에서 저자들이 “self-supervised gap을 메우는 것이 중요”라고 남긴 바로 그 지점에서 등장한 모델이거든요. 오늘은 여기서 마무리하겠습니다! 감사합니다.😊