[MobileOne] An Improved One millisecond Mobile Backbone 논문 리뷰

2022년 6월 Apple에서 mobile friendly한 모델을 제안하는 논문인 An Improved One millisecond Mobile Backbone 을 리뷰해보도록 하겠습니다.

 

1. Introduction

저자들은 실제 mobile에서 latency(inference speed)를 최적화시키는 것을 목표로 하고 있으며 기존의 논문들이 제안했던 mobile-friendly하다는 모델들에 대해 아래와 같은 단점을 지적하였습니다.

• FLOPs와 parameter수가 (상대적으로) 작다고 하여 실제 latency가 빠르지 않음
• skip-connection과 branching은 많은 memory access cost를 발생시키므로 latency가 최적화되지 않음

skip connection & branching examples

그래서 저자들은 (1) mobile device에서 latency의 bottleneck 요인을 분석하고 이를 바탕으로 (2) mobile device에서 inference 속도를 최적화한 모델인 MobileOne을 제안하게 됩니다. MobileOne은 아래와 같이 SOTA accuracy와 latency performance를 도달하였으며 전에 제가 리뷰한 MobileViT-S보다 5배 빠르다고 하네요! latency cost는 CoreML을 이용하여 IPhone12 device에서 측정하였다고 합니다.

Comparisons of Top-1 accuracy on image classification vs latency on an iPhone 12

 

2. Method

2.1 Metric Correlations

모델의 efficiency를 측정하는 (기존의) metric인 parameter수와 FLOPs가 실제 latency와 correlation이 있는 지 확인하였습니다. 저자들은 PyTorch 모델들을 ONNX 형태로 1차 변환하였고 ONNX 모델을 CoreML tool을 이용해 coreml package형태로 2차 변환하여 IPhone12에서 latency를 측정하였습니다.

 

실제로 실험해보니 아래와 같이 대다수의 모델이 parameter수가 높은 반면에 latency는 낮게 측정되었습니다. FLOPs와 latency관계에서도 비슷한 경향을 보인다는 것도 확인하였습니다. 또한 MobileNet이 Transformer계열 모델들과 FLOPs와 parameter 수는 비슷하지만 latency가 낮은 것을 확인하였다고 합니다.  

또한 아래와 같이 Spearman rank correlation측정시에도 latency와 FLOPs는 어느 정도만의 correlation이 존재하고 parameter와 latency는 correlation이 적다는 것을 확인하였습니다. 뿐만 아니라 desktop CPU에서 측정하였을 때도 correlation이 낮게 나오는 것을 알 수 있습니다.

Spearman rank correlation coefficient between latency-flops

위 실험들을 통해 FLOPs(parameter수)와 latency는 correlation이 적다는 것을 확인하였습니다.

 

2.2 Key Bottlenecks

해당 섹션에서는 latency에 악영향을 미치는 network architecture 및 operation에 대해 분석합니다. 

2.2.1 Activation Fucnctions

30 layers구조의 CNN을 구성하였고 해당 CNN에서 activation function에 따른 latency 측정 실험을 하였습니다. 최근에 efficient하다고 제안된 SE-ReLU, Dynamic Shift-Max, DynamicReLUs들은 모두 hardware acceleration에 최적화되지 않아 latency가 높게 측정되었습니다. 그에 비해 ReLU와 GELU는 latency가 낮게 측정됨을 확인하였고 그래서 저자들은 MobileOne에 ReLU activation만 사용한다고 합니다.

Comparison of latency on mobile device of different activation functions

2.2.2 Architectural Blocks

Runtime performance에 큰 (악) 영향을 끼치는 요인은 memory access와 degree of parallelism입니다. memory access는 multi-branch network 구조인 경우에 크게 상승한다고 말합니다. 이유는 각 branch로부터 activation들이 다음 tensor를 연산하기 위해서 저장해야하기 때문입니다. 또한 Squeeze-Excitation(SE) block에서 사용되는 global pooling operation은 synchronization cost가 크기 때문에 runtime performance에 악영향을 줍니다.

 

그래서 위의 내용을 정량적으로 확인하기 위해 30 layers구조의 기본적인 CNN에서 skip-connection과 SE block을 추가했을 때 inference time을 측정하였습니다. 아래와 같이 skip-connection과 SE block이 추가되었을 때 latency가 높아지는 것을 확인하였으며 이를 바탕으로 MobileOne에서는 skip-connection을 완전히 제거하였고 SE block사용은 최대한 제한하였습니다. (SE block이 accuracy 향상에 좋다고 하여 쓴다고는 하네요)

Ablation on latency of different architectural blocks

2.3 MobileOne Architecture

MobileOne의 구조는 위의 분석을 통해 디자인되기도 하였지만 또 다른 중요한 구조는 re-parameterization이 가능한 형태를 사용하여 training-time과 inference-time구조가 다르다는 점입니다. 

 

2.3.1 MobileOne Block

아래 그림에서 보이듯이 RepVGG논문에서는 training시에는 multi-branch구조로 있다가 inference-time에는 여러 branch의 layer(또는 operation)을 하나의 conv layer로 reparameterization하게 됩니다. 이렇게 reparameterization하게 되면 하나의 con layer로 연산 가능하기 때문에 inference time에는 memory access을 많이 하는 multi-branch구조를 제거시킵니다. 

Structural re-parameterization of a RepVGG block

RepVGG에서는 reparameterization을 Standard Conv layer에 적용하였는데요. MobilONE은 RepVGG와 다르게 MobileNet-V1 의 block인 3x3 depthwise conv와 1x1 pointwise conv로 block을 구성하였습니다. 해당 block으로 RepVGG처럼 reparameterizable한 skip-connection추가 된 형태로 training하게 되고 testing때는 reparameterization하게됩니다.

 

MobileOne block has two different structures at train time and test time

그리고 저자들은 over-parameterization branch를 제안하였는데 위 그림과 같이 over-parameterization factor \(k\) 추가하여 해당 (depthwise or pointwise) conv-BN을 연속해서 \(k\)개 쌓도록 하였습니다.  Conv-BN이 여러 개 쌓인 구조는 사이에 비선형 함수인 activation이 없기 때문에 결국 linearlity성격을 가지므로 하나의 conv로 치환(fusing) 가능하기 때문에 저자들은 이와 같이 over-parameterization방법을 제안하게 되었습니다. \(k\)는 hyperparameter로 1~5 사이의 값을 가집니다. 

 

Conv-BN이 어떻게 folding(fusing) 될 수 있는지 수학적으로 알아봅니다. 

• \( W' \in \mathcal{R}^{C_{out} \times C_{in} \times K \times K}, b' \in \mathcal{R}^D \): Weight and bias
  • \(C_{out}\) ,\(C_{in}\): Output and input channel dimension
  • \(K\): Kernel size

Conv-BN layer는 linear operations으로만 구성되어 있기 때문에 위의 Conv layer의 weight \(W'\)가 아래와 같이 하나의 Conv weight로 folding가능합니다.

• \( \widehat{W} = W' * \frac{\gamma}{\sigma}, \widehat{b} = ( b' - \mu ) * \frac{\gamma}{\sigma} + \beta \)

결국, inference시에 모든 branch안의 존재하는 batch norm layer는 이전 Conv layer에 folding됩니다.

 

또한 skip connection(multi-branch)은 RepVGG의 reparameterization방법을 통해 folding된다. 즉, 1x1 conv은 \(K-1\) zero padding이 추가되어 \(K \times K\) Conv과 똑같은 shape으로 만들어 여러 skip connection(multi-branch)는 하나의 Conv layer weight로 합(summation)쳐진다. 수식으로 표현하면 \(W = \sum^M_i \widehat{W}_i, b= \sum^M_i \widehat{b}_i \) 와 같고 \(M\)은 branch의 수를 의미합니다.

 

저자들이 주장하는 Reparameterization과 over-parameterization의 효과를 주장하기 위해 아래와 같은 실험을 진행하였습니다.

Effects for  Reparameterization & over-parameterization

Table 6을 통해 over-parameterization factor \(k\)에 변화에 따른 성능 비교를 하였을 때 MobileOne의 scale이 큰 경우(MobileOne-S1) \(k\)이 커짐에 따라 큰 성능 향상 효과를 보지 못했지만 MobileOne의 scale이 작은 경우에는(MobileOne-S1) \(k\)가 4일때 0.5%의 성능향상 효과를 보았습니다.

Table 7을 통해서는 reparameterizable한 branch를 추가하였을 때 model의 scale variant에 상관없이 성능이 더 높게 나왔다는 결과를 보여주고 Figure 4에서는 multi-branch를 사용하고 k가 4일 경우  train, validation loss가 가장 낮게 학습된다는 실험 결과를 보여주었다. 위의 3개의 실험결과를 통해 저자들은 주장한 방법론에 대한 성능 검증을 하였습니다.

 

2.3.2 Modle Scaling

저자들은 width scale에 따라 5개의 MobileOne 버전을 아래와 같이 제안하였다. 

MoblieOne network specifications

 

2.4 Training

작은 모델일수록 regularization이 크면 underfitting되기 때문에 regularization을 약하게 주어야 합니다. 그래서 저자들은 regularization term인 weigt decay coefficient의 값을 annealing하게 됩니다. (weight decay을 아예 없애는 것이 아닌!) 저자들은 learning rate에 cosine annealing scheduler를 사용할 뿐만 아니라 weight decay에 똑같은 scheduling을 적용하였다고 합니다. 추가적으로 (Exponential Moving Average)EMA와 progressive learning curriculum 방법을 추가하여 성능을 아래와 같이 높였다고 합니다.

EMA: training parameter들이 moving average를 유지시키는 방법으로 원래 잘 학습하던 방향에서 크게 벗어나지 않도록 함
Progressive learning: epoch이 증가할수록 input image size가 커지며 augmentation강도가 커지게 하여 점점 어려운 이미지의 특성을 학습하도록 하는 방법론

Ablation on various train settings for MobileOne-S2 showing Top-1 accuracy on Imagenet

 

2.5 Benchmarking

iPhone12으로 latency를 측정하는 데 있어서 모델 그 자체만 실행되는 시간을 측정하는 것은 불가능하다고 합니다. 그래도 최대한 모델의 latency를 정확히 측정하기 위해 아래와 같은 셋팅으로 실험 진행하였습니다.

• Benchmarking하기 전에 model graph을 load하고 input tensor를 preallocate하였으며 model warmup을 한번 진행함
• Benchmarking시에는 모델을 1000번 run하여 lowest, highest latency를 측정함
• Benchmarking시에는 다른 모든 application은 종료하였음
• 모든 모델의 latency 측정 결과는 lowest latency임
  • lowest latency가 다른 process의 interrupt를 받지 않았다고 가정하기 때문에 해당 결과로 표를 작성함

그리고 Desktop CPU latency도 측정하였는데 이는 Ubuntu desktop(2.3 GHz Intel Xeon Gold 5118 processor)에서 결과를 도출하였습니다. 

 

3. Experiments

3.1 Image Classification on ImageNet-1K

실험셋팅은 다음과 같습니다.

• 300 epochs
• 256 batch sizes
• SGD optimizer with momentum
• Cross entropy loss with label smoothing (0.1 smoothing factor)
• Initial learning rate 0.1 with cosine annealing scheduler
• Initial weight decay 1e-4 and annealed to 1e-5 with cosine annealing scheduler
• AutoAugmentation for MobileOne-S2, 3, 4
• Standard augmentation for MobileOne-S0, 1
• EMA weight averaging with decay constant of 0.9995 

MobileOne의 성능 비교 및 결과는 아래와 같습니다. CNN기반의 모델과 Transformer계열 네트워크 모두와 비교하였습니다.

Performance of various models on ImageNet-1k validation set

 

Comparison of Top-1 Accuracy on ImageNe

3.2 Object detection on MS-COCO

SSDLite의 backbone을 MobileOne으로 변경하여 성능 비교를 하였습니다. Input resolution은 320x320이며 200 epochs 학습하였다고 합니다. (자세한 training configuration은 논문 참고 바랍니다!) 성능 측정은 MS COCO validation set을 사용하였으며 mAP@IoU 0.50:0.05:0.95 기준으로 하였습니다. 아래 backbone에 따른 성능 비교 테이블을 확인하였을 때 MobileOne의 퍼포먼스가 가장 뛰어남을 확인 가능합니다. 

Quantitative performance of object detection on MS-COCO

3.3 Semantic Segmentation on Pascal VOC and ADE 20k

Deeplab V3의 backbone을 MobileOne으로 변경하여 semantic segmentation성능 비교를 하였습니다. VOC와 ADE 20k dataset에 대해 모두 50 epochs 학습하였으며 자세한 training configuration은 논문 참고 바랍니다) 성능 측정은 mean intersection-over-union (mIOU)기준으로 하였으며 아래와 같이 MobileOne이 좋은 성능을 내었습니다.

Quantitative performance of semantic segmentation on Pascal-VOC and ADE20k datasets