ViT: Vision Transformer

17 min readMay 19, 2024

Transformer 구조는 자연어 처리 분야에서 표준으로 자리 잡았지만, 컴퓨터 비전 분야에서는 적용이 제한적이었습니다. 기존의 컴퓨터 비전에서는 Transformer의 attention 메커니즘을 CNN(Convolutional Neural Network)과 함께 사용하거나, CNN 구조를 유지하면서 특정 구성 요소를 Transformer로 대체하는 방식으로 활용되었습니다.

이 연구에서는 CNN에 대한 의존성 없이 이미지 패치의 sequence에 직접 적용된 순수 Transformer가 이미지 분류 작업에서 뛰어난 성능을 발휘할 수 있음을 보여줍니다. Vision Transformer (ViT) 라는 이 새로운 모델은 대량의 데이터로 사전 훈련된 후, ImageNet, CIFAR-100, VTAB 등의 중소형 이미지 인식 벤치마크에 적용되었을 때 CNN 모델에 비해 좋은 결과를 달성했습니다.

이미지 분류에 Transformer 구조를 성공적으로 적용(ICLR 2021)

핵심은 이미지를 패치(16x16)로 분할하고 각 패치의 선형 임베딩을 Transformer의 입력으로 사용하는 방식이라는 것입니다. 이미지 패치는 NLP에서 단어(토큰)처럼 취급됩니다. 모델은 이미지 분류 작업을 위해 지도 학습 방식으로 훈련됩니다.

ImageNet과 같은 중간 규모 데이터셋으로 훈련된 ViT 모델은 강력한 정규화 없이는 ResNet보다 낮은 정확도를 보였습니다. 이는 Transformer가 CNN이 가진 translation equivariance 및 locality와 같은 Inductive Bias(귀납적 편향)을 가지고 있지 않기 때문입니다. 따라서 충분한 데이터 없이는 일반화 성능이 떨어집니다.

Inductive Bias은 머신 러닝 모델이 학습 과정에서 특정 가정이나 선호도를 갖는 것을 의미합니다. 즉, 제한된 데이터에서 일반화된 결론을 도출하기 위해 사용하는 일종의 ‘힌트’입니다.

CNN은 이미지의 작은 영역(local) 정보를 추출하는 convolutional filter를 사용합니다. 이는 이미지의 특정 부분이 다른 부분과 독립적으로 의미를 가질 수 있다는 가정에 기반합니다. 이미지를 이동시켜도 CNN의 출력은 동일하게 유지됩니다.

하지만 대규모 데이터셋(1,400만~3억 개 이미지)으로 훈련된 ViT는 뛰어난 성능을 보입니다. 대규모 학습이 Inductive Bias를 극복하는 것을 확인했습니다.

ViT는 대규모 데이터셋으로 학습될 경우 Transformer가 컴퓨터 비전 분야에서 CNN을 대체할 수 있는 강력한 모델임을 보여줍니다.

RELATED WORK

Cordonnier et al. (2020)는 입력 이미지에서 2x2 크기의 패치를 추출하고 전체 self-attention을 적용했습니다. ViT와 매우 유사하지만, ViT는 대규모 사전 학습을 통해 vanilla Transformer가 최첨단 CNN보다 뛰어난 성능을 낼 수 있음을 보여주는 데 중점을 둡니다. 또한, Cordonnier et al. (2020)는 2x2 픽셀의 작은 패치 크기를 사용하여 저해상도 이미지에만 적용 가능한 반면 ViT는 중간 해상도 이미지에도 적용 가능합니다.

Chen et al. (2020a)는 이미지 해상도와 색상 공간을 줄인 후 이미지 픽셀에 Transformer를 적용한 image GPT (iGPT) 모델을 제안했습니다. iGPT는 생성 모델로서 비지도 학습 방식으로 훈련되었으며, ImageNet에서 최대 72%의 정확도를 달성했습니다.

Bello et al. (2019)는 이미지 분류를 위해 feature map을 향상시키는 데 self-attention을 사용했으며, Hu et al. (2018)과 Carion et al. (2020)은 object detection을 위해 CNN 출력을 self-attention으로 추가 처리했습니다. 이 외에도 video processing, image classification, unsupervised object discovery, unified text-vision tasks 등 다양한 분야에서 CNN과 self-attention을 결합하는 연구가 진행되었습니다.

기존에도 CV분야에 Transformer를 적용하려는 시도는 많았지만 표준 ImageNet 데이터셋보다 더 큰 크기의 데이터셋에서 image recognition 실험을 진행했다는 것이 핵심

Vision Transformer (ViT) 모델 구조

위에 있는 ViT의 구조를 다시 자세히 살펴보겠습니다. 기존 Transformer의 Encoder 부분을 활용해서 제작됐습니다.

이미지 패치 분할: ViT는 입력 이미지를 고정된 크기의 패치로 분할합니다.
선형 임베딩: 각 패치는 선형 투영(linear projection)을 통해 벡터 형태로 변환됩니다. 즉, 각 패치는 고차원 공간에 매핑됩니다.
위치 임베딩 추가: 각 패치의 위치 정보를 유지하기 위해 위치 임베딩(position embedding)을 추가합니다. Transformer는 순서 정보를 알 수 없기 때문에, 각 패치가 이미지에서 어떤 위치에 있었는지 알려주는 정보를 추가해야 합니다.
Transformer 인코더 입력: 패치의 선형 임베딩과 위치 임베딩을 결합하여 생성된 벡터 시퀀스를 표준 Transformer 인코더에 입력합니다.
CLS 토큰: 분류 작업을 수행하기 위해 학습 가능한 “CLS 토큰”을 시퀀스에 추가합니다. 이 토큰은 Transformer 인코더를 통과하면서 전체 이미지에 대한 정보를 취합하게 됩니다.
분류: Transformer 인코더 출력에서 분류 토큰에 해당하는 벡터를 사용하여 이미지를 분류합니다.

Transformer의 구조를 알고 있으면 어렵지 않게 이해할 수 있습니다.

METHOD

ViT는 자연어 처리에서 성공적인 성능을 보여준 Transformer 구조를 최대한 활용합니다. NLP 분야에서 사용되는 확장 가능한 Transformer 아키텍처와 효율적인 구현 방식을 거의 그대로 사용할 수 있다는 장점이 있습니다.

Eq. 1은 이미지 패치를 Transformer 입력 형태로 변환하고 위치 정보를 추가하는 과정을 나타냅니다.
Eq. 2, 3은 Transformer 인코더의 구조와 각 레이어의 동작을 정의합니다.(MSA는 multi-head self-attention)
Eq. 4는 Transformer 인코더 출력에서 이미지를 나타내는 특징 벡터를 추출하는 과정을 보여줍니다.
MLP에는 두번의 GELU activation function을 지나갑니다(transformer는 ReLU).

ViT는 대규모 데이터셋에서 pre-training된 후, 작은 규모의 데이터셋에 대해 특정 작업(downstream task)을 수행하도록 fine-tuning됩니다. pre-training된 예측 헤드를 제거하고, 새롭게 초기화된 분류 헤드를 추가하여 미세 조정합니다.

CNN의 Inductive Bias

CNN은 이미지 처리에 특화된 Inductive Bias를 갖고 있습니다.

Locality: Convolutional filter는 이미지의 local neighborhood, 즉 작은 영역만을 고려합니다. 이는 가까운 픽셀들이 서로 관련성이 높다는 가정에 기반합니다.
Two-dimensional neighborhood structure: CNN은 이미지의 2차원 구조를 유지합니다. Convolution 연산은 픽셀의 spatial relationship을 고려하여 이미지의 공간 정보를 효과적으로 추출합니다.
Translation equivariance: 이미지를 이동시켜도 CNN의 출력은 동일하게 유지됩니다. 즉, 객체의 위치가 바뀌어도 동일하게 인식할 수 있습니다.

ViT의 Inductive Bias

ViT는 CNN보다 이미지 특징에 대한 귀납적 편향이 훨씬 적습니다.

MLP 레이어: ViT의 MLP 레이어만 locality와 translation equivariance를 가집니다.
Self-attention 레이어: Self-attention 레이어는 이미지 전체 패치를 고려하기 때문에 global 정보를 학습합니다.
2차원 구조: ViT는 이미지 패치화 과정과 미세 조정 시 해상도 조정을 위한 위치 임베딩에서만 2차원 구조를 제한적으로 사용합니다.
Position embeddings: 초기 위치 임베딩은 패치의 2차원 위치 정보를 가지고 있지 않습니다. 따라서 패치 간의 spatial relationship은 학습을 통해 처음부터 새롭게 학습되어야 합니다.

Hybrid Architecture

CNN 특징 맵 활용: ViT는 입력으로 원본 이미지 패치 대신 CNN에서 추출된 특징 맵을 사용할 수 있습니다. 이러한 Hybrid Architecture는 CNN의 특징 추출 능력과 Transformer의 강력한 모델링 능력을 결합한 것입니다.
패치 임베딩: Hybrid model에서 패치 임베딩은 CNN 특징 맵에서 추출된 패치에 적용됩니다.
1x1 패치: 특별한 경우로, 패치 크기를 1x1로 설정할 수 있습니다. 이는 CNN 특징 맵의 spatial dimensions을 flatten하여 Transformer dimension으로 투영하는 것과 동일합니다.

Hybrid Architecture는 CNN의 장점을 활용하여 ViT의 성능을 향상시킬 수 있는 방법입니다. CNN은 이미지의 low-level 특징을 효과적으로 추출하고, Transformer는 global context를 고려하여 이미지를 분류하는 데 탁월합니다.

ViT의 Fine-tuning 및 고해상도 활용

ViT는 대규모 데이터셋에서 사전 훈련된 후, 작은 규모의 데이터셋에 대해 특정 작업(downstream task)을 수행하도록 미세 조정됩니다. 사전 훈련된 prediction head를 제거하고, 새롭게 초기화된 분류 헤드를 추가하여 미세 조정합니다.

fine-tuning 시 pre-training보다 높은 해상도를 사용하면 성능이 향상될 수 있습니다. 고해상도 이미지를 입력할 때 패치 크기는 동일하게 유지되므로, 입력 시퀀스 길이가 길어집니다. ViT는 메모리 제약 내에서 임의의 길이 시퀀스를 처리할 수 있지만, 사전 훈련된 위치 임베딩은 더 이상 의미가 없을 수 있습니다. 따라서 사전 훈련된 위치 임베딩을 원본 이미지의 위치에 따라 2D Interpolation을 수행합니다. 2D Interpolation은 알려진 데이터 지점을 기반으로 알려지지 않은 지점의 값을 추정하는 방법입니다.

EXPERIMENTS

ViT의 이미지 표현 학습 능력을 ResNet, Hybrid 모델과 비교 평가하기 위해 실험을 진행합니다. 다양한 데이터셋 크기에 따른 ViT, ResNet, Hybrid 모델의 성능을 비교하고, 각 모델의 데이터 요구량을 파악합니다. 또한, 사전 훈련 비용 대비 ViT의 성능 우수성을 입증하고, self-supervision 방식을 적용한 ViT의 가능성을 살펴봅니다.

데이터셋: 모델의 확장성을 평가하기 위해 다양한 크기의 데이터셋을 사용합니다.

ILSVRC-2012 ImageNet (ImageNet): 1,000개 클래스, 130만 장의 이미지
ImageNet-21k: 21,000개 클래스, 1,400만 장의 이미지
JFT (Sun et al., 2017): 18,000개 클래스, 3억 300만 장의 고해상도 이미지

데이터셋 중복 제거: Kolesnikov et al. (2020)의 연구를 따라 downstream task의 테스트 셋과 겹치는 사전 훈련 데이터셋을 제거합니다.

Downstream task: 사전 훈련된 모델을 다양한 벤치마크 작업에 전이하여 성능을 평가합니다.

ImageNet (original & ReaL labels): 원본 검증 레이블 및 정제된 ReaL 레이블 (Beyer et al., 2020) 사용
CIFAR-10/100: (Krizhevsky, 2009)
Oxford-IIIT Pets: (Parkhi et al., 2012)
Oxford Flowers-102: (Nilsback & Zisserman, 2008)

BERT에서 사용된 설정을 기반으로 한 ViT 모델 변형을 보여줍니다. “Base”, “Large”, “Huge” 모델을 사용하며, 입력 패치 크기를 나타내는 표기법을 사용합니다 (예: ViT-L/16).

COMPARISON TO STATE OF THE ART

ViT 모델을 다른 SOTA 이미지 분류 모델들과 비교합니다.

JFT-300M 데이터셋으로 사전 훈련된 ViT 모델은 모든 데이터셋에서 ResNet 기반 모델보다 높은 성능을 보여줍니다.
ImageNet-21k 데이터셋으로 사전 훈련된 ViT 모델도 좋은 성능을 달성합니다.
ViT 모델은 ResNet 모델보다 사전 훈련에 필요한 계산 리소스가 훨씬 적습니다.

ViT는 이미지 분류 작업에서 가장 좋은 성능을 달성하며, ResNet 기반 모델보다 적은 계산 리소스로 학습 가능합니다. ViT는 대규모 데이터셋으로 사전 훈련될 때 더욱 뛰어난 성능을 보입니다.

ViT가 기존 모델 대비 성능과 효율을 모두 가지는 것을 확인할 수 있습니다.

PRE-TRAINING DATA REQUIREMENTS

ViT 모델의 성능에 대한 사전 훈련 데이터셋 크기의 영향을 분석하고, ViT와 ResNet 모델의 데이터 효율성을 비교합니다. ResNet보다 이미지에 대한 Inductive Bias가 적은 ViT 모델에게 데이터셋 크기가 얼마나 중요한지 확인합니다.

욎쪽 그래프는 ViT 모델이 큰 데이터셋으로 사전 훈련될 때 ResNet보다 성능이 뛰어나다는 것을 보여줍니다. 또한 오른쪽 그래프는 ViT 모델이 작은 데이터셋에서 ResNet보다 더 쉽게 과적합되지만, 큰 데이터셋에서는 더 좋은 성능을 보인다는 것을 보여줍니다.

ImageNet (가장 작은 데이터셋)으로 사전 훈련된 경우, ViT-Large 모델은 ViT-Base 모델보다 성능이 낮습니다. JFT-300M (가장 큰 데이터셋)으로 사전 훈련해야 ViT-Large 모델의 장점이 드러납니다.

따라서 ViT는 데이터 효율성이 낮아 작은 데이터셋에서는 과적합될 수 있지만, 충분한 데이터가 주어진다면 더 우수한 성능을 발휘할 수 있습니다.

ViT는 동일한 계산량을 사용했을 때 일반적으로 ResNet보다 높은 성능을 달성합니다(계산 효율성이 좋다).
Hybrid 모델은 작은 모델 크기에서 유용할 수 있지만, ViT의 확장성이 더 뛰어나므로 큰 모델에서는 ViT가 더 효율적입니다.

또한 FT-300M 데이터셋을 사용하여 다양한 모델의 확장성(Scaling)을 비교 분석합니다. 데이터셋 크기가 모델 성능의 병목 현상이 되지 않는 환경에서, 각 모델의 사전 훈련 비용 대비 성능을 평가했습니다.

ViT는 ResNet보다 성능/계산량 trade-off 측면에서 우수합니다. ViT는 동일한 성능을 달성하기 위해 ResNet보다 약 2–4배 적은 계산량을 사용합니다. 실험 범위 내에서 ViT는 성능 포화 상태에 도달하지 않았으며, 이는 향후 더 큰 모델을 탐색할 필요가 있음을 의미합니다.

ViT는 높은 성능과 계산 효율성을 동시에 달성하는 확장 가능한 모델입니다.

INSPECTING VISION TRANSFORMER

ViT 모델이 이미지 데이터를 어떻게 처리하는지 이해하기 위해 내부 표현을 분석합니다.

필터 시각화

ViT의 첫 번째 레이어는 패치를 저차원 공간에 선형 투영합니다. 왼쪽은 학습된 임베딩 필터의 주요 구성 요소를 보여줍니다. 이 경우에 잘 학습된 CNN의 앞쪽 레이어를 시각화했을 때와 유사한 결과가 나타납니다 .즉, 잘 학습된 CNN과 같이 이미지 인식에 필요한 edge, color 등의 low-level 특징들을 잘 포착하고 있음을 파악할 수 있습니다.

위치 임베딩 유사도 분석

위 그림은 모델이 이미지 내 거리를 위치 임베딩의 유사성으로 인코딩하는 법을 학습한다는 것을 보여줍니다. 즉, 더 가까운 패치는 더 유사한 위치 임베딩을 갖는 경향이 있습니다. 또한, 행-열 구조가 나타납니다. 같은 행/열에 있는 패치는 유사한 임베딩을 가집니다. 마지막으로, 더 큰 그리드의 경우 sin파 구조가 나타나기도 합니다.

Attention Distance

attention 가중치를 기반으로 정보가 통합되는 이미지 공간의 평균 거리를 계산합니다. 이 “attention 거리”는 CNN의 수용 영역 크기와 유사합니다. 일부 헤드는 가장 낮은 레이어에서 이미 이미지의 대부분에 주목하며, 정보를 전역적으로 통합하는 기능이 실제로 모델에서 사용됨을 보여줍니다. 다른 attention 헤드는 낮은 레이어에서 지속적으로 작은 attention 거리를 가집니다. 이러한 매우 국소적인 attention은 Transformer 이전에 ResNet을 적용하는 hybrid 모델에서 덜 두드러집니다.

결론적으로,

ViT는 패치의 미세 구조를 효과적으로 나타내는 임베딩 필터를 학습합니다.
ViT는 위치 임베딩을 통해 이미지 내 거리와 공간적 관계를 효과적으로 인코딩합니다.
ViT는 self-attention 메커니즘을 통해 이미지 전체의 정보를 통합하고, 의미적으로 관련된 영역에 주목합니다.

SELF-SUPERVISION

레이블 없이 데이터 자체에서 학습하는 방식. NLP 분야에서 Transformer의 성공에 큰 역할을 했습니다. BERT에서 사용된 Masked Language Modeling 방식을 모방하여 ViT를 자기 지도 학습합니다. 이미지의 일부 패치를 가리고, 모델이 가려진 패치를 예측하도록 훈련합니다. (마스킹 비율: 50%)

self-supervised pre-training으로 사전 훈련된 ViT-B/16 모델은 ImageNet에서 79.9%의 정확도를 달성했습니다.
이는 처음부터 학습하는 것보다 2% 향상된 결과이지만, 지도 학습(Supervised Learning)으로 사전 훈련된 모델보다는 4% 낮은 수치입니다.

저자들은 self-supervised pre-training은 ViT 모델의 성능을 향상시킬 수 있는 잠재력을 보여주지만 지도 학습 방식에 비해 성능이 낮으며, 더 많은 연구가 필요하다고 주장합니다.

CONCLUSION

핵심 내용

기존 연구들과 달리, 초기 패치 추출 단계를 제외하고는 아키텍처에 이미지 특화 편향을 도입하지 않았습니다. 이미지를 패치 시퀀스로 해석하고, NLP에서 사용되는 표준 Transformer 인코더로 처리했습니다. 이러한 단순하지만 확장 가능한 전략은 대규모 데이터셋으로 사전 훈련될 때 놀라울 정도로 효과적입니다. ViT는 많은 이미지 분류 데이터셋에서 최고의 성능을 달성하거나 능가하며, 사전 훈련 비용이 상대적으로 저렴합니다.

향후 과제

ViT를 객체 탐지, 이미지 분할과 같은 다른 컴퓨터 비전 작업에 적용해야 합니다. Carion et al. (2020)의 연구 결과와 함께 이 연구는 이러한 접근 방식의 가능성을 보여줍니다. 초기 실험에서 self-supervised pre-training으로 사전 훈련된 모델의 성능 향상을 확인했지만, self-supervised pre-training과 대규모 지도 학습 사이에는 여전히 큰 차이가 있습니다. ViT를 더욱 확장하면 성능이 향상될 가능성이 있습니다.