안녕하세요. 강희승입니다. 지난주 Transformer에 이어, Computer Vision 연구에 Transformer를 적용한 ViT에 대해서 리뷰하려고 합니다. ViT는 현재 VLM에서도 많이 채택되어 활용되기 때문에, 다시 한번 복습하고자 해당 논문을 선정하였습니다.
1. Introduction
“Attention is All You Need”에서 Transformer 구조가 처음 등장하고, NLP 분야에서는 해당 구조가 표준으로 자리를 잡았습니다. Vision 분야에서는 여전히 Convolution 구조가 지배적이었지만, NLP에서의 Transformer의 성공에 영감을 받아, Convolution과 Self-Attention을 결합하려는 시도가 있었습니다. 또한, Convolution을 완전히 Self-Attention으로 대체하는 방식의 일부 연구도 있었습니다.
이러한 연구들은 이론적으로는 효율적인 구조지만, Image를 처리하기 위해 설계된 복잡한 Attention Pattern은 그 당시 Hardware Accelerator 상에서 효율적으로 학습하기 어려웠습니다. 따라서, large-scale recognition에서는 ResNet-like 구조들이 여전히 SOTA 자리를 꿰찼습니다.
여기서 저자는 기존 연구들이 Self-Attention이라는 구조 자체에 집중한 접근과 다르게, NLP에서 Transformer가 data scaling과 computation efficiency가 가능하게 된 이유에 집중하였습니다. 기존 연구들의 경우, 이미지를 처리하기 위해 Self-Attention의 구조를 변형하였지만, 저자는 Input을 변형하는 접근을 하였습니다. 즉 image를 약간의 수정을 통해 Transformer의 Self-Attention에 그대로 적용하여 해당 구조의 이점을 취하였습니다.
1.1 Inductive Bias
Transformer는 이런 CNN의 Inductive Bias가 결여되어 있다고 저자는 언급합니다.
여기서 Inductive Bias란, ‘학습데이터에서 보지 못한 데이터에도 generalization이 가능하도록, 어떠한 입/출력의 관계가 어떠한 구조를 가질 것이다’ 는 가정을 모델에 사전 부여하는 것입니다. Logistic regression을 예로 들자면, input과 output이 linear 관계에 있을 것이라는 Inductive Bias를 가집니다. 해당 가정을 기반으로 Decision Boundary가 Linear라는 형태를 가지게 됩니다.
CNN이 가지는 대표적인 Inductive Bias는 다음과 같습니다.
- Locality CNN은 Convolution Kernel을 Sliding Window 기법으로 연산을 수행합니다. FC layer의 경우, 하나의 pixel이 전체 pixel과 연결될 수 있지만, 해당 연산 기법은 pixel 간 연결 범위를 kernel의 크기만큼 영역을 제한합니다. 따라서, 연결 구조 자체가 제한됨으로써 locality bias가 주입됩니다. 이러한 특징을 기반으로 CNN layer를 여러 차례 거치고 Receptive Field가 커지면서, 초기 layer의 저수준의 Local 특징들을 결합하여 점점 큰 구조를 학습하도록 유도될 수 있습니다.
- Translation Equivariance 해당 bias는 입력 이미지가 이동하면 feature map도 같은 방향으로 이동하는 성질입니다. CNN이 Convolution 연산을 수행할 때, weight sharing을 하기 때문에, 어떤 특징이 어떤 위치에 있던, 동일한 특징을 추출합니다. 따라서, 이미지가 이동하면서 특징에 따른 feature map의 활성화 위치 또한 동일하게 이동합니다.
Transformer는 input에 대해 전역적으로 Self Attention을 수행합니다. 즉, 특정 토큰이 모든 토큰을 참고합니다. 또한, Convolution 연산이라는 특정 연산을 수행하는 것이 아니기 때문에, CNN의 Inductive Bias가 결여되어 있다고 생각할 수 있습니다.
따라서, Vision 분야에서 Transformer를 적용하였을 때, 충분한 데이터가 있지 않다면 일반화가 어렵습니다.
2. Method
그럼에도 불구하고, 저자는 Transformer의 이점을 챙기면서도, Vision에서 Image recognition을 일반화하기 위해 Vision Transformer 줄여서 ViT를 제안하였습니다.
2.1 Patch Embedding
NLP에서 Transformer는 1D Sequence를 기반한 Token Embedding을 입력으로 받습니다. 이를 Image에 적용하기 위해, flatten된 2D image의 x \in \mathbb{R^{H\times W \times C}} 에 대해, 2D patch x_p \in \mathbb{R^{N\times (P^2 \cdot C)}} 로 변환합니다.
Patch resolutions (P,P)
N = HW/P^2 로 패치의 개수를 의미
저자의 방법론에서 사용하는 Transformer는 입출력은 D 차원의 고정된 latent representation을 사용합니다. 따라서, 각 patch들을 flatten한 P^2C 차원의 벡터를 trainable linear projection을 통해 D 의 dimension을 가지는 Embedding Space에 투영합니다.
2.2 [CLS] TOKEN
저자는 BERT ( BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding ) 에서 제안한 [CLS] token의 개념을 ViT에 적용하였습니다.
[CLS] token이란 learnable embedding으로, patch sequence의 가장 앞단에 추가되어 학습됩니다. 해당 token은 Transformer Encoder의 Self-Attention을 거치면서 이미지 패치들의 정보를 참조하여, 이미지 전체를 대표하도록 학습되는 token이라고 할 수 있습니다. 이후 해당 token만 추출하여, MLP로 구현된 Classification head를 거쳐 예측을 수행합니다.
구체적으로 (1) 수식에서 z_0은 초기 CLS token으로 이미지 패치들의 정보가 없는 상태로 시작합니다. 이후 (2),(3) 즉 Transformer의 MSA;Multi-head Self-Attention과 MLP 수행하여, CLS token이 모든 이미지 패치들을 참조하도록 하고 해당 token을 비선형 변환을 통해 CLS token의 representation을 정제합니다. 마지막으로 (4)에서 최종 output y를 0번째 token 즉 [CLS] token으로 선택합니다.
2.3 Position Embedding
ViT에서도 동일하게 Position Embedding을 적용하여, 각 패치들에 Position Information을 주입합니다. 다만, 여기서 흥미로운 점은 2D Position Embedding을 사용하는 것이 아닌 1D Position Embedding을 적용하였습니다. 저자가 Ablation을 통해 실험한 결과, 2D-aware 방식이 1D 방식보다 이점을 발견하지 못했다고 합니다.
2.3.1 Ablation
저자는 총 4가지 Position Embedding을 설계하고 실험하여, 왜 2D 방식을 사용하지 않았는지 구체적으로 제시하였습니다.
다음은 저자가 설계한 4가지 Embedding 방식입니다.
- No Position Embedding
- Bag of Patches
- 1D Position Embedding
- Default
- 패치의 flatten order
- 2D Position Embedding
- 이미지 수직 수평을 고려
- Relative Position Embedding
- 절대적 위치보다 상대적 위치를 고려
- 1D relative Attention을 사용
- offset = p_q - p_k 로 특정 패치의 q, 다른 패치의 k embedding의 차를 두 패치간 상대적 거리로 정의하였다.
추가로, Position Embedding을 (1)Transformer의 입력되기 전 한번만 추가, (2) 모든 transformer layer의 시작 부분에서 자신만의 embedding을 더하여 학습, (3) 모든 transformer layer의 시작부분에 공유되는 embedding을 더하여 학습하는 방식을 추가로 설계하였습니다. 위치 정보를 한번만 제공해도 충분한지, 또는 layer 마다 위치 정보를 제공하는 것이 유리한지 확인하기 위해 설계되었습니다.
Table의 결과를 보면 알 수 있듯이, positional embedding을 수행하지 않은 경우와 수행한 경우와의 차이는 크지만, 3가지의 Embedding 방식은 큰 차이가 없는 것을 확인할 수 있습니다. 저자는 ViT의 Transformer Encoder가 Pixel-Level 이 아닌 Patch-Level의 input을 사용하기 때문에, 공간 정보를 Encoding하는 방식의 차이가 상대적으로 덜 중요하다고 해석했습니다. 즉, Positional Encoding은 반드시 필요하나, 방식의 차이보단, 해당 Patch의 Representation이 더욱 중요하다고 해석하였습니다.
2.4 Transformer Encoder
Transformer Encoder는 Patch Embedding 을 Input으로, Multi-head Self-Attention과 MLP(2 layers with a GELU)를 교대로 수행합니다. 각 Block들의 앞에는 Layer Normalization을 수행하고 끝에는 Residual Connection을 적용합니다.
2.5 Fine-Tuning and Higher Resolution
ViT를 활용하는 방법은 다음과 같습니다. Large dataset에 대해 사전학습된 ViT를 기반으로, downstream task로 전이학습을 수행합니다. 이때, 사전학습된 prediction head를 제거하고, zero-init된 Feedforward layer를 추가하여 학습을 합니다.
ViT의 경우, sequence 길이가 학습에 사용된 길이보다 길어도 수행이 가능합니다. 따라서, 더 큰 해상도의 이미지도 동일한 patch size로 처리가 가능합니다. 다만, 기존 학습된 Positional Encoding이 의미가 없어지기 때문에, 기존 사전 학습된 Embedding을 기반으로 2D interpolation을 수행합니다.
저자는 resolution adjustment와 patch extraction이 ViT에서 image 2D 구조에 대한 Inductive Bias가 주입되는 유일한 지점이라고 합니다.
3. Experiments
저자는 그 당시 Vision 분야에서 SOTA인 ResNet을 주요 비교 대상으로 지정하였습니다.
Setup
저자는 다음과 같이 세 가지 모델을 구성하였습니다.
Dataset
- Pretraining
1000개의 classes와 130만 개의 이미지로 구성된 ImageNet-1k dataset, 21,000 classes와 1,400만 개의 이미지로 구성된 ImageNet-21K dataset 과 18,000개의 classes와 3억 300만개의 high resolution 이미지로 구성된 JFT를 사용하였습니다. 이때, 데이터셋에서 Downstream task의 dataset과 중복되는 데이터를 제거하고 사전학습을 진행합니다. - Fine-tuning ImageNet, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102를 기반으로 fine-tuning 및 Validation을 수행합니다.
3.1 Comparison to SOTA
비교 대상 첫 번째인 BiT는 Large ResNet을 기반으로 supervised transfer learning을 수행한 모델입니다. 두 번째는 Noisy Student로 EfficientNet을 기반으로 준지도학습을 이용하여 학습시킨 모델입니다. ViT 모델들의 scale이 커질 수록, 성능이 오르며, ViT-L/16의 경우 기존 SOTA 성능을 모두 능가하며, Huge 모델의 경우 더 큰 향상을 이루었습니다.
또한, ViT의 경우 사전학습 시 기존의 SOTA 모델에 비해 훨씬 적은 compute resource를 필요로 하였습니다. 가장 큰 Huge의 경우에도, BiT-L에 비해 약 4배 가량 적은 compute를 사용하였습니다.
3.2 Pre-Training Data Requirements
ViT는 large scale 데이터인 JFT-300M dataset으로 사전학습하였을 경우, 좋은 성능을 보입니다. 여기서, 저자는 Resnet보다 적은 Inductive Bias를 가지는 ViT가 dataset의 크기가 얼마나 치명적으로 작용하는지 실험하고자 하였습니다.
- 가장 먼저, ImageNet-1k, ImageNet-21k, JFT-300M 순서로 점차 데이터셋의 scale을 증가시키며 사전 학습을 수행하였습니다. Figure3에서 볼 수 있듯, 작은 데이터셋을 기반으로 학습시켰을 경우, BiT가 ViT를 outperform하지만, 더 큰 데이터셋을 기반으로 사전학습할 시, ViT가 더 높은 성능을 보이는 것을 알 수 있습니다.
- 다음으로, JFT-300M을 9M, 30M, 90M,300M 으로 random subsampling 하여 학습을 수행하였습니다. 해당 실험에서도 동일하게 작은 데이터셋(9M)에서 ViT는 BiT에 비해 낮은 성능을 보이지만, 30M 부터 BiT 작은 모델의 성능을 능가하기 시작해서, 더 큰 데이터셋을 사용할 경우 ViT가 더 높은 성능을 보였습니다.
저자는 Convolution Inductive Bias는 작은 데이터셋에서 더욱 용이하지만, 큰 데이터셋에서는 충분한 데이터를 기반으로 관련 패턴을 학습하는 것이 beneficial 하다고 해석하였습니다.
Conclusion
저자는 Image-specific Inductive Bias를 Transformer에 주입하는 것이 아닌, image를 sequence of patch로 취급하여 standard Transformer를 Vision 분야에 그대로 적용하였습니다. 결과적으로, large dataset을 기반으로 사전 학습을 수행할 때 좋은 성능을 보일 수 있었으며, 기존 SOTA의 성능도 능가하였습니다.
이러한 결과는 CNN 기반의 시각적 Inductive Bias 없이도 Large dataset과 training scale이 보장된다면, Transformer(ViT)가 이미지에 필요한 visual representation을 data로부터 학습할 수 있음을 보였습니다. 또한, standard Transformer을 적용함으로써, 구조적 호환성과 뛰어난 scaling 특성 덕에 이후 이미지와 텍스트의 alignment를 맞추고 결합하는 VLM의 Vision Encoder로 많이 채택되어 활용되는 것 같습니다.
이상으로 ViT 리뷰를 마치겠습니다. 읽어주셔서 감사합니다.
