An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale(ViT)(ICLR 2021)
Abstract
- `Transformer` 구조가 NLP task에서 사실상 기준이 되는 동안 Computer Vision에서의 응용은 제한적
- Vison에서 `Attention`은 Convolutional networks와 함께 적용하거나 전체 구조를 그대로 유지하면서 Convolution Network의 특정 구성 요소를 대체하는 데 사용
- CNN에 의존할 필요가 없으며 Image patches의 sequences에 직접 적용된 pure transformer가 Image classification task에서 매우 우수한 성능을 보임
- 대량의 데이터에 대해 사전 훈련되고 중형 또는 소형의 Image recognition benchmarks(ImageNet, CIFAR-100, VTAB 등)에 transfer될 때 `Vision Transformer(ViT)`는 다른 SOTA CNN 모델들과 비교했을 때 훈련에 훨씬 적은 컴퓨팅 리소스를 필요로 하면서도 우수한 결과를 얻음
Introduction
`Self-Attention`에 기반한 architecture들, 특히 `Transformers`는 NLP에서 자주 선택되는 모델이 되었다. 이에 대한 주요한 접근 방법은 거대한 text corpus로 pre-train한 후 작은 task들에 대해서는 fine-tune하는 것이다. 또한 Transformer의 계산적인 효율성과 확장성 덕분에 전례 없는 크기의 100B가 넘는 파라미터들에 대한 학습이 가능해졌다.
그러나 Vision 분야에선 여전히 CNN 기반의 모델이 지배적이다. 최근에는 NLP의 성공에 영감을 받아 CNN과 유사한 architecture와 Self-Attention을 결합하는 시도들이 진행되었으며 이는 일부 Convolution Networks를 대체하거나, 이를 완전히 대체하기기도 하였다. 후자의 경우 이론적으로 효율적이지만, 특수한 Attention patterns을 사용하기 때문에 최신 하드웨어 가속기에는 아직 효과적으로 확장되지 못하였다. 따라서 대규모 Image Recognition 분야에서는 여전히 고전적인 ResNetlike architecture들이 SOTA 모델이다.
이미지를 `patch`로 분할하고, 이러한 패치들의 선형 embedding sequence를 Transformer의 입력으로 제공해 최소한의 수정만으로 이미지에 직접 표준 Transformer를 적용하는 실험을 진행하였다. 이미지의 patch는 NLP에서의 tokens(words)와 동일한 방식으로 취급하며, 지도 학습 방식으로 Image classification에 대해 모델을 학습시켰다.
이 과정에서 강력한 정규화 없이 ImageNet과 같은 중간 크기의 데이터셋에 대해 학습하는 경우 `ResNet`보다 낮은 성능을 보였는데, 이는 CNN의 inductive biases가 부족하여 일반화가 잘 이루어지지 않아 어쩌면 당연한 결과이다. 그러나 대규모 모델(1400만~3억 이미지)에 대한 학습의 경우 inductive biases보다 더 효과적이었다. 따라서 ViT는 충분한 규모로 사전 학습을 한 후 데이터 수가 적은 작업에 적용하면 우수한 결과를 얻을 수 있다.
Method
Vision Transformer(ViT)
전체적인 모델의 구조는 먼저 이미지를 고정된 크기의 patch들로 분할하고, `linearly embedding`을 수행한 후 `position embedding`을 추가하여 결과 sequence를 standard Transformer Encoder의 입력 값으로 활용한다. 또한 classification을 수행하기 위해 학습 가능한 `classification token`을 sequence에 추가한다. 이러한 ViT의 구조를 순서대로 나타내면 다음과 같다.
- 이미지(HxWxC)를 N(=HW/P^2)개의 패치(PxP)로 분할하여 patch sequence(Nx(PxPxC)) 구축
- Linear Projection을 통해 각 패치를 flatten한 벡터를 D차원으로 변환 → Patch Embedding
- Class Embedding과 Patch Embedding에 Position Embedding을 더함
- Transformer Encoder의 input으로 활용해 class embedding에 대한 output인 image representation 도출
- MLP에 image representation을 input으로 넣어 이미지의 class를 분류
- standard Transformer는 `token embeddings`의 1D sequence를 입력받기 때문에, 2D images를 다루기 위해 image에 대해 H x W x C 크기의 이미지를 N x (P x P x C) 크기의 patch로 변환하여 reshape을 수행
- H : 이미지의 높이
- W : 이미지의 넓이
- C : 이미지의 채널
- N(=HW/P^2) : patch의 개수
- P x P : 고정된 patch의 크기
- Transformer는 모든 layer에서 일정한 잠재 벡터 크기 D를 사용하기 때문에 patch를 flatten하고 훈련 가능한 linear projection을 수행해 D차원에 매핑하여 D차원으로 변환 → 이러한 projection의 output을 `patch embedding`이라고 표현
- `Positional Embedding`은 위치 정보를 유지하기 위해 patch embedding에 추가
- 2D-aware position embeddings을 사용했을 때 눈에 띄는 성능 향상이 없었기 때문에 1D position embeddings 사용
- `BERT`의 [class] token과 유사하게 학습 가능한 embedding을 embeded patch들의 sequence에 추가
- Transformer Encoder는 `Multi-Head Self-Attention(MSA)`와 `MLP` Block으로 구성
- MLP는 2개의 Layer로 구성되며 `GELU` activation function 활
- `LayerNorm(LN)`이 모든 block 이전에 적용되며, `Residual Connection(Skip Connection)`이 모든 block 뒤에 적용
- embeded patch들의 상태는 Transformer Encder의 출력에서 `image representation` y로 사용
- 마지막으로 classification head가 pre-training 및 fine-tuning 과정에서 모두 Transformer Encoder의 출력 값에 부착
- classification head는 pre-training 시 하나의 hidden layer와 함께 MLP로 구현
- fine-tuning 시 하나의 linear layer로 구현
Inductive Bias(귀납적 편향)
- training에서 보지 못한 데이터에 대해서도 적절한 귀납적 추론이 가능하도록 하기 위해 모델이 가지고 있는 가정들의 집합
- DNN의 기본적인 요소들의 inductive bias
- `Fully connected` : 입력 및 출력이 모두 연결되어 있어 특별한 relational inductive bias 가정 x
- `Convolutional` : CNN은 작은 크기의 동일한 kernel을 통해 이미지를 지역적으로 보기 때문에 `locality`와 `transitional invariance` 특성을 가짐
- `Recurrent` : RNN은 입력 데이터들이 시간적 특성을 가지고 있다 가정하므로 `sequentiality`와 `temporal invariance` 특성을 가짐
- Transformer는 CNN 및 RNN보다 상대적으로 inductive bias가 낮음
- MLP layer에서만 local 및 translationally equivalent를 가지며 Self-Attention layer는 global layer
- 2차원 neighborhood structure는 모델을 시작할 때 이미지를 patch로 자르고 해상도가 다른 이미지의 position embedding들을 조정하기 위한 fine-tuning 시에만 매우 드물게 사용
- 그 외에는 초기화 시점의 position embedding들에는 patch의 2D 위치에 대한 정보가 없으며 patch 간의 모든 공간 관계를 처음부터 학습해야함
Hybrid Architecture
- image patch에 대한 대안으로 CNN의 Feature map을 활용하여 hybrid architecture를 구성 가능
- Feature map은 이미 raw image의 공간적 정보를 포함하고 있기 때문에 hybrid architecture는 패치 크기를 1x1로 설정 가능
- 1x1 크기 패치를 사용할 경우 feature map의 공간 차원을 flatten하여 각 벡터에 linear projection을 적용
Experiments
Datasets
- ViT는 다음과 같은 3개의 데이터셋을 기반으로 pre-train을 수행
- ImageNet-1k(1k classes, 1.3M images)
- ImageNet-21k(21k classes, 14M images)
- JFT(18k classes, 303M images)
- 다음과 같은 benchmark tasks를 downstream task로 하여 pre-trained ViT의 representation 성능을 검증
- ReaL labels, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102, 19-task VTAB classification suite
Model Variants
- ViT는 총 3개의 volume 및 다양한 패치 크기에 대해 실험을 진행
- ViT-L/16 : ViT Large 모델 및 16x16 크기의 패치
- Baseline CNN의 경우 `ResNet(BiT)` 사용
- `Batch Normalization` layer를 `Group Normalization`으로 변경
- standardized convolution을 사용하여 transfer learning에 더 적합한 Big Transformer(BiT) ResNet 사용
Comparison to State-Of-The-Art
- ViT-H/14, ViT-L/16의 성능을 최신 CNN 모델들과 비교
- 대규모 ResNet으로 지도 전이 학습을 수행하는 `BiT-L`
- 라벨이 제거된 ImageNet과 JFT-300M에서 Semi-Supervised Learning을 사용하여 훈련한 대규모 EfficientNet인 `Student Noisy`
- Noisy Student는 ImageNet에서, BiT-L은 해당 논문에서 활용하는 다른 데이터셋에서 SOTA 모델
- JFT에서 pre-training한 ViT-L/16 모델이 모든 downstream task에서 BiT-L보다 나은 성능을 보임
- ViT-H/14 모델은 ViT-L/16 모델보다 향상된 성능을 보였으며, BiT-L보다 학습 속도가 훨씬 빠름
- 19-task VTAB classification suite를 3개의 그룹으로 나눠 추가 실험 진행
- Natural : Pets, CIFAR, etc
- Specialized : medical, satellite imagery
- Structured : tasks that require geometric understanding like localization
- 전체 데이터뿐만 아니라 각 그룹에서도 ViT-H/14가 좋은 결과를 도출
Pre-training Data Requirements
ViT는 대규모 JFT-300M 데이터셋에서 사전 학습된 경우 우수한 성능을 발휘한다. ResNet보다 vision에 대한 inductive bais가 적은데, 데이터 세트의 크기가 얼마나 중요할지 2가지의 실험을 수행했다.
- pre-training 데이터셋의 크기에 따른 fine-tuning 성능 확인(ImageNet, ImageNet-21k, JFT-300M)
- 작은 데이터 세트에서의 성능 향상을 위해 최적화 진행 → weight decay, dropout, label smoothing
- 가장 작은 데이터 세트인 ImageNet으로 pre-trained을 수행한 결과 BiT CNN이 ViT보다 성능이 앞서며, 정규화를 수행하였음에도 불구하고 ViT Large 모델이 ViT Base 모델보다 성능이 떨어짐
- 그러나 데이터 세트의 크기가 커질수록 ViT 모델이 BiT보다 성능이 앞서며, 더 큰 ViT 모델이 효과가 있음
- 9M, 30M, 90M개의 랜덤한 JFT-300M 데이터 세트의 부분 데이터 세트와 전체 데이터 세트에 대해 모델 학습을 진행
- 더 작은 부분 데이터 세트에 대해서는 추가적인 정규화를 진행하지 않고, 하이퍼파라미터는 동일하게 설정하여 정규화가 아닌 모델 고유의 속성을 평가
- 그러나 `early-stopping`은 사용하여 훈련 중에 달성한 최고의 검증 정확도를 확인
- 실험 결과를 통해 CNN의 inductive bias가 소규모 데이터 세트에서는 유용하지만, 대규모 데이터 세트에서는 데이터로부터 패턴을 학습하는 것으로도 충분함
Scaling Study
- JFT-300M 데이터 세트를 통해 pre-training cost 대비 transfer 성능을 확인
- pre-training cost란 TPUv3 accelerator에서 모델의 inference 속도를 비교한 것
- ViT는 성능과 cost의 trade-off에서 ResNet을 압도
- ViT는 동일한 성능을 달성하기 위해 약 2~4배 적은 컴퓨팅 파워를 사용
- cost가 증가할수록 ViT와 hybrid의 성능과 cost의 trade-off 차이는 감소
- hybrid는 CNN의 feature map을 활용한 ViT 모델
Conclusion
- Computer Vision에서 `Self-Attention`을 사용하는 이전 작업과 달리 초기 patch extraction step을 제외하고는 이미지별로 inductive biases를 architecture에 도입하지 않음
- 대신, 이미지를 일련의 `patch sequence`로 해석하고, NLP에 사용되는 표준 Transformer Encoder로 처리
- 초기 결과는 고무적이지만 아직 많은 과제가 남아 있으며, 그 중 하나는 `Detection` 혹은 `Segmentation`과 같은 Computer Vision task에 ViT를 적용하는 것
- 또한 또다른 과제는 `Self-supervised pre-training` 방법을 탐구하는 것
- 초기 실험에서는 Self-supervised pre-training보다 개선된 결과를 보였지만, 여전히 Self-supervised와 large-scale의 `Supervised pre-training` 방법에는 큰 격차가 존재
Reference
[2] ViT 이미지
