[8] MOBILEVIT: LIGHT-WEIGHT, GENERAL-PURPOSE,AND MOBILE-FRIENDLY VISION TRANSFORMER

[Paper] https://arxiv.org/pdf/2110.02178.pdf

 

[Github] https://github.com/apple/ml-cvnets

GitHub - apple/ml-cvnets: CVNets: A library for training computer vision networks

 

이번 포스팅에서는 mobileViT 논문에 대해 리뷰하겠습니다 :) 

 

 

Abstract  

비전 분야에서 많이 쓰이는 대표적인 신경망 중 하나가 바로 CNN이죠! CNN은 spatial inductive bias가 있기 때문에 다양한 비전 task에 더 적은  파라미터로 학습할 수 있습니다. 하지마 CNN의 단점 중에 하나가 feature들이 locally하다는 점입니다. global한 정보를 학습하기 위해서 self-attention 기반의 Vision Transformer(ViT)가 많이 사용되고 있는데 CNN에 비하면 ViT의 파라미터 수는 매우 많습니다. 따라서 본 논문에서는 이런 점들을 해결하기 위해!

mobile vision task를 위해 가볍고 latency가 짧으며 CNN과  ViT의 장점만을 결합한 신경망인 MobileViT를 설계했습니다.

본 논문에서 제안하는 MobileViT는 global 정보를 처리하는 데 있어 일반적인 Transformer와 다른 관점을 제시합니다. 그 결과 MobileViT가 CNN 및 ViT 보다 훨씬 뛰어난 성능을 보인다는 것을 다양한 실험을 통해 입증했습니다.

 

 

1. Introduction  

본 논문에서 제안하는 MobileViT는 소스가 제한된 모바일 기기에서 vision task를 실행할 수 있도록 ViT를 "가볍고, 빠르게" 만든 모델입니다. Transformer 계열의 신경망이 성능을 향상시키기 위해 선택한 방법은 파라미터 수를 증가시키는 것이었습니다. 하지만 이 방법은 모델 사이즈와 latency를 매우 증가시킨다는 단점이 있습니다.  최근까지도 비전분야에서 CNN이 활발히 사용되고 있는 만큼 CNN의 장점도 놓칠 수 없기 때문에 MobileViT는 CNN과 ViT의 장점을 결합하여 이러한 단점을 개선할 수 있도록 설계하였습니다.

- CNN: Spatial inductive bias를 가지며, data augmentation에 덜 민감합니다.

- Transformer : 모델이 무거우며, 최적화가 어렵습니다. 부족한   inductive bias문제와 over-fitting을 방지하기 위해  data augmentation이 필수적입니다. 또한, down stream task를 위해 무거운 decoder가 필수적입니다. 하지만 global information을 processing 할 수 있습니다. 

MobileViT는 장점들만 취해 locally하고 global한 정보 모두 encoding 할 수 있고, light weight와 low latency 모델입니다. 또한 성능면에서도 CNN보다 높은 성능을 달성했다고 합니다.

 

그렇다면 어떻게 CNN과 Transformer를 결합했을까요? 아주 간단하게만 살펴보겠습니다.

일반적인 CNN은unfolding -> local processing -> folding의 과정을 거치치만 MobileViT는 local processing부분을 transformer를 이용한 global processing으로 대체했다고 합니다. 때문에 이 부분은 CNN과 ViT의 특성을 모두 갖게되고 적은 파라미터수와 간단한 training recipe만으로 높은 성능을 달성할 수 있었습니다.

 

 

2. MobileViT: A Light-Weight Transformer  

논문에서는 MobileViT의 구조에 대해 설명하기 전에 ViT에 대해 간단히 설명하고 있습니다. ViT에 대한 설명은 아래 글을 참고하면 좋습니다 :)

VIT: https://ga02-ailab.tistory.com/147

2.1 MobileVit Architecture  

2.1.1 MobileViT block

MobileViT block의 구조는 위 그림과 같습니다. 이 부분의 역할은 적은 파라미터로 local, global 정보를 배울 수 있도록 하는 것입니다. 

먼저 input tensor 를 n*n convolution과 point-wise(1*1) convolution을 거치게 해 local spatial 정보를 학습합니다. 이렇게 해서 나온 tensor를 X_L 이라고 하겠습니다. 이때, long-range non-local dependency를 가지도록 학습됩니다. long-range란 거리가 먼 image patch들 끼리도 정보를 주고 받을 수 있도록 하는 것을 말합니다. 이를 위해 가장 많이 쓰이는 방법은 dilated-convolution입니다. 하지만 이 방법은 dilation rate를 잘 정해줘야 한다는 단점이 있습니다. 또 다른 방법이 바로 self-attention입니다. multi-head self-attention은 이미 그 효과가 여러차례 입증되었지만 반복적으로 설명했듯이 heavy-weight와 sub-standard optimizability 문제가 있습니다. 바로 이 부분 때문에 ViT가 spatial inductive bias가 부족한 것이죠.

MobileViT가 spatial inductive bias와 함께 global representatioin을 배우도록 하기 위해  X_L을 unfold 시켜 이미지와 동일하게 3채널의 non-overlapping flattend patch인 X_U를 만듭니다.  이때, X_U는 아래와 같이 표현됩니다. 

N=HW/P(patch 수)

P=wh(각 patch의 dimension)

patch들 사이의 relationship은 여러 개의 transformer를 거쳐 encoding 되게 되고, 이렇게 global information을 학습합니다. encoding된 output은 X_G로 표현되고 이 과정들을 수식으로 표현하면 아래와 같습니다.

X_U는 n*n크기의 local information 을 가지고 있는 patch들입니다. 최종 output인 X_G는 서로 다른 patch들 간의 global representation을 encoding함으로 X_G의 각 pixel은 모든 pixel의 정보를 encoding한다고 말할수 있습니다. 그림으로 살펴보겠습니다. 

- 빨간 pixel은 파란 pixel들과 정보를 공유합니다.

- n*n convolution을 통해 파란 pixel이 인접한 pixel들과 정보를 공유합니다. 

=> 그 결과, 빨간 pixel이 입력 이미지 전체 pixel과 정보를 공유할 수 있습니다. 

위 과정들 덕분에 MobileViT는 patch의 순서와 pixel의 공간정보도 잃지 않은채로 학습 할 수 있습니다. 그리고 최종 출력이었던 X_G를 다시 fold해 X_F를 얻습니다. X_L을 local, global 정보를 포함한 채로 원래의 차원으로 projection하기 위해 point-wise convolution 을 사용하고 MobileViT block의 입력 X와 concat한 후 n*n convolution을 통해 feature를 fuse합니다.

 

2.1.2 Relationship to convolution

standard convolution은  unfolding -> local representation을 위한  matrix multiplication -> folding 세 단계를 거칩니다.  MobileViT는 local processing단계를 보다 심층저인 global processing을 위해 transformer layer로 대체합니다. 그 결과, MobileViT는 convolution의 속성 중 하나인 spatial bias를 가지고 , 따라서 MobileViT block을 convolution transformer라고 말할 수 있습니다. 이러한 구조는 컴퓨터, 휴대폰 상관없이 다양한 기기에서 MobileViT를 사용 할 수 있다는 장점이 있습니다.

 

2.1.3 Light-weight

기존의 convolution과 transformer를 결합한 신경망들은 모두 파라미터 수가 많아 무거웠습니다. MobileViT는 어떻게 light-weight를 달성 할 수 있었을까요?

기존의 convolution과 transformer를 결합한 신경망을 먼저 살펴보겠습니다. 아래 그림처럼 기존의 신경망은 spatial information을 latent로 바꿉니다.  

transformer를 적용하는 단계에서 인접 pixel을 축으로 하여 stack하고 pixel값들을 linear projection을 이용해 latent space로 보내는 embedding 연산이 image의 specific inductive bias를 잃게 하는 것입니다. 

반면에, MobileViT는 convolution과 transformer를 두 개의 특성을 모두 살려 global representation을 배울 수 있도록 하기 때문에 light weight가 가능하답니다.

 

2.1.4 Computational cost

MobileViT와 ViT의 게산 복잡도를 비교해보면 각각 O(N^2P d) and O(N^2d)입니다. 언뜻보면 MobileViT가 더 복잡해보이지만 실제로는 MobileViT가 DeIT 모델보다 약 1/2배의 FLOPs를 가지고 ImageNet-1k 데이터셋에서 더 높은 정확도를 달성했습니다. 이것이 가능한 이유 역시 윗절에 설명한 것처럼 convolution과 transformer의 장점을 살려 결합했기 때문이라고 할 수 있습니다. 

 

2.1.5 MobileViT architecture

본 논문에서는 MobileViT 를 3가지의 다른 타입으로 설계했습니다. 모델의 사이즈에 따라 XXS, XS, S 세 종류가 존재합니다. 

 

2.2 Multi-Scale Sampler for Training Efficiency  

기존 ViT계열의 모델들이 사용하던 multi-scale representation 전략은 fine-tuning이었습니다. 다양한 scale 별로 네트워크를 fine-tuning합니다. 이러한 접근법은 positional encoding이 다양한 scale에 따라 interpolate 되어야 하고 모델의 성능이 이 interpolate 방식에 영향을 받을 수 받게 없게 됩니다. 즉, 최대한 다양한 scale을 입력으로 받아 학습해야 성능이 높아진다는 얘기입니다. 하지만 MobileViT 는 CNN과 비슷하게 동작하기 때문에 positional embedding이 필요없고 다시 말하면, fine-tuning 또한 필요없게 됩니다.

하지만 multi-scale training이 다양한 CNN의 성능향상에 효과적이라는 것은 이미 분명하기 때문에 완전히 버릴 수 없어 MobileViT 에도 변형된 multi-scale training 방법을 사용합니다. 

기존 CNN에서 사용하던 방법들은 미리 정해진 입력 사이즈 중 하나를 iteration마다 선택해 학습하는데, 이는 batch size가 가장 큰 입력 사이즈로 고정되어 버리기 때문에 작은 사이즈에서는 GPU 사용률이 떨어지는 단점이 있습니다. 그래서 MobileViT는 다음과 같은 전략을 사용합니다. 

먼저 아래 spatial resolution S가 정렬된 채로 주어집니다.

다음은 t번째 training iteratioin에서 각 GPU가 하나의 scale을 랜덤하게 선택하고 batch size는 아래 식으로 설정해줍니다. 

그 결과 작은 scale에서는 batch가 커지기 때문에 빠르게 학습할 수 있게됩니다. 아래 그림은 일반적인 방법과 multi-scale sampler를 비교한 것인데, 그림 (b)에서 알 수 있듯이  multi-scale sampler에서 모델 업데이트 수가 적어 전체 epoch time도 줄어드는 것을 확인할 수 있습니다. 

본 논문에서는 multi-scale sampler 값으로 {(160, 160),(192, 192),(256, 256),(288, 288),(320, 320)}을 사용하고 있습니다.

 

 

 

3. Experimental Results  

3.1 CNN 모델과의 비교

다른 light cnn보다 우수한 성능을 보이고 있습니다. 특히 (b)를 보면 가장 적은 파라미터수를 가지고 있음에도 가장 높은 정확도를 달성한 것을 확인할 수 있습니다. 또한 XXS, XS, S 중 가장 많은 파라미터를 가진 S도 다른 CNN들에 비하면 파라미터 수는 적고 정확도는 높습니다.

 

3.2 ViT 모델과의 비교

역시 MobileViT가 가장 높은 성능을 보이고 있습니다. 

 

3.3 Inference Speed

아이폰 12를 이용하여 측정한 결과입니다. MobileNet이 가장 빨랐고 transformer 계열 중에서는 MobileViT가 가장 빨랐습니다. 그런데 GPU에서는 MobileViT가 가장 느렸는데 이는 MobileViT가 다른 모델들보다 큰 입력 사이즈(256)를 사용하기도 하고 shallow+narrow 특성을 모두 가지기 때문이라고 합니다. MobileViT 연산들이 하드웨어에 최적화되게 구현된다면 MobileViT의 속도는 더 빨라질 것이라고 기대해볼 수 있을 것 같습니다.

 

다음은 MobileViT 모델들 끼리의 비교입니다.

3.4 Patch size에 따른 성능 변화

위 그림은 patch size에 따른 inference 속도와 정확도입니다. 같은 사이즈의 모델이라고 해도 patch size에 따라 변화가 크기때문에 값을 잘 설정해주는 것이 중요합니다. 

 

3.5 CNN의 kernel size와 patch size

아래 그림에서 n은 kernel의 크기를, h와 w는 patch의 크기에 해당하는 값입니다.

h와 w가 n보다 크면 하나의 patch에 있는 각 pixel들이 Convolution으로 부터 patch에 속한 모든 pixel 정보를 얻을 수 없기 때문에 local information representation능력이 떨어지게 됩니다. 이는 전체적인 정확도 하락으로 이어질수밖에 없습니다.