https://arxiv.org/abs/2010.11929
An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
While the Transformer architecture has become the de-facto standard for natural language processing tasks, its applications to computer vision remain limited. In vision, attention is either applied in conjunction with convolutional networks, or used to rep
arxiv.org
해당 논문을 보고 작성했습니다.
Abstract
Transformer 구조가 자연어 처리 task에서 대세가 되었지만, computer vision에선 여전히 한계가 있습니다. vision에서는 attention이 convolution network와 결합되거나, 전체 구조는 유지한 채로 convolutional network의 특정 요소를 대체하는 정도로 사용되었습니다. 저자들은 CNN을 필수적으로 사용할 필요가 없고 image patch의 sequence에 바로 transformer를 사용하는 것이 image classification task에서 잘 동작함을 보입니다. 대규모 data로 pre-train 하고 다양한 mid-size or small image recognition benchmark (ImageNet, CIFAR-100, VTAB, etc.)로 transfer 할 때, Vision Transformer (ViT)는 최신 convolutional network보다 상당히 적은 학습 연산 resource를 사용하면서도 뛰어난 성능을 보여줍니다.
Introduction
computer vision에서 convolutional architecture가 여전히 지배적으로 많이 사용되고 있습니다. NLP에서의 성공에 영감을 받아, CNN-like architecture에 self-attention을 결합하거나, convolution을 완전히 대체하는 다양한 연구들이 등장했습니다. 후자의 경우, 이론적으로 효율적이지만 특수한 attention pattern을 사용하기 때문에 현대 가속기에서 효과적으로 확장되지 않았습니다.
NLP에서 Transformer의 확장 성공에 영감을 받아, 저자들은 가능한 한 최소한의 수정으로 standard Transformer를 image에 바로 사용하는 실험을 수행했습니다. 이를 위해, 저자들은 image를 patch로 나누고 각 patch를 linaer embedding sequence으로 만들어 Transformer의 input으로 사용했습니다. image patch는 NLP의 token과 같이 사용됩니다. 저자들은 supervised fashion에서 image classification을 수행하도록 model을 학습했습니다.
저자들이 mid-sized dataset으로 학습할 때, 이 model은 비슷한 크기의 ResNet보다 몇 퍼센트 낮은 정확도를 보입니다. 이는 좋지 않은 결과로 볼 수 있지만, Tranformer는 CNN이 가지고 있는 translation equivariance, locality와 같은 inductive basis가 존재하지 않기 때문에, 충분한 양의 data로 학습되지 않으면 일반화가 잘 되지 않습니다.
하지만 대규모 dataset으로 model을 학습한다면 상황이 달라집니다. 저자들은 대규모 학습이 inductive bias를 능가한다는 것을 발견했습니다. 저자들의 Vision Transformer (ViT)는 충분한 양으로 pre-train 되고 적은 양의 datapoint로 특정 task로의 transfer가 수행되었을 때, 좋은 결과를 보였습니다.
Related Work
self-attention을 image에 적용하는 Naive application은 각 pixel이 다른 pixel에 attend되어야 합니다. pixel 수의 제곱만큼 cost가 들기 때문에, realistic 한 input size를 사용할 수 없게 됩니다. 그러므로 image 처리에서 Transformer를 사용하기 위해서 여러 시도가 있었습니다. 각 query pixel의 인접한 pixel에만 self-attention을 적용하는 연구가 있었습니다. 이러한 local multi-head dot-product self attention block을 convolution의 대체로 사용하는 연구도 있었습니다. 이러한 다양한 specialized attention architecture들은 computer vision task에서 유망한 결과를 얻었지만, hardware accelerator에 효율적으로 적용하기 위해선 복잡한 engineering을 필요로 합니다.
저자들의 model가 가장 연관 있는 model로, input image에서 2 x 2 patch를 추출하고 full self-attention을 적용하는 연구가 있습니다. 이 model은 ViT와 매우 유사하지만, 저자들의 연구는 large scale pre-training이 vanilla transformer가 최신 CNNs에 경쟁력 있거나 오히려 더 좋은 성능을 보이도록 만들어 준다는 것을 잘 입증합니다. 그리고 유사했던 model은 small patch size인 2x2 pixel을 사용하여 낮은 resolution image에만 사용할 수 있지만, 저자들의 model은 medium-resolution image도 다룰 수 있습니다.
이뿐만 아니라 convolutional neural networks (CNNs)을 self-attention 형태와 결합하는 다양한 연구도 있었습니다. 예를 들어 feature map을 augmentation 하여 image classification을 수행하거나, CNN의 output에 self-attention을 수행하는 연구들이 있었습니다.
또 다른 최근 관련된 model은 image GPT (iGPT)입니다. iGPT는 image resolution과 color space를 줄인 후 image pixel에 Transformer를 수행합니다. model은 생성 model로 unsupervised fashion에서 학습되며, 생성된 representation을 fine-tune 하거나 linaer 하게 조절하여 classification performance를 수행할 수 있고, ImageNet에서 72% 정도의 정확도를 보여줍니다.
Method
저자들의 model 구조는 최대한 original Transformer와 유사하도록 만들었습니다. 이러한 의도적인 간단한 setup의 이점은 확장 가능한 MLP Transformer 구조와 그 효율성을 그대로 사용할 수 있다는 것입니다.
Vision Transformer (ViT)
model 구조는 위와 같습니다. standard Transformer는 1D token embedding sequence를 input으로 받습니다. 2D image를 다루기 위해, 저자들은 image x ∈ R^{H x W x C}를 flattened 2D patches sequence x_p ∈ R^{N x (P^2 x c)}로 reshape 하였습니다. 여기서 (H, W)는 original image의 resolution을 의미하고, C는 channel 수를 의미합니다. (P, P)는 각 image patch의 resolution을 의미하고, N = HW/P^2가 patch 수가 되며 이는 Trasformer의 input sequence length가 됩니다. Transformer는 모든 layer에 constant latent vector size D를 사용하기 때문에, 저자들은 patch를 flatten 한 다음 학습 가능한 linear projection을 사용하여 D차원으로 mapping합니다. 이 projection의 output을 patch embedding이라 부릅니다.
BERT의 [class] token과 유사하게, 저자들은 학습 가능한 embedding을 embedded patches sequence 앞에 붙입니다(z_0^0 = x_{class}). Transformer encoder의 output z_L^0은 image representation y로 사용됩니다. pre-training 할 땐, classification head로 1개 hidden layer가 있는 MLP를 사용하고 fine-tuning 할 땐 single linear layer를 사용합니다.
Position embedding은 positional information을 얻기 위해 patch embedding에 추가합니다. 저자들은 더 향상된 2D-aware position emedding을 사용했을 때 큰 성능 향상을 얻지 못했기 때문에, standard learnable 1D position embedding을 사용했습니다. 최종 embedding vector sequence를 encoder의 input으로 사용합니다.
Transformer encoder는 multi-headed self-attention (MSA)와 MLP block의 교차로 구성됩니다. LayerNorm (LN)은 각 block 이전에 사용되며, 각 block 이후에 residual connection을 사용했습니다. MLP는 GELU non-linearity와 함께 2개 layer를 사용합니다.
- Inductive bias
저자들은 Vision Transformer가 CNN에 비해 훨씬 더 적은 image-specific inductive bias를 갖는다는 사실에 주목합니다. CNN에서는 locality, 2차원 neighborhood structure, translation equivaraince이 전체 model에 걸쳐 각 layer마다 내장되어 있습니다. ViT에서는 MLP layer만 local & translationally equivariant를 가지고 있으며, self-attention layer은 global 합니다. model 시작 부분에서 image를 patch로 cutting 하고, fine-tuning 할 땐 서로 다른 resolution의 image를 조정할 때를 제외하고는 2차원 neighborhood structure를 사용하지 않습니다. 그 외에는 초기화 시 position embedding은 patch의 2D position에 대한 정보를 가지고 있지 않으며, patch 사이 공간적 관계는 처음부터 학습됩니다.
- Hybrid Architecture
raw image patch를 사용하는 대신, input sequence로 CNN의 feature map을 사용할 수도 있습니다. 이러한 hybrid model에서 patch embedding projection E가 CNN feature map에서 추출된 patch에 적용됩니다. 특별한 경우, patch는 1x1 size가 가능하며, 이는 feature map의 공간적 차원을 단순히 flatten 하고 Transformer dimension으로 project 하여 input sequence를 얻는다는 것을 의미합니다. classification input embedding과 position embedding은 위에서 언급한 것과 같이 추가됩니다.
Fine-Tuning and Higher Resolution
일반적으로, 저자들은 ViT를 대규모 dataset으로 pretrain한 다음 downstream task로 fine-tune 합니다. 이를 위해, 저자들은 pre-trained prediction head를 제거하고 zero-initialized D x K feedforward layer를 붙였습니다. 여기서 K는 downstream의 class 수입니다. pre-train보다 더 높은 resolution으로 fine-tune 하는 것이 종종 이점을 가져다줍니다. 더 높은 resolution의 image를 feed 할 때, patch size는 동일하게 유지하여 더 긴 sequence length가 됩니다. Vision Transformer는 메모리 제약 내에서 임의의 sequence length를 처리할 수 있지만, pre-trained position embedding은 더 이상 의미가 없어지게 됩니다. 그래서 저자들은 original image의 location에 맞춰 pre-trained position embedding의 2D interpolation을 수행했습니다. 이 resolution 조정과 patch extraction은 Vision Transformer에 2D 구조에 대한 inductive bias를 vision transformer에 주입하는 유일한 지점입니다.
Experiments
저자들은 ResNet, Vision Transformer (ViT)와 hybrid model의 representation learning 성능을 평가했습니다. 각 model이 필요로 하는 data 양을 이해하기 위해, 다양한 크기의 dataset으로 pre-train 하였고 여러 benchmark task에서 평가를 진행했습니다. model을 pre-train 할 때 사용되는 연산량을 고려하면, ViT는 더 낮은 pre-train cost를 가지고 대부분의 recognition benchmark에서 최신 model과 유사한 성능을 보여 매우 좋은 모습을 보여줍니다.
Model Variants
저자들은 ViT 구성을 BERT에서 사용된 구조를 기반으로 하며, 위와 같습니다. "Base와 "Large" model은 BERT에서 바로 가져왔으며, "Huge"는 더 큰 model입니다. ViT-L/16의 경우 "Large" version이며 16x16 input patch size를 사용한다는 것을 의미합니다. Transformer의 sequence length는 patch 크기의 제곱에 반비례하므로, patch size가 작은 model일수록 연산량은 더 많아집니다.
baseline CNN으로 저자들은 BN layer를 Group Normalization으로 대체한 ResNet을 사용했습니다. 이러한 수정은 transfer를 향상시키고, 수정된 model을 "ResNet (BiT)"라 부릅니다. hybrid의 경우, 중간 feature map을 patch size가 1 pixel인 ViT에 feed 했습니다. 다양한 sequence length에 대한 실험을 진행하기 위해, 저자들은 1) regular ResNet50의 stage 4 출력을 사용하거나, 2) stage 4를 제거하고 동일한 수의 layer를 stage 3에 배치하여(총 layer 수를 유지), 이 확장된 stage 3의 output을 사용했습니다. 두 번째 경우가 sequence length가 4배 더 길고, ViT의 연산량이 더 많습니다.
Comparison to State of the Art
저자들은 먼저 가장 큰 model인 ViT-H/14와 ViT-L/16을 최신 CNN model과 비교했습니다. 첫번째 비교 대상은 대형 ResNet을 사용한 supervised transfer learning을 수행하는 Big Transfer (BiT)입니다. 두 번째는 semi-supervised learning을 사용하여 ImageNet과 JFT-300M에서 label이 제거된 상태로 학습된 대형 EfficientNet인 Noisy Student입니다. 현재 Noisy Student는 ImageNet에서 가장 좋은 성능을 보여주고, BiT-L이 다른 dataset에서 가장 좋은 성능을 보여줍니다.
실험 결과는 위와 같습니다. JFT-300M로 pre-train된 ViT-L/16이 BiT-L보다 모든 task에서 뛰어난 성능을 보여주며, model의 크기가 더 작고 학습할 때 연산량이 훨씬 더 적게 사용됩니다. 더 큰 model인 ViT-H/14는 특히 더 어려운 dataset인 ImageNet, CIFAR-100, VTAB에서 더 뛰어난 성능을 보여줍니다.
이 외에도 다양한 실험이 존재합니다.
Conclusion
저자들은 image recognition을 수행하기 위해 Transformer를 바로 사용하는 연구를 수행했습니다. computer vision에서 self-attention을 사용하는 이전 연구들과 다르게, 저자들은 초기 patch extraction step을 제외하고는 image-specific inductive bias를 model에 적용하지 않았습니다. 대신, image를 patch의 sequence로 변환하고 standard Transformer encoder로 처리했습니다. 이러한 간단한 방식은 pre-train 하고 대규모 dataset을 사용했을 때 상당히 좋은 성능을 보였습니다. Vision Transformer는 많은 image classification dataset에서 더 적은 연산량으로 가장 좋은 성능을 보여줍니다.
하지만 detection이나 segmentation과 같은 다른 computer vision task에 ViT를 사용하는 것은 여전히 challenge로 남아있습니다. 저자들의 초기 실험에서 self-supervised pre-training의 성능을 향상했지만, 여전히 self-supervised와 supervised pre-train 사이의 gap 차이가 존재합니다.
