Transformer 구조는 NLP에서 사실상 표준적인 방식으로 자리잡았지만 컴퓨터비전에서는 아직 활용도가 높지 않습니다. 컴퓨터비전에서 적용된 방식은 대부분 CNN과 함께 쓰이거나 일부를 대체하는 식인데요, 이 논문에서는 CNN에 의존하지 않는 순수한 Transformer만 가지고도 image classification 문제를 잘 해결할 수 있음을 보입니다.
Proposed Method
확장하는 Transformer 구조를 쉽게 차용하고 효율적인 구현을 활용하기 위하여 원래 Transformer 구조를 변형 없이 그대로 사용하였습니다. Transformer는 token embedding의 1D sequence를 입력으로 받기 때문에 2D 이미지를 patch로 잘라 각각의 patch를 flatten하여 마치 token처럼 사용합니다.
BERT의 [class] 토큰과 유사하게, learnable embedding을 embedded patche 시퀀스의 앞에 붙입니다(). Position embedding 또한 마찬가지로 patch embedding에 추가합니다.
이 구조를 표현하는 식은 아래와 같습니다.
Hybrid Architecture
이미지를 단순하게 패치로 잘라 embedding 하는 대신, CNN에서 나온 feature map을 사용할 수 있을 것입니다. 2D feature map을 flatten하여 input sequence로 사용한 hybrid model도 실험하였습니다.
Fine-Tuning and Higher Resolution
이 ViT는 큰 데이터셋에서 pre-train된 후에 작은 데이터셋(태스크)에 fine-tuning됩니다. 이 때 prediction head를 제거하고 zero-initialized fc 레이어를 다시 추가하였습니다. Fine-tuning시에는 더 큰 해상도의 이미지를 사용하는 것이 도움이 될 때가 있기 때문에, ViT에서는 패치 크기는 유지하고 sequence의 길이를 늘리는 방식을 사용하였습니다.
Experiment
Setup
Datasets
Pre-training은 ImageNet(1k classes, 1.3M images), ImageNet-21k(21k classes, 14M images), JFT(18k classes, 303M HR images)에 수행하고 다른 데이터셋에 transfer 하였습니다.
Transfer한 데이터셋으로는 ImageNet(original validation labels), ImageNet(cleaned ReaL labels), CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102을 사용했습니다. 이 밖에도 19개의 task를 모아놓은 VTAB classification suite에도 실험하였습니다.
Model Variants
모델은 BERT 구조의 속성을 그대로 차용하되 크기를 바꿔서 ViT-B~H 세 가지로 만들었습니다.
비교를 위한 baseline CNN으로는 ResNet(BiT)을 사용했습니다.
Comparison to State-of-the-Art
아래 결과와 같이 ViT 모델은 ImageNet 벤치마크에서 SOTA 성능을 보여주는 두 모델(BiT, Noisy Student)에 비해 짧은 pre-training 시간을 가지는데 비해 비슷하거나 더 나은 성능을 보여줍니다.
아래 결과는 VTAB의 19개 task를 분류에 따라 Natural, Specialized, Structure로 나눠 평가한 것입니다. 대부분의 결과에서 더 나은 성능을 보였습니다.
Pre-training Data Requirements
ViT 모델은 JFT-300M 데이터셋에 pre-training 했을 때 잘 동작했습니다. 이미지에 대한 inductive bias가 부족한 설계인데, 이 데이터셋의 크기는 얼마나 중요할까요? 이를 확인하기 위해 pre-training 데이터셋의 크기를 바꾸면서 실험해보았습니다.
결과적으로 few-shot transfer에서 만족스러운 결과를 얻었습니다. 하지만 pre-training 데이터셋의 크기가 작을 때 baseline에 비해 성능은 높지 않았습니다.
Scaling Study
아래 실험은 모델의 크기와 pre-training epochs에 종속적인 pre-training compute에 따른 transfer accuracy를 비교한 것입니다. 일반적으로 ViT는 ResNet에 비해 더 나은 trade-off를 보여주었고, ResNet feature를 사용하는 Hybrid 모델은 pre-training compute가 작을 때 ViT에 비해 약간 더 나았지만 compute가 커질 수록 차이는 줄었습니다. ViT는 compute의 증가에 따라 성능이 saturation하지 않는 것으로 보아 더 많은 리소스를 사용할 때 더 큰 성능의 증가를 기대할 수 있습니다.
5. Inspecting Vision Transformer
Transformer가 어떻게 이미지 데이터를 처리하는지 살펴보기 위해 internal representation을 관찰해보았습니다.
(left) Embedding filter의 top principal components를 의미합니다. 이들은 각 patch에 존재하는 low-dimensional representation에 대한 basis function처럼 보입니다.
(center) 한 이미지의 input patch들에 대한 position embedding의 유사도를 표현한 그림입니다. 자세히 보면 같은 row, 같은 column에 있는 patch들은 유사한 embedding을 가지는 것을 볼 수 있습니다.
(right) ViT는 가장 낮은 레이어세서도 이미지의 모든 부분으로부터 정보를 볼 수 있습니다. 이 그림은 레이어별로 attention weight를 기반으로 어떤 정보가 모이고 있는지 표현한 것입니다. 이 거리는 CNN에서 receptive field에 해당합니다.
6. Self-Supervision
NLP에서 Transformer 모델의 성공은 그 scalability 뿐만 아니라 large-scale pre-training 덕분입니다. 이 실험에서는 BERT와 유사하게 masked patch prediction 문제를 self-supervision으로 하여 학습시켜보았습니다. 작은 ViT-B/16 모델을 self-supervised로 학습시킨 결과 ImageNet에서 79.9%의 성능으로 from scratch에 비해 2% 향상을, 하지만 supervised training에 비해 4% 낮은 결과를 얻었습니다.
Conclusions
이전에 self-supervision을 컴퓨터비전에 적용한 다른 work들과 다르게, 이번 논문에서는 이미지에 대한 어떤 inductive bias도 사용하지 않은 구조를 제안하여 SOTA 성능을 달성할 수 있음을 보였습니다. 이 단순하고 scalable한 구조는 큰 데이터셋에 pre-training될 때 아주 잘 동작했습니다.
하지만 여전히 챌린지는 남아있습니다. 하나는 object detection이나 segmentation과 같은 다른 task에 적용하는 것이고, self-supervised pre-training 방법을 더 연구하는 것입니다.
