[5주차 / 임정준 / 논문리뷰] An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale
#01 Abstract
Transformer 구조가 NLP 분야에서 사실상 'Standard' 모델이 되었지만, CV 분야에서는 그 적용이 미진하였다. Vision 분야에서는 CNN 구조의 몇 부분만 대체하거나 혹은 네트워크 간의 결합을 할 때 'attention'을 활용하였다. 하지만 본 논문에서는 이미지들을 패치들로 쪼개어 Sequence 형태로 순수 Transformer에 적용하였을 때 이미지 분류 Task에서 그 성능이 SOTA 수준급으로 도달할 수 있음을 보여주었다.
본 리뷰에서는 해당 모델은 'ViT'에 대해 알아보고, 일반적인 CNN 네트워크와 어떠한 차이점이 존재하는지 알아보도록 하겠다. 더욱 더 자세한 이해를 위해서는 기본적으로 Transformer에 대한 전반적인 지식을 이해하고 오는 것을 권장합니다 !
#02 Introduction
지금까지 일반적인 CNN network에 self-attention이 결합된 시도, 그리고 self-attention으로 layer 들을 전부 대체한 시도는 있었으나 후자 모델의 경우 attention pattern으로 인하여 이론적으로는 효율적일 수는 있으나 스케일링이 효과적으로 되지 않아 하드웨어 가속을 하는데 어려움이 있었다. 따라서 고전적인 ResNet과 같은 Architecture들이 아직 SOTA 모델의 자리를 지키고 있는 이유였었다.
NLP 분야에서 널리 사용되던 Transformer를 CV 분야에서 효율적으로 활용하기 위하여, 이미지들을 patch들로 분리하고, Transformer의 input으로 패치들 간의 선형 임베딩 sequence를 넣어주었다. 이미지 패치들이 NLP 분야의 token과 같은 역할을 하는 것이다.
Transformer는 Inductive bias(귀납적 편향)이 CNN처럼 내재되어 있지 않아, translation equivariance 및 locality와 같은 가정들에서 어려움이 있어, 불충분한 데이터로 학습 시에는 일반화로 잘 수렴되지 않는 경향이 있다. Inductive bias란 만나지 못한 상황을 해결하기 위해 '추가적인 가정'을 상정하는 것인데, CNN의 경우 Locality를 이용하여 공간적(Spatial) 특징을 잘 활용하고(Conv layer의 filter), RNN 계열 모델의 경우 순차적(Sequential) 특징을 잘 활용하여
상기 이미지처럼 Translation Invariance를 나타낸 것인데, CNN의 경우 세 이미지를 모두 같은 것이라고 할 수 있지만(translation-able) Transformer의 경우 순차성(Sequential)에 강점을 보이는 모델이기에 세 이미지에 대하여 모두 같은 것이라고 충분히 일반화되지 않는 측면이 있다.
하지만 대용량의 Dataset에서 학습이 이루어진다면, Inductive Bias를 충분히 커버할 수준으로 큰 학습이 이루어지는 것을 확인할 수 있다. 따라서 충분한 스케일에서 Pre-trained model을 생성하고, 이를 이용해 Transfer Learning을 한다면 다른 데이터셋들에서도 훨씬 좋은 성능을 낼 수 있음 역시 확인할 수 있었다.
#03 Method
##03 -1 Vision Transformer(ViT)
ViT는 Transformer 모델의 Encoder 부분을 가져와서 이를 Image-specific Transformer로 바꾸어 내었다. 이미지를 고정된 사이즈로 분리하고 여기에 Position embedding 및 [Class] embedding을 추가한다. 그리고 이를 Standard Transformer encoder에 넣어 Sequence Vector 결과값을 반환한다. 그 후, [Class] Token을 활용하여 이를 MLP Head에서 처리하여 알맞은 Image Classification 결과값을 반환한다. Input Feature는 다음과 같다.
Input Dimension:
- Token Embedding Sequence, 1D Input으로 Reshape
: H*W*C --> N * P*P * C, 단 N=H*W/P*P , Transformer에 효과적인 Sequence 길이로서 Serving
[ (H, W): original image 해상도 / C : 채널 수 / (P, P) : 각 이미지 패치의 해상도 ]
- Constant Latent Vector Size: D
: D 차원으로 mapping하여 linear projection이 train 가능하도록 함
: Patch Embedding이 output으로 출력
- [CLS] Token
: BERT의 [class] token과 유사하게, 학습 가능한 (trainable) embeded patch(z_0)의 Sequence를 추가
: z_L - L 번째 layer의 첫번째 토큰(Classification Head)
--> 본 Head는 MLP에서 하나의 hidden layer로 pre-training, fine-tuning 시에 하나의 단일 layer로 구현
- Position embeddings
: Positional information(위치 정보)를 유지하기 위하여 Patch Embedding을 추가
: Standard Learnable 1D position embeddings 활용
* Transformer Encoder
: Multiheaded self-attention layers로 구성
: MLP blocks 적용
: 기존 모델과 다르게 Layer Normalization & Residual Connection을 모든 block에 들어가기 '전'에 적용
: ReLU 대신 비선형함수 GELU 사용
+) GELU
GELU란 RELU의 결정론적인(deterministic) 효과를 살리면서도, 확률론적인(stochastic) 결과를 도출해내는 비선형함수이다.
식을 보게되면, 기존의 input에 표준정규분포의 CDF(Cumulative Distribution Function)를 이용하여 확률적인 값을 계산해내지만 그 값은 특정 '값'으로 나오는 결정론적인 것이라고 할 수 있다. ReLU의 smoothing 버전이라고 할 수 있는데, 2번째의 그래프는 GELU 함수를 근사시킨 그래프로 그 특징을 확인할 수 있다.
* Inductive Bias
: 다른 CNN 모델들보다 이미지에 대해 덜 inductive하기 때문에, 모델 전반적으로 Locality를 확보할 수 있는 장치를 마련해야 한다
: Vit에서는 오직 MLP layer만 local & translationally equivariant(공간적 해석능력 확보)
<--> self-attention layer는 Global함(해석능력 확보 X)
* Hybrid Architecture
: Raw image patch들의 대안으로, CNN의 feature map으로부터 input sequence를 구성할 수 있다.
: 본 Hybrid model에서는 patch embedding projection이 CNN에서 추출된 patch들이 적용되었다.
--> 1x1의 Spatial Size를 가지고 있는데, 이는 Transformer 차원이 단순히 1차원화(flatten)된 feature map의 공간적 차원을 요하기 때문이다.
##03-2 Fine-Tuning & Higher Resolution
일반적으로, 큰 Dataset에서 ViT를 미리 학습시키고 그 후 작은 Task에 대해 Fine-tuning하는 방식으로 수행된다. 이를 위해서 pre-trained prediction head를 제거하고, 0으로 시작하는(zero-initialized) D x K 차원의 feedforward layer를 부착한다. (K는 하위 Task의 Class 개수) 이 방식은 Higher Resolution task에 대하여 fine-tuning하는데 유리한 것으로 나타났다.
높은 해상도를 넣을 때 patch size를 유지하면 그 sequence 길이가 늘어난다. 여기에는 ViT가 Sequence 길이를 임의적으로 다루는 것이 가능하다. 하지만 Fine-Tuning 시, pre-trained position embedding은 더 이상 의미가 없는 것으로 드러났다. 따라서, 2D interpolation(2차원 보간)을 통하여 pre-trained position embedding을 얻어내었다.
#04 Experiment
ResNet, ViT, 그리고 Hybrid 모델들의 학습 성능을 구현하고 비교하였으며, 데이터의 필요조건에 대해 이해하기 위하여 다양한 데이터셋에 대하여 pre-training 해보고 해당 모델들을 벤치마킹하였다. 결론적으로 ViT는 다른 모델들과 비교했을 때 더 적은 pre-training cost를 가지고도 SOTA 수준의 벤치마킹을 달성하였으며, self-supervision(자기-감독)을 활용한 작은 실험 역시 수행되었다.
BERT와 유사하게, pre-training dataset의 크기에 따라 Base/Large/Huge에 따라 모델을 구분하였으며, 해당 모델의 차원과 사이즈는 상기 Table과 같다. 또한 Notation으로, 예를 들어 ViT-L/16은 16*16 input patch size의 ViT-Large model을 의미한다.
Baseline인 CNN에 대하여, BatchNorm 대신 Group Normalization을 사용하였으며, Standardized Convolution을 활용하였다.
+) Group Normalization: Batch Size의 크기가 작을 때 Batch Normalization 성능이 저하된다는 점에 주목하여 이를 개선한 표준화
: C=G, 혹인 G=1이면 Layer Normalization과 동일
+) Standardized Convolution: Convolution layer에서 weight(가중치)들에 대해 Normalization을 진행하여 loss 경사를 smoothing시키는 기법
또한, hybrid 모델에서는 ResNet50에서의 다양한 시도를 통하여 Intermediate feature-map을 확보하였다.
* Training
- Adam(with Beta_1=0.9 , Beta_2=0.999, 4096 batch size, weight decay=0.1)
- Linear learning rate warmup & decay
* Fine-tuning
- SGD with momentum(512 batch size)
- 해상도 변경(ViT-L/16: 512 , ViT-H/14 : 518)
- Few-shot Accuracy 활용 / Fine-tuning Accuracy 역시 Fine-tuning 후 측정
: Regularized least-squares regression
+) Cost 절감을 위하여 linear few-shot accuracy도 활용
BiT-L(large ResNet에서 supervised transfer learning이 활용된 모델), Noisy Student(semi-supervised learning으로 학습된 large EfficientNet)과도 비교하여 보았을 때 ViT가 더 좋은 성능을 내고 있는 것을 확인할 수 있다. 하지만, 아키텍쳐 외 다른 parameter(e.g. training schedule, optimizer, weight decay etc.)들로 인해 성능에 영향이 있을 수 있으므로 다양한 실험을 진행한 결과, 더 적은 resource로도 대부분의 데이터셋에서 더 좋은 성능을 이끌어냈다는 것 역시 확인할 수 있다.
##04-1 Dataset Size?
데이터셋의 크기를 달리 하여 실험을 진행했었고, 작은 데이터셋에서는 최적화를 위하여 weight decay 조정, dropout 및 label smoothing을 진행하였다. 그러나 결론적으로 작은 데이터셋에서는 비교적 경쟁력있는 성능을 내지 못하였다.
이에 대한 원인으로는 직관적으로 Inductive Bias의 결여라고 설명할 수 있는데, 비교적 큰 데이터셋에서는 연관 패턴을 직관적으로 학습하는 것이 충분하고, 오히려 직관적이라는 해석이 가능하다.
##04-2 Scaling Study
첫째로, ViT가 ResNet에 비하여 모든 Trade-off에서 우세한 것을 확인할 수 있다. 둘째, Hybrid 모델이 ViT모델보다 약간 앞서나 큰 모델로 갈수록 그 차이가 줄어들며 오히려 성능이 좋아지는 구간도 존재한다. 이는 Hybrid 모델의 CNN에서 나온 1*1 sequence들에서 나온 Convolutional local feature가 어느 사이즈에서든 ViT를 도와주도록 하는 것을 알 수 있다. 마지막으로, ViT는 시도된 범위 내에서는 '포화(Saturate)'되지 않은 것으로 관측되는데, 이는 추후 Scaling effort를 더 높여볼 수 있는 기제로 작동할 수 있겠다.
##04-3 ViT Inspection
상기된 이미지 내용을 기반으로, Self-attention 모델이 낮은 layer들에서 ViT로 하여금 정보를 통합시킬 수 있도록 하는 것을 확인할 수 있다. 실제로 이렇게 나타난 'attention distance'는 CNN network에서 초반 layer들에서 관측되는 Inspection의 형태와 유사하다는 것을 알 수 있다. 또한 depth가 깊어질수록 attention distance가 길어져 해당 이미지에 대해 '전체적인' 접근을 한다는 사실을 파악할 수 있다.
또한, 두번째 Position embedding similarity(유사도)에서 알 수 있듯, 가까운 Patch 간의 유사도 및 공간 정보가 잘 학습됐다는 것을 확인할 수 있다. 즉, 낮은 Layer들에게서도 Localization 효과를 보이며, Similarity representation을 통하여 해당 위치 정보 역시 잘 학습했다는 것으로 해석할 수 있겠다.
##04-4 Self-Supervision
NLP tasks들에서 수행되는 것처럼, 한 문장에서 특정 단어를 masking 후 이를 prediction하는 Task와 같이 Vision task에서도 self-supervision 방식을 수행하였다.
해당 방식으로 Pre-training된 ViT-B/16 모델의 경우 79.9%의 정확도를 보였으나 이는 일반적은 supervised pre-training 보다는 4% 정도 성능이 뒤떨어진다.
#05 Conclusion
Image-specific Inductive Bias에서 벗어나, Standard Transformer Encoder에 적용할 수 있도록 이미지를 patch들의 sequence라고 해석하였다. 그리고 이를 Self-Attention에 적용하여 여타 SOTA CNN 모델들과 견줄 정도의 성능을 발휘하였다. 이렇게 간단하고 강력한 모델은 큰 데이터셋에서 pre-training 되었을 때 그 우수한 성능이 드러나는 것을 확인하였다.
추후 Detection 및 Segmentation에서의 적용, 그리고 self-supervised pre-training과 같이 아직 연구가 미진했던 분야에 대한 도전, 그리고 Scaling Study에서 나온 성능 향상 Motivation이 남아있다. Computer Vision의 Transformer적 적용은 Text 분야가 당면해있던 과제들과도 연결될 수 있다고 생각한다. CNN 구조의 소멸, 혹은 아키텍처를 결합하여 CV 분야와 NLP 분야 간의 효율적인 Computation 등 생성 모델이나 실시간으로 다중 정보를 받아들이는 모델에서 더 좋은 성능을 내지 않을까라고 내심 생각하고 있다.