[컴퓨터 비전] Vision Transformers
"An Image is Worth 16 x 16 words : Transformers for Image recognition At scale"
1. Vision Transformer 란?
ViT는 Google Research 팀이 제안한 모델로, 기존 자연어 처리(NLP) 분야에서 성공한 Transformer 구조를 이미지 인식에 그대로 적용하는 방식입니다.
이 논문의 핵심 아이디어 중 하나는 "이미지를 여러 개의 작은 패치(Patch)로 나누고 이를 토큰처럼 처리하여 Transformer로 학습하는 것"입니다. 이미지 패치를 토큰 처리하여 기존의 Transformer를 최소한으로 수정으로 이미지 분야에 적용시켰습니다.
ViT를 이용해서 기존 CNN 기반 네트워크가 아니라 컴퓨터 비전 분야에서 Transformer의 활용 가능성을 확장할 수 있음을 보여주었습니다.
2. ViT의 구조 및 특징
2-1. 2가지 임베딩 벡터
기존의 Transfomer를 사용하기 위해서는 2D 이미지를 Transformer의 입력 형태에 맞게 변형하는 과정이 필요합니다. 이 논문에서는 "patch embedding", "position embedding" 이렇게 두가지 임베딩 벡터를 제시하고 있습니다.
Patch embedding은 이미지를 토큰을 사용하기 위한 방법입니다. 이미지를 P x P 패치로 나눕니다. 그리고 flatten을 진행하여 Transformer의 input size에 맞게 Linear Projection을 시킵니다. 또한 추가적으로 NLP의 BERT처럼, 특별한 [class] 토큰을 시퀀스 앞에 추가합니다. 이 토큰은 이미지 전체를 대표하는 벡터로 학습되어 최종 분류에 사용됩니다.
position embedding은 Transformer가 패치 순서를 인식할 수 있도록 벡터로, patch embedding에 해당 임베딩을 추가합니다. position embedding은 일반적인 1D learnable 위치 임베딩을 사용합니다. 실험 결과 복잡한 2D-aware 임베딩을 사용했더니 별다른 성능 차이를 보이지 않았습니다.
수식적으로 다음처럼 표현한다.
여기서 MSA는 Multi Self Attention이고, LN는 Layer Norm, MLP는 Multi Layer Perceptron을 의미한다.
2-2. ViT의 특징 : 적은 Inductive Bias
Inductive Bias는 모델이 학습하기 전에 세상에 대해 미리 가지고 있는 "기본적인 가정이나 전제"를 의미합니다. 미리 특정한 가정을 통해서 데이터를 분석하도록 설계하여 좀 더 General한 모델을 학습될 수 있도록 유도하도록 합니다.
CNN은 이미지 처리 적합한 Inductive bias를 가지고 있습니다. 1) Locality : 가까운 픽셀들은 의미 있는 정보를 공유할 것이다. 2) Translation invariance : 패턴은 어디에 있든 같은 것으로 인식 되어야 된다. 이 두가지 특징으로 가지고 있습니다.
하지만 Transformer는 이런 Locality, Translation invariance 특징을 가지고 있지 않습니다. 처음부터 글로벌하게 이미지 특징을 학습합니다. 이에 공간적인 관계를 표현하기 위해서는 Positional embedding과 학습 데이터로 새롭게 습득해야 되는 특징을 가지고 있습니다.
2-3. Hybrid Architecture
CNN을 같이 사용하는 Hybrid Architecture로도 사용이 가능하다. Patch embedding을 생성할 때, CNN feature map을 이용해서 projection을 수행할 수 있다.
2-4. Fine Tuning 방법
Fine Tuning을 진행하기 위해서 미리 학습된 prediction head를 제거하고, Task에 맞는 feedforward 레이어를 추가하여 학습하면된다.
이때 미리 학습된 모델보다 큰 해상도를 사용했을 효과가 더 좋았다. 더 큰 해상도를 사용할 때에 주의 할 점이 존재한다. 같은 patch size에서 큰 해상도를 사용하게 되면 입력 데이터의 길이가 길어질 수 있다. 이 때 학습된 Positional embedding의 의미가 사라질 수도 있다는 점이다. 이를 해결하기 위해서는 2D interpolation을 이용한다.
3. 실험 방법
3-1. 모델에 따른 성능 실험
논문에서는 ResNet, ViT, Hybrid 모델을 사용해서 성능을 측정했다.
ResNet 모델은 기존 ResNet 구조 살짝 바꾼 ResNet(BiT)을 사용했다. Batch Normalization 대신 Group Noarmalization을 사용하고 Transfer learning을 통해 성능을 향상 시킨 모델이다.
사용된 ViT모델은 ViT-Base, ViT-Large, ViT-Huge를 사용했다. Base와 Large는 BERT에서 그대로 사용했고, Huge만 여기서 만들어서 사용했다. 이에 대한 자세한 내용은 아래와 같다.
Layers | Hidden size D | MLP size | Heads | Params | |
ViT-Base | 12 | 768 | 3072 | 12 | 86M |
ViT-Large | 24 | 1024 | 4096 | 16 | 307M |
ViT-Huge | 32 | 1280 | 5120 | 16 | 632M |
Hybrid 모델은 ViT 입력 특징점으로 ResNet50의 중간 feature map으로 사용해 만들었다.
실험은 여러 데이터셋의 분류 모델의 Accuracy로 비교 했습니다. pre-train 데이터셋으로 Google의 JFP(JFT-300M), I21k(ImageNet-21k)로 학습했다고 합니다.
밑 결과를 보면 ViT-H/14 모델이 가장 좋은 성능을 내는 것을 볼 수 있습니다. 특히 ResNet기반의 BiT-L 모델에 비해서 좋은 성능을 내는 것을 보여줍니다. 하지만 pretrain을 위한 계산량을 비교하면 현저히 ViT-H/14가 BiT-L보다 적은 것을 확인할 수 있습니다.
3-2. 데이터 크기에 따른 성능 결과
Transformer의 경우에는 CNN과 달리 이미지에 적합한 Inductive bias가 없어 많은 데이터로 학습을 시켜 공간적인 관계를 학습을 시켜주어야 된다. 이에 데이터 크기에 따라서 ViT에 얼마나 큰 성능 차이를 보이는 가에 대해서 실험을 진행했다.
사용 데이터로는 ImageNet-1k (130만장), ImageNet-21k (2100만장), JFT-300M (3억장)을 사용했다. 오른쪽으로 갈수록 많은 데이터를 의미한다.
아래의 결과를 보면 작은 데이터로 학습했을 때에는 ViT가 ResNet에 비해서 적은 성능을 보여준다. 하지만 데이터가 많아 질수록 ViT의 성능이 ResNet과 비슷한 성능을 내는 것을 확인할 수 있다.
또한 ResNet는 특정 데이터량을 넘어가면 성능적 한계를 보이지만, ViT는 데이터량이 넘어갈수록 성능이 비례해서 향상되는 것을 확인할 수 있다.
해당 결과로 인해 ViT의 적은 Inductive Bias는 많은 학습 데이터로 커버가 가능하다는 것을 알 수 있습니다.
4. Inspecting Vision Transformer
4-1. low 단계에서의 embedding filter
flattent한 patch를 linear projection하는 레이어를 분석하면 아래처럼 나온다. 이를 보면 이미지의 low한 특징을 잘 뽑아내는 것을 확인할 수 있다.
4-2. Position embedding
아래 그림은 각 행과 열의 position embedding과 모든 embedding의 코사인 유사도를 나타내는 것입니다. 이를 보면 비슷한 위치에 있는 embedding은 유사한 값을 가지고, 또한 같은 행과 열의 embedding도 유사한 값을 가지는 것을 확인할 수 있습니다. 이를 통해서 position embedding이 잘 학습된 걸 확인할 수 있습니다.
또한 Grid를 많이 사용하게 되면, 기존 Transformer에서 사용되는 hand-crafted 2D-aware embedding과 유사한 구조를 가지는 것을 확인할 수 있습니다.
4-3. Self-Attention
ViT는 낮은 계층에서도 이미지 전체를 참조할 수 있습니다. CNN에서는 초기 레이어에는 receptive field가 좁고, 레이어가 깊어질수록 receptive field가 넓어져 전체적인 특징을 파악합니다. 하지만 ViT는 Self-Attention이 존재하여 처음부터 전체적인 특징을 파악합니다. 이를 증명하기 위해서 레이어마다 어텐션 가중치를 기반으로 Attention distance를 계산하여 했습니다.
밑에 그래프는 레이어 depth에 따른 Attention distance 입니다. 결과를 보면 초기 레이어부터 전체적인 부분을 보는 Head가 존재하는 것을 확인할 수 있습니다.
5. Self-Supervised Pre-Training
기존의 NLP 계열 (BERT, GPT)는 Self-Supervised 방식으로 학습이 효과적인 것으로 알려졌습니다. 이에 ViT도 Self-Supervised 방식이 유의미한지 실험을 진행했습니다.
NLP 계열과 마찬가지로 입력 패치 중에서 일부를 마스킹하여, 이를 복원하는 하는 방식으로 학습을 진행했습니다.
이에 결과는 ImageNet에서 79.9%로 낮은 결과를 보여주었습니다. 하지만 ViT가 supervised 방식이 아니여도 어느 정도는 학습이 진행된다는 것을 확인할 수 있습니다.
6. 정리
ViT (Vision Transformer) : 이미지를 여러 개의 작은 패치(Patch)로 나누고 이를 토큰처럼 처리하여 Transformer로 학습한 모델
ViT의 특징
- 적은 Inductive Bias : CNN에 비해서 적은 Inductive Bias을 가지지만, 학습 데이터량으로 충분히 공간적인 특징을 학습시킬 수 있다.
- 전역 이미지를 고려 : self-attention으로 CNN과 달리 처음부터 전체적인 구조를 보고 학습을 진행한다.
- 뛰어난 확장성 : CNN에 비해 좀더 유연한 구조로 모델의 크기를 키울수록 성능이 꾸준히 확장된다.
- 시각적 해석이 쉬움 : Attention Map을 통해 ViT가 어디에 주목하고 있는 지 시각화가 가능
- Self-attention의 가능성 : supervised learning에 비해 성능은 낮지만, 어느 정도 효과를 보인다.