Vision Transformer (2)

Vision Transformer (1)

---

Experiments

How fine-tuning

실험 시에는 ViT 모델을 큰 데이터셋에 대해 사전훈련 하고 작은 데이터셋에 대해 fine-tuning 하는 과정을 거칩니다. Fine-tuning 을 위해서 사전훈련된 $z_L^0$에 붙은 MLP 를 제거하고 0으로 초기화된 $D\times K$ 선형변환 층을 추가합니다. ($K$는 fine-tuning 데이터셋 클래스 개수) 또한, fine-tuning 시에 고해상도 이미지를 사용하는 것이 좋다고 알려져 있기 때문에 패치 크기를 일정하게 유지한 채 큰 크기의 이미지를 사용합니다. 패치 크기가 고정되니 상대적으로 시퀀스 길이가 더 늘어나겠죠. 단 이러한 경우에는 사전훈련을 통해 학습한 position 임베딩 텐서가 무의미해지므로 고해상도 크기와 원래 이미지 크기에 따른 위치에 알맞게 2차원 보간을 해줍니다.

Setup

모델의 scalability 를 확인하기 위해 1) ImageNet (1.3M 이미지, 1k 클래스), 2) ImageNet-21k (14M 이미지, 21k 클래스), 3) JFT (303M 이미지, 18k 클래스) 데이터셋에 대해 사전훈련하고 fine-tuning 데이터셋과 겹치는 데이터를 제거합니다. 모델로는 Table 1과 같이 Base / Large / Huge 3가지 종류를 사용합니다. 예를 들어 ViT-L/16 은 패치 크기 16의 Large 종류를 말합니다. 특히, 패치 사이즈가 줄어들수록 시퀀스 길이가 늘어나므로 계산량이 더 많아집니다.

CNN 계열의 baseline 으로는 ResNet 모델이 사용되고 hybrid architecture 로는 ResNet50의 stage 4의 출력을 사용합니다. 배치 사이즈는 4096을 사용하고 모든 모델에 대해서 사전훈련에는 Adam optimizer, fine-tuning 에는 SGD momentum optimizer 를 사용합니다. (ResNet 사전훈련시 일반적으로 SGD optimizer 가 사용되나 논문의 세팅으로는  Table 7과 같이 Adam optimizer 성능이 더 좋았다고 합니다.) 

Comparison to SOTA

전이학습 평가를 위한 각 데이터셋 별 ViT 모델과 기존 전이학습 SOTA CNN 모델과의 (BiT-L / Noisy Student) 성능비교는 Table 2와 같습니다. JFT 데이터셋으로 사전훈련한 ViT-L/16 모델은 더 적은 연산을 가지고도 모든 벤치마크에 대해서 BiT-L 보다 높은 성능을 보였습니다. 또한, ViT-H/14 와 같이 큰 모델은 성능이 더욱 향상되었는데요, 중요한 점은 모델이 매우 커졌음에도 소요되는 연산시간이 그렇게 크지 늘지 않았다는 점입니다.

Pretraining data requirements

Table 2를 보면 JFT 데이터셋으로 사전훈련한 모델의 성능은 기존 SOTA에 비해 높지만 이보다 작은 데이터셋인 I21k (ImageNet 21k) 에 대해서는 성능이 ResNet 계열에 비해 낮습니다. 적게 설계된 inductive bias 를 극복하기 위한 적절한 데이터셋 크기를 알아보기 위해 ImageNet $\rightarrow$ ImageNet 21k $\rightarrow$ JFT 순으로 사전훈련하고 ImageNet 데이터셋에 대해 Figure 3과 같이 실험합니다. 이때 작은 데이터셋의 효과를 보정하기 위해 weight decay, dropout, label smoothing 등과 같은 정규화 기법들이 사용됩니다. 참고로 dropout 은 qkv 선형변환에는 적용되지 않고 위치 임베딩이 더해진 이후와 일반 MLP 층에 대해 적용됩니다.

ImageNet 과 같은 작은 데이터셋에 대해서는 ViT-Large 계열은 Base 계열에 비해 성능이 낮지만 JFT 300M 과 같은 매우 큰 데이터셋에 대해서는 가장 큰 Huge 모델이 성능이 제일 좋습니다. 또한, Figure 4는 JFT 300M 데이터셋을 일부분만 사용했을 때의 ImageNet 에 대한 성능으로 마찬가지로 데이터셋의 크기가 클 수록 ViT 계열이 성능이 더 좋아집니다.

Scaling study

이번에는 JFT 300M 데이터셋과 다른 하이퍼파라미터를 동일하게 유지한채 모델 별로의 계산량 / 성능을 비교하는 실험입니다. 결과는 Figure 5로 ViT 계열은 ResNet 계열에 비해 동일한 성능을 얻기 위한 연산량이 배 이상 적습니다. 또한, hybrid 모델의 경우 약간의 연산량 오버헤드가 존재하지만 순수 ViT 모델들에 비해 성능이 더 좋습니다. 특히, 모델이 매우 커지면 성능이 어느정도 포화되었던 ResNET 계열과 달리 ViT 계열은 포화되는 양상을 보이지 않습니다. 즉, 이 실험보다 훨씬 큰 스케일에 대한 성능이 기대되는 부분입니다.

Inspecting vision transformer

ViT 가 어떻게 이미지 데이터를 처리하는지 알아보기 위해 먼저 Figure 7의 왼쪽 그림처럼 패치 임베딩에 사용된 파라미터를 분석하였고 위치 임베딩또한 Figure 7의 가운데 그림처럼 분석합니다. 그림을 보면 비슷한 위치에 존재하는 패치에 대해서는 비슷한 임베딩 결과가 도출되고 같은 줄 / 열에 대해서도 어느 정도 비슷한 패턴을 보입니다. 특히, Figure 10에서처럼 더 큰 그리드에서는 sinusoidal 패턴을 보이기도 하는것으로 보아 전에 언급했던 것처럼 2차원 위치 임베딩이 굳이 필요없이 스스로 잘 학습되었다고 생각할 수 있습니다.

마지막으로 알아볼 부분은 ViT 에서의 self-attention 동작입니다. Self-attention 으로 ViT 에서는 낮은 layer 에서도 특정 위치에서 이미지 전체를 조망할 수 있는데요, 이를 구체적으로 알아보기 위해 Figure 7의 오른쪽 그림처럼 각 층별 attention 거리를 계산합니다. 여기서 attention 거리는 CNN 에서의 receptive field 와 비슷한 개념인데 각 헤드별로 낮은 layer 에 대해서는 attention 거리가 넓게 분포되어 있어 낮은 layer 부터 이미지 전체를 바라보는 몇 개의 헤드들이 존재하며, layer 가 증가하면서 attention 평균거리는 증가합니다. 특이할 점은 ResNet feature map을 이용한 hybrid 모델에 대해서는 attention 거리가 평균적으로 높은데, 이는 feature map 입력 자체가 이미 어느정도 주변 공간정보를 반영되었기 때문이라고 추정할 수 있습니다.

 

---

Vision Transformer (3) - Attention Map

Vision Transformer (4) - Pytorch 구현