MobileViT 논문 리뷰

ICML 2022에 accept된 Apple직원분들의 논문인 MOBILEVIT: LIGHT-WEIGHT, GENERAL-PURPOSE,

AND MOBILE-FRIENDLY VISION TRANSFORMER 을 리뷰해보겠습니다!

 

1. Introduction

MobileViT는 mobile과 같은 하드웨어의 자원이 제한된 곳에서 사용할 수 있도록 만든 작고(Light-weight) 빠른(low-latency) ViT(Visual Transformer) 모델입니다. 위와 같은 성능을 도출하기위해 MobileViT는 CNN과 ViT의 장점을 결합하였다고 합니다. 

• CNN의 장점
  • spatial(local) inductive bias
  • data augmentation에 덜 민감
• ViT의 장점 
  • input-adaptive weighting
  • global processing (spatial inductive bias와 대비되는 특성)

Inductive bias란?? ML 모델에 대해 정확한 예측을 위해 사용하는 추가적인 가정임.
예를 들어 CNN은 convolution filter가 Window sliding을 하기 때문에 local한 영역에서 spatial한(공간적인) 정보를 뽑아내는데요. 이는 "Vistion task는 local한 영역에서 정보를 얻을게 많다"라고 inductive bias가 들어가게 되는 것입니다. 이러한 가정이 옳기 때문에 CNN은 효과적인 성능을 보이는 것입니다. RNN 또한 모델설계과정에서 sequential한 inductive bias가 들어간것입니다. 

그래서 MobileViT block은 local, global한 information모두 효과적으로 encoding할 수 있으므로 accuracy도 높고 latency도 낮습니다.

 

그럼 결합은 어떻게 했느냐?? 물론 기본적으로 CNN을 쓰기도 했지만 MobileViT는 기존 ViT와 ViT의 variants(다른 논문들)과 다르게 global representation을 학습하기 위해 기존의 convolution의 연산을 변경하였습니다. 기본적으로 standard 한 convolution은 unfolding, local processing, folding의 순차적 operation을 포함하지만 MobileViT block은 local processing 부분을 transformer를 이용한 global processing으로 대체하였습니다. 이를 통해 해당 block은 CNN의 특성과 ViT의 특성 모두 가지게 되고 이는 적은 parameter와 간단한 training recipe(basic augmentation)으로도 좋은 성능을 가져왔다고 말합니다.

 

2. MobileViT: A Light-weight Transformer

MobileViT의 구조를 알아보기 전에 ViT먼저 간단하게 알아보고 가죠!

2.1 ViT(Visual Transformer)

아래오 같이 ViT는 이미지 \( X \in \mathbb{R}^{  H \times W \times C }  \)가 sequential한 flatten된 patches \( X_f \in \mathbb{R}^{  N \times PC } \) 로 reshape됩니다. 그리고  \(X_f \)는 linear layer을 통해 \( X_p  \in \mathbb{R}^{ N \times d }\)로  linear projection되어 fixed된 \(d\)-dimension을 가지게 됩니다. 그리고 \(X_p \)을 입력으로 \(L\)개의 transformer block을 학습하게 됩니다.  \(C, H, W\)은 각각 image의 channel, height, width을 의미하고 \( P=wh \)은 임의의 height \(h\)와 \(w\)을 가지는 patch안의 image pixel을 말하고 \(N \)은 patch의 개수를 의미합니다. 

 

예를 들어 아래의 이미지가 15(H)x15(W)x(3)이라면 9개의 patch로 나눠지는 것이고 각 patch는 h = 15(H)/3 = 5, w= 15(W)/3 =5 를 가지며 flatten되어 5(h)*5(w) 25의 차원을 가지게 됩니다. 즉, \(X_f \in \mathbb{R}^{ 9(N) \times 25(P)* 3(C) } \)의 형태를 가지게 됩니다. 그리고 \( X_f \) (보라색 부분!) 는 linear layer의 weight와 곱해져서 \(d\) dimension으로 표현됩니다. linear projection을 통해 sequential한 정보를 학습할 수 있도록 positional encoding하는 것입니다.  (아래의 extra learnable class embedding은 class에 대한 정보를 학습할 parameter라고 인지하시면 됩니다. )

 

https://github.com/gupta-abhay/pytorch-vit

기본적으로 ViT는 CNN이 가지는 spatial inductive bias를 가지지 못합니다. 그래서 visual representation을 학습하기 위해서는 수많은 데이터를 필요로 하는 문제점을 발생시키죠.  이러한 문제점을 해결하기 위해 MobileViT를 제안하게 되는 것이죠!

 

2.2 MobileViT Architecture

MobileViT block

MobileViT block은 적은 parameter로 input tensor \(X \in \mathbb{R}^{ H \times W \times C} \)의 local, global information을 모두 학습할 수 있도록 하는 것을 목표로 합니다. 먼저, Input tensor을 입력으로 MobileViT는 \( n \times n \)의 convolution과 point-wise (or 1x1) convolution을 적용하여 local spatial information을 학습하게 되는 것이죠. Convolution울 적용한 Ouput은  \( X_L \in \mathbb{R}^{ H \times W \times d } \)이며  \( d > C \)을 만족합니다.

 

그리고 MobileViT이 long-range non-local dependencies를 가지게 하기 위한 모델링을 하게됩니다. 

long-range란? patch간의 position이 멀리 떨어져있어도 서로 간의 정보를 주고 받을 수 있도록 함. 반대로 RNN과 같이 단어(token)간의 거리가 멀 경우 서로의 정보를 교환 및 학습하기 힘든 것을 short-range라고 함.

long-range dependency modeling을 위해서는 가장 흔히 쓰는 방법 중 하나는 dilated convolution입니다. 하지만 이러한 방법은 dilation rate에 따라 성능이 크게 좌우되므로 좋지 않죠. 다른 좋은 solution으로는 self-attention방법론이 있는데 이 중에 하나가 ViT의 multi-head attention입니다. 하지만 기존의 ViT는 parameter도 많고 sub-standard optimizability를 가지게 되고 이는 ViT가 spatial inductive bias를 가지지 못하는 이유가 됩니다.  

 

그래서 저자들은  ViT의 multi-head attention으로 long range dependency를 만족시키고 더하여 MobileViT block을 제안하여 기 spatial inductive bias까지 가지도록 하게 하는 것을 목적으로 하는 것입니다.  그래서 \( X_L \)을 unfold시켜 image와 같은 3차원의 non-overlapping flattend patches \( X_U \in \mathbb{R}^{P \times N \times d } \)를 만들어 냅니다. \(N = \frac{HW}{P} \)는 patch의 수, \(P = wh \)는 각 patch의 dimension이며 \( h \leq n\)과 \( w \leq n\)을 만족합니다.  (위에서 n=3임을 기억!) 각 \( p \in \{ 1, ... , P\}\)에 대해, inner-patch간의 relationship은 \(L\)개의 transformer를 거쳐 encoding되면서 global information을 학습하게 됩니다. encoding된 ouput은 \(X_G \in \mathbb{R}^{ P \times N \times d} \)입니다. (기존 ViT와 다르게 positional encoding이 없다는 것을 확인!)

 

\[
X_G(p) = Transformer(X_U (p)), 1 \leq p \leq P , \quad \cdots Eq. (1)
\]

 

 \(n \times n\) convolution을 거쳐나온 \(X_U(p)\)는 \(n \times n\)크기의 local information을 encode하고 있으며 \(X_G(p)\)는 하나의 \(p\)-th location에 대해 다른 \(P\) patch들간의 global infromation을 encoding하고 있으므로 \(X_G(p)\)의 각 pixel은 \(X\)의 모든 pixel을 encoding한다고 말할 수 있으며 이는 MobileViT의 effective한 receptive field의 크기는 \(H \times W\)입니다.

그렇기 때문에 MobileViT는 각 patch의 patch order와 pixel의 spatial order 모두 잘 학습되는 것입니다. 그리고 다시 \(X_G\)를 fold시켜 \( X_F \in \mathbb{R}^{H \times W \times d} \)를 얻습니다. \(X_F \)는 다시 원래의 \(C\)-dimension으로 projection 시키기 위해 point-wise convoution을 사용하게 되고 해당 output은 \(X\)와 concatenation operation을 통해 결합됩니다. 이후 \(n \times n\) convolution을 사용해 concatenated된 features들을 Fusing하게 되죠. 

 

 

2.2.1 MobileViT architecture

MobileViT architecture

위는 MobileViT의 전체 구조입니다. MobileViT의 첫번째 layer는 strided 3x3 convolution을 사용하고 다음으로는 MobileNetv2 (or MV2) blocks과 MobileViT blocks을 사용합니다. activation function은 swish를 사용하였고 \(h = w =2\)로 설정하였습니다. 그리고 MobileViT는 3가지 다른 사이즈의 모델이 있습니다. S: small, XS: extra small, XXS: extra extra small입니다.

2.2.1 Light weight

기존의 ViT는 spatial inductive bias가 없으므로 model의 capacity를 올려 visual representation을 학습하도록 하였습니다. 하지만 MobileViT는 spatial inductive bias를 포함하므로 model의 capacity를 낮출 수 있습니다. 그래서 transformer layer \(L = \{2, 4, 3 \} \), dimension \(d = \{ 96, 120,144\}\)로 설정하였고 각 spatial level은 \(32 \times 32\), \(16 \times 16\), \(8 \times 8\)으로 셋팅하였습니다. (기존의 ViT-based 모델 DeIT는 \(L=12, d=192\)을 사용했었음)

2.3 Multi-scale Sampler for Training Efficiency

기존 ViT-based model은 multi-scale representation을 학습하기위해 fine-tuning을 진행했었죠. 이는 기존의 ViT모델의 positional embedding이 input size에 따라 interpolated되어야 하기때문입니다. (다양한 size의 input을 받아야 성능이 높아지는 구조라고 이해하시면 됩니다.) 하지만, MobileViT는 CNN과 비슷하므로 postional embedding이 필요없는 것이고 이는 Fine-tuning이 필요없다는 뜻입니다.

 

다만, multi-scale training은 기존 CNN의 성능에도 효과적이므로 저자들은 MobileViT에도 해당 방법론을 사용하게 됩니다. spatial resoltuion \( S = \{ (H_1, W_1), ... , (H_n, W_n) \} \)이 정렬된 채로 주어집니다. 그리고 \(t\)-th training iteration에서 각 GPU마다 하나의 spatial resoultion을 랜덤하게 sampling \( (H_t, W_t) \in S \)되고 batch size는 \(b_t = \frac{H_n W_n b}{H_t W_t}를 갖습니다. 그 결과 작은 spatial resolution을 가질경우 큰 batch size이 사용됩니다. 이는 optimization update를 감소시키므로 빠르게 training할 수 있게 합니다.

위의 그림은 standard한 sampler와 multi-scale을 비교한것인데 (b)에서 알 수 있듯이 update되는 횟수가 적으므로 epoch time도 적은것을 알 수 있습니다. 또한 Multi-scale sampler를 통해 0.5%의 성능 향상도 보았다고 합니다.

 

3. Experiment results

3.1 Experiment setting for classification

• ImageNet-1K dataset
• 300 epoch
• 1024 batch size with AdamW optimizer and label smoothing cross entropy(smoothing=0.1)
• multi-scale sampler \(S = \{ (160,160), (192,192), (256,256), (288,288), (320,320) \} \)
• learning rate from 0.0002 to 0.002 for the first 3k iteration, annealed to 0.0002 using consine scheduler
• 0.01 L2 weight decay 
• basic data augmentation (i.e. random resized cropping, horizontal flipping)

3.2 Experiment result for classification

3.3 Experiment setting for object detection

• MS-COCO dataset
• SSD-Lite의 backbone으로 MobileViT사용
  • 다른 backbone들과 성능 비교
  • SSD-Lite는 기존 SSD head의 conv를 separable convolution으로 바꾼 것
• 320 x 320 input resolution 
• AdamW optimizer
• smooth L1 (localization용) and cross entropy (classification용)

3.4 Experiment setting for object detection