ViT 읽기 — 이미지를 16×16 단어로 쪼개 Transformer에 넣기

구글 리서치가 2021년 ICLR에서 발표한 An Image is Worth 16×16 Words(Dosovitskiy et al., arXiv:2010.11929)를 정리한다. 한 줄 요약은 — 이미지를 16×16 패치로 잘라 각 패치를 “단어(토큰)“처럼 취급하고, NLP에서 쓰던 표준 Transformer를 거의 그대로 적용했더니, 데이터만 충분하면 CNN을 이긴다는 것이다. 제목 그대로 “이미지 한 장 = 16×16 단어들의 문장"이라는 발상이 논문 전체를 요약한다.

왜 순수 Transformer를 이미지에?

당시 상황을 알면 이 논문의 무게가 보인다. NLP는 이미 Transformer 천하였다. “큰 데이터로 pre-train → 작은 데이터로 fine-tune” 공식이 표준이 됐고, 파라미터를 100B까지 키워도 성능이 꺾이지 않았다. 반면 비전은 여전히 CNN이 왕이었다. Transformer를 이미지에 써보려는 시도는 많았지만 대부분 CNN에 attention을 곁들이는 수준이었고, 순수 Transformer만으로는 안 된다는 게 통념이었다.

기술적 걸림돌은 분명했다. self-attention은 모든 요소가 다른 모든 요소를 본다. 이걸 픽셀 단위로 하면 픽셀 수의 제곱으로 비용이 폭발한다(224×224면 약 5만 픽셀). 현실적으로 불가능하다.

ViT의 해결책은 단순하다. 픽셀 대신 16×16 패치를 하나의 토큰으로 취급하는 것이다. 224×224 이미지는 패치 196개(14×14)로 줄어, Transformer가 감당할 수 있는 “문장 길이"가 된다.

이미지를 16×16 단어로 — 토큰 만들기

첫 단계는 이미지를 토큰 시퀀스로 바꾸는 것이다.

이미지가 토큰 시퀀스로 변환되는 과정

  • 패치로 분할: 이미지를 16×16 패치들로 자른다. 문장으로 치면 단어 여러 개짜리 문장이다.
  • 선형 투영(patch embedding): 각 패치(16×16×3)를 쭉 펴서 학습 가능한 행렬로 D차원 벡터로 바꾼다. NLP에서 단어를 벡터로 만드는 word embedding과 같은 역할이다.
  • 클래스 토큰: BERT를 그대로 베껴, 맨 앞(0번 자리)에 학습 가능한 특수 토큰 [class]를 붙인다. Transformer를 다 통과한 뒤 이 자리의 출력 벡터가 “이미지 전체를 대표하는 값"이 되고, 여기에 분류기를 붙인다.
  • 위치 임베딩(position embedding): Transformer는 순서 개념이 없어서 각 패치가 원래 몇 번 자리였는지를 더해준다. 흥미롭게도 정교한 2D 방식 대신 그냥 1D 학습형으로도 충분했다 — 모델이 알아서 2D 구조를 학습하기 때문이다.

Transformer Encoder 블록

이렇게 만든 토큰 시퀀스가 Transformer Encoder를 통과한다. 이 블록이 ViT의 심장부이고, L번 반복된다.

Transformer Encoder 블록 구조

한 블록은 두 개의 서브층으로 이뤄진다.

  • z' = MSA(LN(z)) + z — LayerNorm 후 멀티헤드 셀프 어텐션, 그리고 잔차 합산(⊕)
  • z'' = MLP(LN(z')) + z' — LayerNorm 후 MLP(GELU 2층), 그리고 잔차 합산(⊕)

여기서 셀프 어텐션이 “모든 패치가 서로를 보는” 부분이다. CNN은 층을 깊이 쌓아야 넓은 영역을 보지만, 어텐션은 첫 층부터 이미지 전체 맥락을 통합할 수 있다. 각 서브층 앞의 LayerNorm과 뒤의 잔차 연결(⊕)은 블록을 깊게 쌓아도 학습이 무너지지 않게 붙잡아주는 장치다. 마지막에는 [class] 토큰의 출력이 MLP Head로 가서 “새/공/자동차…” 분류가 나온다.

모델 크기는 BERT 네이밍을 그대로 물려받았다. ViT-Base(12층, 86M), ViT-Large(24층, 307M), ViT-Huge(32층, 632M). 표기법에서 ViT-L/16은 “Large 모델 + 16×16 패치"를 뜻한다. 패치가 작을수록 토큰 수가 늘어 더 무겁고 더 정확하다(시퀀스 길이는 패치 크기의 제곱에 반비례한다).

데이터가 inductive bias를 이긴다

이 논문의 진짜 핵심은 여기다. 나머지 실험은 전부 이 한 문장을 뒷받침한다.

**Inductive bias(귀납적 편향)**는 모델 구조에 미리 심어놓은 “이미지는 이런 것이다"라는 가정이다. CNN에는 이게 잔뜩 들어 있다 — locality(가까운 픽셀끼리 관련 있다), translation equivariance(고양이가 어디 있든 같은 고양이), 2D 이웃 구조가 층마다 내장돼 있다. 덕분에 CNN은 적은 데이터로도 잘 배운다. 공짜 사전지식이기 때문이다.

반면 ViT에는 이 사전지식이 거의 없다. 패치로 자르는 것과 fine-tune 때 위치 보정, 딱 두 군데만 2D 지식이 들어가고, 패치들 사이의 공간 관계는 전부 데이터로부터 맨땅에서 배워야 한다. 그래서 결과가 데이터 규모에 따라 갈린다.

  • 작은 데이터(ImageNet, 130만 장)만: ViT가 비슷한 크기의 ResNet보다 몇 %p 뒤진다.
  • 중간(ImageNet-21k, 1400만 장): 엇비슷해진다.
  • 큰 데이터(JFT-300M, 3억 장): ViT가 CNN을 역전한다. 큰 모델일수록 격차가 벌어진다.

즉 “구조에 사전지식을 넣느냐"보다 **“데이터를 충분히 주고 스스로 배우게 하느냐”**가 더 강력하다. 논문의 표현으로는 large scale training trumps inductive bias. 이 교훈이 이후 Foundation Model 시대 전체를 관통하는 철학이 된다.

결과와 내부 들여다보기

JFT-300M로 사전학습한 ViT-H/14는 ImageNet 88.55%, CIFAR-100 94.55%, VTAB(19과제) 77.63%로 당시 최강 CNN들(BiT-L, Noisy Student)을 넘어섰다. 그러면서도 사전학습 연산량은 2~4배 적었다 — ViT-H/14가 2.5k TPUv3-core-days인 데 비해 BiT-L은 9.9k, Noisy Student는 12.3k였다. 정확도도 이기고 비용도 싼 것이다. 공개 데이터인 ImageNet-21k로 학습한 ViT-L/16조차 클라우드 TPUv3 8코어로 약 30일이면 훈련할 수 있어 재현 가능성도 열어놨다.

블랙박스가 아니라는 것도 보여준다. 첫 임베딩 필터는 CNN 초기 필터처럼 그럴듯한 기저 함수(edge, 색 패턴)를 학습한다. 위치 임베딩은 아무 사전지식 없이 시작했는데도 가까운 패치끼리 비슷한 벡터를, 심지어 행/열 구조까지 스스로 재발견한다 — 손으로 만든 2D 위치 임베딩이 별 이득이 없던 이유다. 어텐션 거리(CNN의 receptive field에 대응)를 보면 낮은 층에서도 일부 head는 이미지 전체를 보고, 층이 깊어질수록 넓어지며, 최종적으로 분류에 의미 있는 영역에 집중한다.

물론 한계도 있다. 진짜 위력을 보려면 JFT-300M급 초대형(그것도 비공개) 데이터가 필요하고, 이 논문은 분류만 다뤘으며(검출·분할은 이후 다른 연구들의 몫), 자기지도 학습과 지도학습 사이의 격차도 남아 있다. 그럼에도 ViT는 “비전엔 CNN이 반드시 필요하다"는 통념을 깼고, NLP와 비전을 같은 아키텍처로 통일하며 멀티모달·Foundation Model 시대의 문을 열었다.

한 줄 정리

이미지를 단어처럼 쪼개면 언어 모델용 Transformer가 그대로 비전 모델이 되고, 데이터만 충분하면 CNN의 사전지식 없이도 더 잘한다.

Hugo로 만듦
JimmyStack 테마 사용 중