- NLP분야에서 transformer가 사실상 standard하였지만, vision에서는 매우 드물게 사용되었다. Vision에서는 attention 사용 시, CNN과 혼합해 사용하거나 전체구조에서 몇개의 구성요소만 대체되는 방식으로 사용되어 왔다. 본 연구는 CNN에 의존하는 것이 불필요하며 오직 transformer만으로 image patch들의 sequence를 적용해 image classification task에 매우 잘 동작함을 실험적으로 증명하였다. 많은 dataset을 pre-train한 후 small~midsize의 이미지 인식 bench mark dataset에 전이학습 시, 여타 CNN보다 ViT가 요구계산량은 적으면서도 S.O.T.A를 보여준다.
1. 서론 (Introduction)
-"Self-Attention"기반 구조 특히나 Transformer[Vaswani2017]는 NLP task에서 널리 사용되며 가장 지배적이고 널리 사용되는 방식은 매우 긴 text에 대한 pre-train 이후 작고 구체적인 task의 dataset에 fine-tuning하는 것이다. transformer의 특징 중 하나인 계산의 효율성과 확장성 덕분에 전례없는 크기(100억개 이상의 parameter)의 model을 학습할 수 있다는 점이다. (model과 dataset 크기의 증가에도 성능이 saturating되지 않음)
- 다만 computer vision에서는 CNN방식이 지배적으로 남아있어서 NLP의 성공에 영감을 얻어 CNN에 "Self-Attention"을 합치려는 실험, Conv.layer를 모두 self-attention으로 대체하는 등이 이루어졌다. 후자의 경우, 이론적으로는 효율적이지만 attention의 독특한 계산메커니즘으로 효율적사용이 어렵다. (∵ image 전체를 한번에 넣는 연산 >>> 단어 vector에 대한 연산) 그렇기에 ResNet같은 모델들이 Large-Scale image에서 S.O.T.A로 남아있다.
- 본 연구는 NLP성공에 영감을 받아 standard transformer architecture의 구조를 조금 수정 후, image에 직접 적용하는 실험을 진행하였다. 이를 위해 image를 patch로 나누고, 이 patch들의 linear embedding의 sequence를 transformer의 input으로 집어 넣는다. (patch는 NLP의 transformer에서 token(word)과 동일하게 다루어 짐을 의미한다.)
- 훈련 시, ImageNet과 같은 중간크기의 dataset에 대해 강한 규제화(regularization)없이는 비슷한 크기의 ResNet과 같은 구조의 모델들에 비해 정확도가 떨어졌다. 이에 대해 CNN은 내재된 inductive biases가 있으나 Transformer는 내재된 inductive biases의 부족으로 충분하지 못한 양의 dataset에 대한 훈련은 모델의 일반화성능이 좋지 않을 것이라 예측 하였다.
- 다만, 14M~300M정도의 대규모의 dataset의 경우, inductive biases를 찍어눌러버릴 수 있음을 연구진은 확인했다. ViT가 충분히 큰 dataset에 pre-train된 후 더 작은 dataset에 대해 transfer learning이 이루어지게 될 때, 기존의 S.O.T.A모델들의 성능과 비슷하거나 더 능가하는 결과를 보여준다. cf) the best model reaches the accuracy of 88.55% on ImageNet, 90.72% on ImageNet-ReaL, 94.55% on CIFAR-100, and 77.63% on the VTAB suite of 19 tasks.
2. Related Work
- Transformer[Vaswani2017]는 NLP분야의 기계번역에서 S.O.T.A방법이 되었으며, Large Transformer기반 모델들은 매우 큰 말뭉치(corpora)에 pre-train한 후 task의 목전에 전이학습(fine-tuning)을 진행한다 : BERT[Devlin2019]는 "Denoising Self-Supervised pre-training task"로 사용하고 GPT[Radford2018, 2019, Brown2020]는 "Language modeling pre-train task"로 활용한다.
- image에 대한 Naive한 "Self-Attention" 적용은 각 pixel이 다른 모든 pixel에 주목(attend)할 것을 요구한다. (= 하나의 pixel을 embedding 시, 다른 pixel도 embedding에 참여할 것이 요구된다는 의미) pixel수에서 2차적인 계산복잡도를 야기하며 이로 인해 다양한 input size로 확장되는 것이 어렵다. 즉, image processing에 transformer를 적용하려면 몇가지 근사(approximation)가 필요하다. ∙ local self-attention ∙ sparse attention ∙ 다양한 크기의 block에 scale attention 적용 이런 specialized attention구조들은 computer vision 분야에서 괜찮은 결과를 보여주지만 효율적 구현을 위해서는 복잡한 engineering이 필요하다.
- self-attention을 이용해 feature map을 augmentation하거나 CNN의 output에 attention을 추가적으로 사용하는 등 CNN과 self-attention을 융합하려는 많은 연구도 이루어 졌다.
- 본 연구와 가장 관련있는 모델은 Cordonnier2020논문에서 소개된 모델로 input image에서 2×2 크기의 patch를 추출하여 top에 full self-attention을 적용한다. ViT와 비슷하지만 ViT는 더 추가적인 증명을 진행하는데, 대규모 pre-training은 평범한 transformer가 S.O.T.A와 경쟁력 있게 한다. 추가적으로 Cordonnier2020논문의 모델은 2×2 piexl이라는 작은 크기의 patch를 사용하는데, 이는 small-resolution image에만 적용가능하다는 단점이 존재한다. (우리는 medium-resolution image를 다루지만...)
- 최근의 또다른 관련있는 모델은 iGPT[Chen2020]로 iGPT는 image의 resolution과 color space를 줄인 후 pixel들에 transformer를 적용한 생성모델로써 "Unsupervised"방식으로 훈련되었고 이를 통해 얻어진 representation은 classification을 위해 전이학습이나 선형적으로 탐색될 수 있으며, 이를 예시로 활용 시 ImageNet에서 최대 정확도 72%로 나타난다.
- 본 연구는 추가적으로 기존 ImageNet보다 더 많은 image를 보유한 인식데이터셋을 활용한 연구의 연장선으로 CNN의 성능이 dataset size에 따라 어떻게 달라지는 지에 대한 연구와 대규모 dataset(ImageNet-21K, JFT-300M)에 대한 CNN의 전이학습에 대해 경험적탐구도 진행한다. (이전 연구와 달리 ResNet이 아닌 Transformer를 사용하여.)
3. Method
- model 설계에서 기존 Transformer[Vaswani2017]를 가능한 근접하게 구성했으며, 이는 쉽게 확장가능한 NLP transformer 구조에 대한 이점과 효율적인 구현이 가능하다는 이점이 되어 별도의 설정없이 바로 사용할 수 있게 된다. 3.1 Vision Transformer (ViT) - 기존의 Transformer는 token embedding의 1D sequence를 input으로 받는다. - BERT의 [CLS]token처럼 embedding된 patch들의 가장 앞에 하나의 learnable한 class token embedding vector를 추가한다. - Position Embedding은 image의 위치정보를 유지하기 위해 patch embedding 시 trainable한 position embeddings가 더해진다. image를 위해 개선된 2D-aware position embedding을 사용해 보았지만 1D Position Embedding과의 유의미한 성능향상이 없어서 "1D Position Embedding"을 사용한다. (Appendix D.4) 이렇게 embedding된 벡터들의 sequence를 encoder의 입력에 넣는다.
- Transformer의 Encoder부분은 Multi-Head Attention(식 2,3)층들과 MLP가 교차로 구성되는데, 이는 하나의 image이더라도 차원을 쪼갠 뒤 multi-head를 관찰하는 현상을 볼 수 있다.
3.2 Fine-tuning. &. Higher Resolution
4. Experiments
연구진들은 ResNet, ViT, 그리고 하이브리드 모델의 representation learning capabilities를 평가했다. 각 모델이 요구하는 데이터의 양상을 파악하기 위해 다양한 사이즈의 데이터셋으로 사전훈련을 진행했고 많은 벤치마크 테스크에 대해 평가 실시했다. ViT가 모델의 사전훈련 연산 대비 성능 부분에서 다른 모델들보다 더 낮은 pre-training 비용으로 가장 높은 성능을 달성했다. 마지막으로 self-supervised를 이용한 작은 실험을 수행해 ViT가 self-supervised에서도 가능성이 있음을 보였다.
4.1 Setup ∙Datasets - 모델 확장성(scalability)을 조사하기 위해 1,000개 class와 1.3M image가 있는 ILSVRC-2012 ImageNet dataset, 21,000 class와 14M image가 있는 superset ImageNet-21k(Deng2009), 18,000 class와 303M의 고해상도 image의 JFT(Sun2017)를 사용.
[Kolesnikov2020]에 이어 downstream task의 testset을 통해 pre-training dataset을 중복 제거. 이 dataset에 대해 훈련된 모델을 다음과 같은 여러 benchmark에 전이학습을 진행한다: 원본 유효성 검사 라벨과 정리된 ReaL 라벨의 ImageNet(Beyer2020), CIFAR-10/100(Krizhevsky2009), Oxford-IIIT Pets(Parkhi2012) 및 Oxford Flowers-102(Nilsback & Ziserman, 2008). 이때, dataset의 전처리과정은 [Kolesnikov2020]를 따른다.
- 또한 19개 작업 VTAB classification suite(Zhai 2019)에 대해서도 평가한다. VTAB는 task당 1,000개의 training example들을 사용해 다양한 task에 대해 낮은 data transfer를 평가한다. task는 3가지로 분류된다: Natural(Pets, CIFAR), Specialized(의료, 위성사진), Structured(localization같은 기하학적 이해가 필요한 작업)
∙Model Variants - 표 1에 요약된 것처럼 BERT(Devlin 2019)에 사용된 ViT 구조를 기반으로 BERT에서 "Base" 및 "Large" 모델을 직접 채택하고 더 큰 "Huge" 모델을 추가한다. 후술에서는 모델의 크기와 입력 patch크기를 나타내기 위해 간략한 표기법을 사용. ex) ViT-L/16은 16 × 16의 input patch size를 갖는 "Large" variant를 의미. Transformer의 시퀀스 길이는 패치 크기의 제곱에 반비례하므로 패치 크기가 작은 모델은 계산 비용이 더 많이 든다. - 기본 CNN의 경우 ResNet을 사용한다. 다만, Batch Normalization(Ioffe & Szegedy, 2015)대신 Group Normalization(Wu & He, 2018)으로 바꾸고 표준화된 Convolution(Qiao 2019)을 사용한다. 이 수정은 transfer를 개선하며[Kolesnikov2020], 이를 "ResNet(BiT)라 부른다.
hybrid 시, 우리는 중간 feature map을 하나의 "pixel"의 patch 크기로 ViT에 전달한다. 다른 시퀀스 길이를 실험하기 위해 (i) 일반 ResNet50의 4단계 output을 가져오거나 (ii) 4단계 제거, 3단계에 동일한 수의 레이어(총 레이어 수 유지)를 배치하고 이 확장된 3단계 출력을 가져온다. 이때, (ii)를 사용하면 sequence 길이가 4배 더 길어지고 ViT 모델의 비용이 더 많아진다.
∙Training. &. Fine-tuning - ResNet을 포함, 모든 모델은 Adam(Kingma & Ba, 2015)을 사용한다. (β1 = 0.9, β2 = 0.999) batch=4096으로 훈련하고 모든 모델의 transfer에 유용한 0.1이라는 높은 weight_decay를 적용한다. (Appendix D.1은 일반적인 관행과 달리 Adam이 ResNets용 SGD보다 약간 더 잘 작동함을 보여준다).
linear Learning rate warmup과 decay를 사용한다. (detail은 Appendix B.1을 참조) 전이학습 시, 운동량이 있는 SGD, batch=512를 모든 모델에 대해 Appendix B.1.1에 소개된 것처럼 사용한다. 표 2의 ImageNet 결과의 경우 ViT-L/16의 경우 512, ViT-H/14의 경우 518로 미세 조정했으며 평균 계수 0.9999999로 [Polyak & Juditsky 1992]를 사용한다(Ramachan 2019, 2020).
∙Metric - few-shot이나 전이학습 정확도를 통해 downstream dataset에 대한 결과에 대해 설명한다. 전이학습 정확도는 각 dataset서 전이학습한 이후 각 모델의 성능을 capture한다. few-shot의 정확도는 train image의 표현을 {-1,1}K 표적 벡터에 매핑한 정규화된 최소 제곱 회귀(MLE) 문제를 해결하여 얻을 수 있다. 이 공식을 사용하면 폐쇄형(closed form)으로 정확한 솔루션을 복구할 수 있습니다. 주로 전이학습의 성능에 중점을 두지만, 전이학습비용이 너무 많이들어서 신속한 평가를 위해 선형 퓨샷 정확도를 사용하기도 한다.
4.2 Comparision to S.O.T.A - 먼저 가장 큰 모델인 ViT-H/14 및 ViT-L/16을 최신 CNN과 비교한다. 첫 번째 비교할 점은 대규모 ResNets로 감독된 전송 학습을 수행하는 Big Transfer(BiT)이다. 두 번째는 Noisy Student(Xie 2020)으로, 레이블이 제거된 상태에서 ImageNet 및 JFT-300M에서 준지도 학습을 사용하여 훈련된 대규모 EfficientNet이다.
- 표 2의 결과는 JFT-300M에서 pre-train된 작은 ViT-L/16 모델은 모든 작업에서 BiT-L을 능가하는 동시에 훈련에 훨씬 적은 계산 리소스를 필요로 함을 보여준다. 대형 모델인 ViT-H/14는 특히 ImageNet, CIFAR-100 및 VTAB suite 등 까다로운 dataset에서 성능을 더욱 향상시킨다.
흥미로운 점은, 이 모델은 이전기술보다 pre-train에 훨씬 적은 시간이 소요되었다. 하지만 pre-train의 효율성은 architecture의 선택뿐만 아니라 training schedule, optimizer, weight_decay 등 같은 다른 parameter에도 영향을 받을 수 있다는 점에 주목한다. Section 4.4에서는 다양한 architecture에 대한 성능과 computing에 대해 제어된 연구를 제공한다.
- 마지막으로, ImageNet-21k dataset으로 pre-train된 ViT-L/16 모델은 대부분의 dataset에서도 우수한 성능을 발휘하면서 pre-train에 필요한 리소스는 더 적습니다. 약 30일 내에 8개의 코어가 있는 표준 클라우드 TPUv3를 사용하여 교육될 수 있다. - 그림 2는 VTAB task를 각 그룹으로 분해하고, 이 benchmark의 이전 S.O.T.A 방법인 BiT, VIVI - ImageNet 및 Youtube에서 공동 교육된 ResNet(Tschannen 2020) 및 S4L - ImageNet에서 supervised 및 semi-supervised 학습(Zhai 2019)과 비교한다. ViT-H/14는 Natural 및 Structure task에서 BiT-R152x4 및 기타 방법을 능가하는데, Specialized에서 상위 두 모델의 성능은 유사하다.
4.3 Pre-training Data Requirements - ViT는 대규모 JFT-300M dataset에서 pre-train 시 성능이 우수하며 ResNet보다 Inductive bias가 적다. 그렇다면 dataset의 크기는 얼마나 중요할까? 에 대해 2가지 실험을 진행한다.
① 크기가 증가하는 dataset에 대해 ViT 모델을 pre-train한다: ImageNet, ImageNet-21k 및 JFT-300M. 소규모 dataset의 성능을 향상시키기 위해 weight_decay, dropout, label-smoothing이라는 3가지 기본적인 Regularization parameter들을 최적화한다. 그림 3은 ImageNet을 pre-train된 결과를 보여준다 cf. (다른 dataset에 대한 결과는 표 5에 나와 있다). ImageNet pre-train모델도 전이학습되어있으나 ImageNet에서는 다시 전이학습이 진행된다. (전이학습 시 해상도가 높아지면 성능이 향상되기 때문). 가장 작은 dataset인 ImageNet에서 pre-train 시, ViT-Large 모델은 (moderate) regularization에도 불구하고 ViT-Base 모델에 비해 성능이 떨어진다. JFT-300M만 있으면 더 큰 모델의 이점을 최대한 누릴 수 있는데, 그림 3은 또한 다양한 크기의 BiT 모델이 차지하는 성능 영역을 보여준다. BiT CNN은 ImageNet에서 ViT를 능가하지만 dataset이 클수록 ViT가 앞서는 것을 알 수 있다.
② 9M, 30M 및 90M의 랜덤 하위 집합과 전체 JFT-300M dataset에 대한 모델을 교육한다. 이때, 더 작은 하위 집합에 대해 추가적인 정규화를 수행하지 않고 모든 설정에 대해 동일한 hyper-parameter를 사용한다. 이 방식으로, regularization의 효과가 아니라 본질적인 모델 특성을 평가합니다. 하지만, Early Stop을 사용하고, training 중 달성한 최고의 validation accuracy를 알려준다. 계산을 절약하기 위해 전체의 전이학습 정확도 대신 few-shot linear accuracy를 알려주며, 이는 그림 4에는 결과가 나와 있다. ViT는 소규모 dataset에서 비슷한 계산 비용으로 ResNet보다 더 적합하다. ex) ViT-B/32는 ResNet50보다 약간 빠르다.
9M 하위 집합에서는 성능이 훨씬 떨어지지만 90M+ 하위 집합에서는 성능이 더 우수한데, ResNet152x2 및 ViT-L/16에 대해서도 마찬가지이다. 이 결과는 Convolution의 Inductive Bias가 작은 dataset에 유용하다는 직관을 강화한다. 하지만, 큰 dataset의 경우 관련 패턴을 데이터에서 직접 학습하는 것만으로도 충분하고, 심지어 유용하다.
4.4. Scaling Study - JFT-300M의 transfer 성능을 평가하여 다양한 모델에 대해 제어된 scaling 연구를 진행한다. 이 설정에서 데이터 크기는 모델의 성능에 병목 현상을 일으키지 않으며, 각 모델의 accuracy/pre-train cost를 평가한다. model set은 다음과 같다. - 7 epochs에 대해 사전훈련된 ResNets(R50x1, R50x2 R101x1, R152x1, R152x2, R152x2) - 7 epochs에 대해 사전 훈련된 R16 및 R1450 - 14 epochs에 대해 사전 훈련된 R152x2 및 R200x3 - 14 epochs에 대해 사전 훈련된 ViT-B/32, B/16, L/32, L/16 - 14 epochs에 대해 사전 훈련된 R50+ViT-L/16 (hybrid의 경우 모델 이름 끝 숫자는 patch size가 아닌, ResNet 백본의 총 downsampling비율을 나타낸다).
- 그림 5에는 총 사전 교육 계산 대비 이전 성능이 나와 있다(compute detail: Appendix D.5 ; model별 detail: Appendix의 표 6). 이에 대해 몇 가지 패턴을 관찰할 수 있다: i) ViT는 accuracy/computing 절충에서 ResNets를 압도한다. ViT는 동일한 성능을 얻기 위해 약 2~4배 적은 컴퓨팅을 사용한다(평균 5개 이상의 dataset). ii) hybrid는 적은 computing으로 ViT를 약간 능가하나, 더 큰 모델에서는 그 차이가 사라진다. 이 결과는 Convolution의 local feature processing이 모든 크기에서 ViT를 지원할 것으로 예상할 수 있기 때문에 다소 놀라운 결과라 할 수 있다. iii) ViT는 시도한 범위 내에서 포화되지 않는 것으로 나타나 향후 확장 노력에 동기를 준다.
4.5. Inspecting Vision Transformer - ViT는 image data처리방법이해를 위해 internal-representation을 분석한다. ViT의 첫층은 flattened patch를 낮은 차원의 공간(Eq.1)으로 linearly project한다. 그림 7(왼쪽)은 학습된 embedding filter의 주요 구성요소를 보여준다. 구성요소는 각 patch내 fine-structure를 낮은 차원으로 표현하기 위한 그럴듯한 기본함수와 유사하다. - 투영한 후, 학습된 position embedding이 patch의 representation에 추가된다. 그림 7(가운데)은 모델이 position embedding의 유사성에서 image내 distance를 encoding하는 방법을 학습하는 것을 보여준다. 즉, 더 가까운 patch는 더 유사한 position embedding을 갖는 경향이 있다.
- 또한 행-열(row-column) 구조가 나타난다. 동일한 행/열에 있는 패치는 유사한 임베딩을 갖는다.
- 마지막으로, sin파의 구조는 때때로 더 큰 grid(Appendix D)에 apparent한데, position embedding이 2D-image topology를 표현하는 방법을 배운다는 것은 직접 설계된 2D-aware embedding이 개선되지 않는 이유를 설명한다.(Appendix D.4)
- Self-Attention을 통해 ViT는 전체적으로 정보를 통합할 수 있는데, 가장 낮은 층에서도 이미지를 생성할 수 있다. 연구자들은 신경망이 이 수용력을 어느정도까지 사용하는지 조사한다. 특히 attention weight를 기반으로 정보가 통합되는 image space의 average distance를 계산한다(그림 7, 오른쪽). 이 "attention distance"는 CNN의 receptive field size와 비슷하다.
- 연구자들은 일부 head가 이미 낮은 층에 있는 대부분의 image에 주목을 하는것을 발견, 정보를 global하게 통합하는 특징이 모델에 실제로 사용되는것을 보여준다. 다른 attention heads는 낮은 층에서 일관되게 attention distance가 작다. 이렇게 고도로 local하게 된 attention은 transformer(그림 7, 오른쪽)이전에 ResNet을 적용하는 hybrid 모델에서 덜 두드러지며, 이는 CNN의 초기 convolution층과 유사한 기능을 수행할 수 있음을 보여준다. 또한 신경망의 깊이에 따라 attention distance가 증가하는데, Global하게 모델이 clasification과 의미론적(semantically)으로 관련이 있는 image의 영역에 주목함을 발견할 수 있다. (그림 6).
4.6. Self-Supervision - Transformer는 NLP 작업에서 인상적인 성능을 보여주지만, 그들의 성공의 대부분은 뛰어난 확장성뿐만 아니라 대규모 self-supervised pretraining(Devlin 2019; Radford 2018)에서 비롯됩니다. 또한 BERT에서 사용되는 masked language modeling task를 모방해 self-supervision을 위한 masked patch prediction에 대한 예비 탐색을 수행한다. "Self-supervised pre-training"을 통해, 우리의 작은 ViT-B/16 모델은 ImageNet에서 79.9%의 정확도를 달성하고, 처음부터 교육에 대해 2%의 상당한 개선을 이루나, 여전히 supervised pre-training에 비해 4% 뒤떨어져 있습니다. Appendix B.1.2에는 추가 detail이 존재한다.
5. Conclusion
- 연구진들은 image 인식에 대한 transformer의 직접적인 적용을 탐구하며 이전의 연구들에서 computer vision에 "self-attention"을 사용한 것과 달리 초기 patch생성을 제외하고 image별 inductive biases를 architecture에 도입하지 않는다. 대신, image를 일련의 patch로 해석해 기본 Transformer의 Encoder로 처리한다. 이는 단순하지만 확장가능하기에 대규모 dataset에 대한 pre-train과 결합하게 되면 매우 잘 작동한다. 따라서 ViT는 많은 S.O.T.A모델과 비슷하거나 능가하지만 pre-train에 상대적으로 cheap하게 작동한다.
- 본 연구의 이런 시초적인 결과는 고무적이라 할 수 있으나 많은 과제가 여전히 남아있는데, ① Detection, Segmentation 등의 Computer Vision에 적용하는 것과 ② "Self-Supervised pre-training"에 대한 방법을 찾는 것이다. 우리의 이 초기실험은 "Self-Supervised pre-training"에 개선을 보여주긴 하지만 대규모 "Supervised pre-training"보다는 여전히 큰 격차가 존재하기에 ViT의 추가적인 확장은 성능향상의 여지를 보여준다.
😶 부록 (Appendix)
[A] Multi-Head Self Attention
[B] Experiment Details
이를 위해, trainingset의 작은 부분을 validation set으로 설정한다. (CIFAR은 2%, ImageNet은 1%)ResNet의 경우, [Kolesnikov2020]을 따르며 모든 fine-tuning 실험은 384resolution에서 실행된다. [Kolesnikov2020]에 따르면, training과 다른 resolution으로fine-tuning을 하는 것이 일반적이다.
[C] Additional Results
논문에 제시된 수치에 해당하는 세부 결과로 표 5는 논문의 그림 3에 해당하며 크기가 증가하는 dataset에서 사전 훈련된 다양한 ViT 모델의 전이학습 성능을 보여준다: ImageNet, ImageNet-21k 및 JFT-300M. 표 6은 논문의 그림 5에 해당하며 다양한 크기의 ViT, ResNet 및 하이브리드 모델의 전송 성능과 사전 교육의 예상 계산 비용을 보여준다.
