ViT

Abstract

NLP 분야에선 Transformer 구조가 기준이 되어가고 있지만, Vision 분야에서는 적용이 미약한 편이다. Vision 분야에서는 아직까지 convolution network 구조에 의존하는 추세이고, 본 논문에서는 이러한 cnn 구조에 대한 의존을 줄이고, 이미지 시퀀스를 트랜스포머 구조에 적용하여 classification task에 적용하여 매우 좋은 성능을 보인다. Large 데이터 셋에 대해서 Pretrain한 뒤에 Mid나 Small 데이터 셋에 대해서 transfer learning을 적용하여 적은 컴퓨팅 자원을 사용하면서도 cnn에 비해 sota에서 훌륭한 성능을 보였다.

1. Introduction

NLP 분야에서 self-attention 기반 transformer 구조가 주로 사용되고 있다. 주로 큰 텍스트 데이터들에 대해서 학습하고, 더 작은 task 데이터 셋에 대해서 fine tunning을 적용한다. Transformer의 능률과 확장성 덕분에 과적합의 문제를 피하면서 더 큰 모델과 데이터 셋을 사용할 수 있게 되었다.

하지만, computer vision 분야에서는 NLP와 맞물려서 일부 convolution network를 self-attention 구조로 대체하는 연구가 진행되고 있지만 최근의 모델들은 이론적으로 효율적이지만 특정한 attention 패턴을 사용하기 때문에 아직 하드웨어 가속기에 대해서 효과적으로 확장되지 못하고 있다.

본 논문에서는 NLP 분야의 Transformer에 영향을 받아 이미지에 대해서 가능한 적게 수정하여 이미지를 Transformer에 적용하는 시도를 진행했다. 그래서 본 논문에서는 이미지를 패치 단위로 잘라서 이 패치들의 선형 embedding 시퀀스를 이미지의 input으로 사용했다.

strong regularization없이 ImageNet과 같은 mid 사이즈의 데이터 셋으로부터 학습했을 때, 제안하는 모델은 ResNet보다 몇 퍼센트 낮은 정확도를 보였다. 이것은 다소 실망스러운 결과를 보일 수 있지만, Transformer는 cnn 구조의 locality나 translation equivariance와 같은 inductive bias가 없으므로 불충분한 데이터로 학습했을 때 일반화를 잘 하지 못하게 되는 것이다.

그러나 매우 큰 데이터 셋으로부터 학습했을 때 본 논문의 저자는 inductive bias를 능가할 수 있다고 주장했다. ImageNet-21K나 JFT-300M에서 학습했을 때, ViT는 최신 여러 모델을 능가하는 성능을 보였다. 특히, ImageNet에 대해서 88.55%, ImageNet-Real에 대해서 90.72%, CIFAR-100에 대해서 94.55%, 19개 task의 VTAB suite에 대해서 77.63%를 달성했다.

2. Realated Work

Vision 분야에서 Naive한 self-attention의 응용은 서로 다른 픽셀에 대해서 attention해야하기 때문에 2차 비용함수 문제로 이어져서 실제 input 크기를 사용하지 않는다. Parmar는 각 픽셀의 이웃 픽셀들에 대해서만 self-attention을 적용했다. 다른 연구로는 Child가 주장한 Sparse Transformer의 이미지에 대한 global self-attention을 통한 확장 가능한 근사화가 있다. 이에 대한 대안 방법중 하나는 다양한 크기의 블록에 scale attention을 적용하는 것도 있다. 이러한 방법들은 이론적으로 유리하지만, 하드웨어 가속기에 대해서 효율적으로 적용하려면 복잡한 로직이 필요하게 되는 문제가 있다.

본 논문은 입력 이미지로부터 2x2 patch를 추출하여 전체적으로 self-attention을 적용한 Cordonnier의 모델에 가장 영향을 받았다. 이 모델은 ViT와 매우 유사하지만, 본 논문에서 제안하는 모델은 large scale pretrain을 장려했고, SOTA의 cnn 모델들에 대해서 경쟁력을 얻을 수 있었다. 그리고 Cordonnier가 제안한 모델은 2x2 픽셀의 작은 패치 사이즈를 사용하여 작은 이미지에만 적용가능한 반면에 제안하는 모델은 middle 이미지에도 적용 가능하다.

또한 최근에는 convolution network와 transformer를 함께 사용하는 방법에 대한 연구도 이뤄지고 있다. 한편, 관련 연구로 image GPT 모델이 있으며, 이 모델은 image resolution과 color space를 줄인뒤에 transformer를 적용했다.

본 논문의 연구는 표준 ImageNet의 데이터 셋보다 큰 규모의 데이터 셋에서 이미지 인식을 수행하는 연구로써, convolution network 기반의 기존 연구가 아닌 transformer를 기반으로 Imagenet-21K와 JFT-300M과 같은 큰 규모의 데이터 셋에 대해서 실증적 탐구를 진행한다.

3. Method

본 논문은 최대한 NLP에서의 Transformer의 구조를 그대로 사용함으로써, 확장 가능한 Transformer의 구조와 효율적인 구현이 즉시 가능하도록 하였다.

3.1 Vision Transformer (VIT)

제안하는 모델의 도식화는 위와 같다. 표준 Transformer는 토큰 임베딩된 1D 시퀀스를 입력으로 받는다. 이를 위해 HXWXC 차원의 2D 이미지를 NX(P2xC) shape의 flattened 2D 패치들의 시퀀스로 reshape했다. 여기서, C는 채널의 수, P는 각 이미지 패치의 resolution을 의미한다. 그리고 N = HW/P2와 같다. Transformer는 D차원의 일정한 잠재 벡터를 모든 레이어에 걸쳐 사용하기 때문에, 본 논문의 저자는 패치를 flatten한 뒤 훈련가능한 linear projection을 통해 D 차원으로 매핑한다.

BERT의 class token과 유사하게 여기서도 학습 가능한 Embedding을 임베딩된 패치들의 시퀀스 앞에 붙인다. 이것의 Transformer Encoder 출력 state는 image representation y를 의미한다. pre-training이나 fine-tuning시에 classification head는 z0L에 붙는다. classification head는 pre-training시에는 하나의 hidden layer와 MLP에 의해 구현되고, fine-tuning시에는 single 선형 layer에 의해 구현된다.

positional 정보를 유지하기 위해 임베딩된 패치에 position embedding 값이 더해진다. 본 저자는 표준 학습 가능한 1D position embedding들을 사용했다. 왜냐하면, 2D에 적응적인 position embedding에서 놀랄만한 성능 향상을 얻지 못했기 때문이다.

Transformer Encoder는 MSA와 MLP 블록으로 구성되어 있고, LayerNorm(LN)은 매 블록을 거치기 이전에 적용되고, residual connection이 모든 블록 이후에 적용된다.

Inductive bias

Vit는 기존 CNN 모델들에 비해 image-specific inductive bias가 훨씬 적다. CNN에서는 2차원적인 이웃 구조와 같은 locality와 translation 등가성이 전체 모델의 각 layer에 보장된다. Vit에서는 오직 MLP 레이어에서만 지역성과 translation 등가성이 적용되고, self-attention 레이어들은 전역적이다. 2차원적 이웃 구조는 매우 드물게 사용되는데, 모델을 시작할때 이미지를 패치형태로 자르고 fine-tuning시에 서로 다른 이미지의 position embedding을 조절할때 사용된다. 그 외에는 패치들의 2d position 정보는 사용되지 않고, 패치간의 공간적인 관계는 학습되어야만 한다.

Hybrid Architecture

raw 이미지 패치에 대한 대안적인 방법으로써, input 시퀀스는 feature map의 형태로 구성될 수 있다. 이와 같은 hybrid model에서는 patch embedding projection E가 CNN featuure map에서 추출된 패치에 적용된다.

3.2 Fine-Tuning and Higher Resolution

본 논문에서는 ViT를 large 데이터 셋에 대해서 pretrain을 진행한 뒤, downstream task로 fine-tune을 한다. downstream시에는 pretrain된 prediction head를 제거하고, zero 초기화된 DxK 차원의 순방향 레이어를 붙인다. 여기서, K는 downstream class 수를 의미한다. 그리고 fine-tune시에는 pretrain된 resolution보다 더 높은 resolution을 사용할 때 더 유리하다고 주장한다. 다만, 더 높은 resolution에 대해서 fine-tune시에는 동일한 patch size를 사용하기 때문에 더 많은 sequence 길이가 생성되고, 이에 따라 position embedding이 의미가 없어지므로 pretrained position embedding에 대해서 원본 이미지만큼 2d interpolation을 적용한다.

4. Experiments

본 논문은 ResNet, Vision Transformer, Hybird 모델에 대해서 테스트를 했다. 요구되는 데이터량을 파악하기 위해 다양한 사이즈의 데이터 셋에 대해서 다양한 테스크 환경에서 테스트를 진행했다. ViT는 적은 pretrain cost를 가지고 sota의 recognition benchmark에서 순조롭게 기록을 낼 수 있었다. 추가적으로 최근에는 self-supervision과 관련해서 작은 실험도 진행하고 있다.

4.1 Setup

Datasets

모델 확장성을 실험하기 위해, 1k classes와 1.3M images를 가진 ILSVRC-2012 ImageNet 데이터 셋과 21k classes와 14M images를 가진 ImageNet-21k 데이터 셋과 18k classes와 303M high-resolution images를 가진 JFT 데이터 셋을 사용했다. 추가적으로 downstream task들의 test set에 대해서 중복 제거를 수행했다. 본 논문의 저자는 이 데이터 셋에 대해서 학습한 모델들을 대상으로 몇개의 benchmark task에 대해서 transfer 하였다. ImageNet, ReaL labels, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102가 이에 해당된다. 이 데이터 셋들에 대해서는 Kolesnikov의 전처리 방식을 따랐다. 그리고 추가적으로 19-task VTAB classification suite에 대해서 평가했다. VTAB은 적은 데이터의 task당 1000개의 학습 예제를 가진 다양한 task로의 transfer을 평가한다. 여기서, task는 3가지 그룹으로 나뉜다.

Natural - Pets, CIFAR, etc.와 같은 task
Specialized - medical, satelite와 같은 task
Structured - localization과 같이 기하학적 이해가 필요한 task

Model Varaints

여기서 ViT-L/16은 16x16 input patch 사이즈를 의미하고, 주의할 점은 patch size의 크기가 작을수록 더 많은 연산 cost가 필요하다는 점이다.

실험에 사용된 ResNet에는 Batch Noramlization 대신 Group Normalization과 standardized convolution을 사용했고, 이러한 시도는 transfer를 향상 시켰고 본 논문에서는 이 모델을 ResNet(BiT)로 명시했다. Hybrid 모델의 경우에는 feature map을 그대로 ViT 모델의 input에 사용했으며, 이 때 patch size는 1로 사용했다. sequence 길이에 따른 실험을 하기 위해 (i)ResNet50의 stage 4의 출력을 사용하거나, (ii)stage 4를 제거하고 stage 3와 같은 layer의 수(전체 layer의 수를 유지하기 위해)를 사용하여 이 확장된 stage 3의 출력을 사용했다. 여기서 옵션 (ii)의 결과는 4배 더 킨 시퀀스를 생성하여 ViT 모델이 더 비싸졌다.

Taining & Fine-Tuning

전체 모델들을 학습시에는 Adam Optimizer를 사용했다. commomn practice 대상으로는 ResNet 모델에 대해서 Adam이 SGD보다 더 유용하다고 주장한다. 그리고 fine-tuning 시에는 Momentum SGD를 사용했고 ImageNet의 결과는 higher resolution으로 fine tuning하여 결과를 얻었다.

Metrics

few-shot 또는 fine-tuning 정확도를 통해 downstream dataset의 결과를 명시했다. fine-tuning 정확도는 각각의 데이터 셋에 대해서 fine-tuning한 후의 각 모델의 성능을 통해 구했고, few-shot 정확도는 훈련 이미지 하위 집합의 표현(고정)을 [-1, 1]K 타겟 벡터에 매핑하는 정규화된 최소 제곱 회귀 문제를 해결하여 얻었다. fine-tuning의 비용이 매우 큰 경우 few-shot은 빠르게 결과를 얻을 수 있으므로 본 저자는 이런 경우 linear few-shot 정확도를 사용한다고 했다.

4.2 Comparison to State of the Art

JFT와 ImageNet21K로 학습된 ViT-H/L 모델과 기존 BiT-L(ResNet152x4), Noisy Student(EfficientNet-L2)의 성능을 비교한 결과 더 적은 컴퓨팅 자원으로 더 높은 성능을 얻을 수 있었다.

ViT-L/16 모델의 경우 ImageNet-21k에서 pretrain시에 TPUv3 8core의 자원을 사용했을 때 약 30일 정도 학습 시간이 소요되었다고 한다.

VIVI : ImageNet과 Youtube에서 학습된 ResNet 모델
S4L : ImageNet에서 supervised plus semi-supervised learning 모델

4.3 Pretraining Data Requirements - 4.4 Scaling Study

ViT 모델에 대해서 ImageNet, ImageNet-21k, JFT-300M 순서로 pretrain 진행했다. 더 적은 데이터 셋에 대해서 성능을 높이기 위해 weight decay, dropout, label smoothing 파라미터를 최적화 했다.

JFT-300M 데이터 셋의 경우는 9M, 30M, 90M 순서로 random sampling을 통해 진행했다.

여기서 주목할만한 점은 BiT의 경우 더 적은 데이터 셋에 대해서는 ViT보다 더 높은 성능을 보이지만, 데이터 셋이 더 커짐에 따라 ViT 성능이 더 높아지게 된다는 점이다.

또한, Hybrid 모델을 사용했을 때, 비슷한 컴퓨팅 cost에 비해 더 높은 성능을 얻을 수 있음을 확인할 수 있다.

4.5 Inspecting Vision Transformer

왼쪽은 초기 선형 임베딩 필터의 RGB 값들인데, CNN의 convolution filter와 유사한 형태를 보임을 알 수 있다. 그리고 Position Embedding의 경우에도 학습함에 따라 각 패치의 위치에 따라 가중치들이 학습됨을 알 수 있다. 마지막으로 ViT 모델 네트워크의 깊이가 깊어짐에 따라서 attention distance가 줄어들은것을 확인할 수 있다. 여기서 attention distance는 cnn의 receptive field 크기와 유사하다. 초기 네트워크 레이어는 이미지를 전역적으로 보고, 네트워크의 depth가 깊어짐에 따라 이미지 분류에 관계있는 영역에 집중하는 것을 알 수 있다.

4.6 Self Supervision

BERT의 masked language modeling task처럼 self supervision을 통해 masked patch 예측하는 실험도 진행했다. ImageNet에서 약 79.9%의 성능을 얻었지만, 아직 부족한 결과를 보인다.

5. Conclusion

• NLP의 단순한 Transformer Encoding 구조를 이미지에 적용하여 높은 확장성을 얻으면서도 적은 train cost를 소요할 수 있게 되었다. ViT의 확장성이 더 높은 성능을 이끌 수 있을것이다.

• Challenge
  1) 다른 Vision Task로의 적용
  2) Self Supervision 성능 고도화