ViT: Transformers For Image Recognition At Scale

ABSTRACT

이 논문은 이미지 인식 분야의 기존 중심 구조인 CNN 대신, 이미지를 작은 패치로 나누고 이를 토큰처럼 처리하는 Transformer 기반 모델인 Vision Transformer(ViT)를 제안한다. ViT는 대규모 데이터로 사전 학습한 뒤 전이 학습을 적용하면, ImageNet, CIFAR-100, VTAB 등 다양한 벤치마크에서 기존 최첨단 CNN 모델과 비슷하거나 오히려 더 좋은 성능을 보이면서도 상대적으로 적은 계산 자원으로 학습할 수 있음을 보여준다.

 

1 INTRODUCTION

Transformer는 이미 자연어 처리 분야에서 압도적인 성과를 거두며 표준 모델로 자리 잡았지만, 컴퓨터 비전 분야에서는 여전히 CNN이 주로 사용되었다. 기존 연구들 역시 CNN 구조를 유지한 채 attention을 일부 결합하는 방식에 머물렀고, Transformer를 이미지에 직접 적용하려는 시도는 많지 않았다.

저자들은 이러한 한계를 문제 의식으로 삼아, 이미지를 작은 패치로 나눈 뒤 이를 단어처럼 Transformer에 입력하는 구조를 제안한다. 초기 실험에서는 데이터가 충분하지 않을 경우 CNN보다 성능이 떨어지는 모습을 보였는데, 이는 Transformer가 CNN처럼 지역성이나 이동 불변성과 같은 inductive bias을 거의 갖고 있지 않기 때문이다.

하지만 대규모 데이터로 사전 학습을 진행하여 데이터가 충분히 많아지면, 모델이 직접 패턴을 학습할 수 있기 때문에 inductive bias의 중요성은 점점 줄어들고, 학습 능력 자체가 더 중요한 요소가 된다. 이때 ViT는 기존 CNN을 뛰어넘는 성능을 보여주며 이미지 인식 분야에서도 Transformer 구조가 충분히 경쟁력이 있음을 입증한다.

 

2 RELATED WORK

먼저 Transformer는 자연어 처리 분야에서 표준 모델로 자리 잡으며, 대규모 사전 학습과 전이 학습을 통해 뛰어난 성과를 꾸준히 보여 왔다. 하지만 이 구조를 이미지에 그대로 적용하면, 모든 픽셀이 서로 attention을 수행해야 하기 때문에 계산량이 픽셀 수에 비례해 제곱으로 증가하는 문제가 발생한다. 이로 인해 현실적인 해상도의 이미지에는 적용하기가 쉽지 않았다.

이러한 한계를 해결하기 위해 기존 연구들은 attention 범위를 제한하거나, 계산을 근사하는 다양한 방법을 제안해 왔다. 일부 연구는 국소 영역에만 attention을 적용했고, 또 다른 연구들은 sparse 구조나 축 방향 attention을 활용해 연산량을 줄이고자 했다. 이런 방식들은 일정 수준 이상의 성과를 내기는 했지만, 구조가 복잡해지고 구현이 까다로워진다는 단점이 있었다.

ViT와 가장 유사한 연구로는 2×2 크기의 작은 패치 단위로 attention을 적용한 모델이나, 픽셀 기반 생성 모델인 iGPT가 있다. 그러나 이러한 모델들은 처리할 수 있는 해상도에 제약이 있거나, 이미지 분류 성능이 제한적이라는 한계를 가지고 있었다.

최근에는 ImageNet보다 훨씬 큰 데이터셋을 활용해 성능을 끌어올리는 연구 흐름도 활발하게 이어지고 있었다. 이러한 많은 연구들은 대규모 데이터로 CNN을 사전 학습하는 방식에 집중해 왔지만, ViT는 CNN이 아닌 Transformer 자체를 이미지에 직접 적용했다는 점에서 차별성을 가진다.

 

3 METHOD

 

ViT의 핵심 목표는 원래 Transformer 구조를 가능한 유지하는 것이다. 기본 아이디어는 이미지를 작은 패치로 나누고(겹치는 부분없이 패치 크기에 맞춰 나눈다고 한다), 이를 단어처럼 벡터로 변환해 Transformer에 입력하는 것이다. 각 패치는 선형 변환을 통해 임베딩되고, 위치 정보를 유지하기 위해 positional embedding이 더해진다. 또한 BERT처럼 class token을 추가해, 이 토큰이 전체 이미지를 대표하도록 설계되었다. Transformer 내부 구조 역시 기존과 동일하다. Multi-head self-attention과 MLP 블록이 반복되며, LayerNorm과 residual connection을 그대로 사용한다.

이러한 설계는 CNN과 달리, 이미지에 특화된 inductive bias가 거의 없다. CNN은 구조적으로 지역성과 이동 불변성을 갖고 있지만, ViT는 대부분의 공간 정보를 데이터로부터 직접 학습해야 한다. 그래서 데이터가 적을 때는 성능이 떨어지지만, 데이터가 많아질수록 강해진다.

 

ViT는 fine-tuning 단계에서 더 높은 해상도의 이미지를 사용해 성능을 높이는데, 이때 해상도가 높아지면 패치 개수가 많아지면서 기존 위치 임베딩을 그대로 사용할 수 없게 된다. 이를 해결하기 위해 위치 임베딩을 2차원 보간(interpolation) 방식으로 재조정해 새 해상도에 맞춘다.

Inductive bias
: CNN은 작은 커널로 주변 픽셀을 중심으로 특징을 추출하고, 같은 필터를 전체 이미지에 공유해 사용한다. 이로 인해 위치에 상관없이 같은 패턴을 인식할 수 있으며, 구조적으로 지역성과 이동 불변성이라는 inductive bias를 갖는다.

 

4 EXPERIMENTS

저자들은 ImageNet, ImageNet-21k, JFT-300M과 같은 서로 다른 규모의 데이터셋에서 ViT를 사전 학습한 뒤, 여러 이미지 분류 벤치마크(ImageNet, CIFAR-10/100, Pets, Flowers, VTAB 등)에 전이 학습하여 성능을 비교한다.

 

ViT는 BERT 구조를 기반으로 Base, Large, Huge 세 가지 크기로 설계되었으며, 패치 크기와 모델 크기에 따라 계산량이 달라지도록 구성되었다. 패치 크기가 작을수록 시퀀스 길이가 길어지기 때문에 계산량이 증가한다.

학습에는 Adam과 대규모 배치를 사용하고, fine-tuning 단계에서는 SGD와 고해상도 입력을 활용한다.

 

이 논문은 "전체 모델"을 다시 학습해 평가하는 fine-tuning 방식과, 모델을 고정한 채 "선형 분류기만 학습"하는 few-shot 방식, 두 가지 방법으로 성능을 측정했다. Few-shot 평가는 빠르게 모델의 표현력을 확인하기 위한 보조 지표로 사용되었다.

 

SOTA 비교 실험에서는 ViT를 당시 최고 성능을 기록하던 BiT (ResNet 기반 대규모 전이 학습) 와 Noisy Student (EfficientNet 기반 반지도 학습) 모델과 직접 비교한다. 그 결과 ViT-H/14는 ImageNet, CIFAR-100, VTAB 등 대부분의 벤치마크에서 기존 모델과 비슷하거나 더 좋은 성능을 달성한다. 특히 ViT가 훨씬 적은 학습 비용으로 이러한 성능을 달성했다.

 

ViT는 CNN에 비해 이미지에 대한 inductive bias가 거의 없기에, 학습 데이터의 규모에 특히 민감하다. 따라서 저자들은 ViT가 잘 작동하기 위해서는 얼마나 많은 데이터가 필요한가를 위해 데이터 크기를 달리한 여러 실험을 진행했다.

먼저 ImageNet, ImageNet-21k, JFT-300M처럼 점점 규모가 커지는 데이터셋으로 ViT를 사전 학습한 뒤 성능을 비교했다. 작은 데이터셋에서는 정규화를 적용했음에도 불구하고, 큰 ViT 모델이 오히려 작은 모델보다 성능이 낮게 나타났다. 하지만 데이터 규모가 커질수록(JFT-300M 수준에 이르러서야) ViT가 기존 CNN 모델을 앞서는 성능을 보였다.

 

다음으로 저자들은 JFT-300M에서 일부 데이터만 무작위로 뽑아 학습하는 실험도 진행했다. 이때는 정규화 효과를 배제하고, 순수하게 모델 자체의 특성을 비교하는 데 초점을 맞췄다. 작은 데이터셋에서는 ViT가 ResNet보다 더 쉽게 과적합되는 경향을 보였으나, 데이터가 충분히 많아지면 ViT는 CNN보다 더 좋은 성능을 기록했다.

이 실험을 통해 데이터가 적을 때는 CNN이 가진 구조적 편향이 큰 장점으로 작용하지만, 데이터가 충분히 많아지면 이러한 편향은 더 이상 필수가 아니고 오히려 Transformer처럼 유연한 구조가 더 큰 성능을 낼 수 있다는 것이다. 또한 few-shot 환경과 VTAB 저데이터 실험에서도 ViT는 안정적인 성능을 보여주었으며, 매우 적은 데이터 상황에서도 전이 학습 가능성이 높다는 점을 확인할 수 있었다.

 

 

ViT는 모델과 데이터가 커질수록 성능이 꾸준히 상승하는 특징을 보였다. 반면 ResNet은 일정 규모를 넘어서면 성능 향상이 둔화되는 경향이 나타났다. 특히 같은 계산 비용 조건에서도 ViT가 더 좋은 성능을 내는 경우가 많았는데, 이는 Transformer 구조가 대규모 학습 환경에서 더 효율적으로 확장된다는 점을 보여준다.

 

 

ViT는 무작위로 패치를 보는 것이 아니라, 물체의 윤곽이나 중요한 영역에 자연스럽게 집중하는 모습을 보였다. 초기층에서는 작은 패턴과 질감을 중심으로 attention이 형성되고, 깊은 층으로 갈수록 전체 구조와 의미적인 영역에 집중하는 경향이 나타났다. 이는 CNN에서 계층적으로 특징을 학습하는 방식과 유사하다. 또한 class token은 중요한 패치들에 집중하며, 전체 이미지를 대표하는 벡터로 잘 작동하고 있음이 확인되었다.

 

저자들은 MoCo v3 기반의 contrastive learning 기법을 사용하여, 라벨이 없는 ImageNet 데이터만으로 ViT를 사전 학습하였다. 이를 통해 정답 라벨 없이도 모델이 의미 있는 표현을 스스로 학습할 수 있는지를 검증하고자 했다.

실험 결과 이러한 방식으로 학습된 ViT는 지도 학습으로 사전 학습한 모델과 비교해도 성능 차이가 크지 않은 것으로 나타났다. 이는 ViT가 대규모 라벨 데이터에 크게 의존하지 않더라도 충분히 효과적인 표현을 학습할 수 있음을 보여주며, 감독 학습뿐 아니라 자기지도 학습 환경에서도 활용 가능한 범용적인 구조임을 보여준다.

 

Contrastive learning
: 의미적으로 유사한 데이터 쌍은 임베딩 공간에서 가깝게, 서로 다른 데이터는 멀어지도록 학습하는 표현 학습 방법이다. 라벨이 없어도 데이터 간 관계만으로 특징 표현을 학습할 수 있어 self-supervised learning에서 널리 사용된다.

 

 

5 CONCLUSION

결론에서 저자들은 ViT의 핵심 기여를 다시 한 번 정리한다. ViT는 CNN을 사용하지 않고도 Transformer 구조만으로 이미지 인식 과제에서 최고 수준의 성능을 달성할 수 있음을 처음으로 보여준 모델이다. 특히 충분히 큰 데이터와 함께 학습될 경우, ViT는 기존 CNN 기반 모델과 유사하거나 그 이상의 성능을 보이면서도 구조적으로는 더 단순하다는 장점을 가진다. 이러한 결과는 향후 비전 분야의 중심 아키텍처가 CNN에서 Transformer 기반 구조로 점차 이동할 가능성을 보여준다.

또한 저자들은 향후 연구 방향으로 더 큰 모델과 더 많은 데이터, 그리고 자기지도 학습 기법과 결합한 ViT 연구가 중요한 흐름이 될 것이라고 전망한다.