MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer

1 INTRODUCTION . Self-attention-based models , especially vision transformers ( ViTs ; Figure 1a ; Dosovitskiy et al. , 2021 ) , are an alternative to convolutional neural networks ( CNNs ) to learn visual representations . Briefly , ViT divides an image into a sequence of non-overlapping patches and then learns interpatch representations using multi-headed self-attention in transformers ( Vaswani et al. , 2017 ) . The general trend is to increase the number of parameters in ViT networks to improve the performance ( e.g. , Touvron et al. , 2021a ; Graham et al. , 2021 ; Wu et al. , 2021 ) . However , these performance improvements come at the cost of model size ( network parameters ) and latency . Many real-world applications ( e.g. , augmented reality and autonomous wheelchairs ) require visual recognition tasks ( e.g. , object detection and semantic segmentation ) to run on resource-constrained mobile devices in a timely fashion . To be effective , ViT models for such tasks should be light-weight and fast . Even if the model size of ViT models is reduced to match the resource constraints of mobile devices , their performance is significantly worse than light-weight CNNs . For instance , for a parameter budget of about 5-6 million , DeIT ( Touvron et al. , 2021a ) is 3 % less accurate than MobileNetv3 ( Howard et al. , 2019 ) . Therefore , the need to design light-weight ViT models is imperative . Light-weight CNNs have powered many mobile vision tasks . However , ViT-based networks are still far from being used on such devices . Unlike light-weight CNNs that are easy to optimize and integrate with task-specific networks , ViTs are heavy-weight ( e.g. , ViT-B/16 vs. MobileNetv3 : 86 vs. 7.5 million parameters ) , harder to optimize ( Xiao et al. , 2021 ) , need extensive data augmentation and L2 regularization to prevent over-fitting ( Touvron et al. , 2021a ; Wang et al. , 2021 ) , and require expensive decoders for down-stream tasks , especially for dense prediction tasks . For instance , a ViT-based segmentation network ( Ranftl et al. , 2021 ) learns about 345 million parameters and achieves similar performance as the CNN-based network , DeepLabv3 ( Chen et al. , 2017 ) , with 59 million parameters . The need for more parameters in ViT-based models is likely because they lack image-specific inductive bias , which is inherent in CNNs ( Xiao et al. , 2021 ) . To build robust and high-performing ViT models , hybrid approaches that combine convolutions and transformers are gaining interest ( Xiao et al. , 2021 ; d ’ Ascoli et al. , 2021 ; Chen et al. , 2021b ) . However , these N 3P Flatten image patches Linear N d ⊕ Positional encoding Transformer L× Linear Logits ( a ) Standard visual transformer ( ViT ) MobileViT block Local representations Transformers as Convolutions ( global representations ) Fusion W H C X Conv-n× ny Conv-1x1 W H d P N d Transformer L× P N d W H d Conv-1x1 W H C W H 2C Conv-n× n W H C Y Unfold Fold Conv-3× 3 ↓ 2 MV2 MV2 ↓ 2 MV2 2× MV2 ↓ 2 MobileViT block L = 2 h = w = 2 MV2 ↓ 2 MobileViT block L = 4 h = w = 2 MV2 ↓ 2 MobileViT block L = 3 h = w = 2 Conv-1× 1 Global pool→ Linear Logits Output spatial→ dimensions 128× 128 64× 64 32× 32 16× 16 8× 8 1× 1 ~w ( b ) MobileViT . Here , Conv-n× n in the MobileViT block represents a standard n × n convolution and MV2 refers to MobileNetv2 block . Blocks that perform down-sampling are marked with ↓ 2 . Figure 1 : Visual transformers vs. MobileViT hybrid models are still heavy-weight and are sensitive to data augmentation . For example , removing CutMix ( Zhong et al. , 2020 ) and DeIT-style ( Touvron et al. , 2021a ) data augmentation causes a significant drop in ImageNet accuracy ( 78.1 % to 72.4 % ) of Heo et al . ( 2021 ) . It remains an open question to combine the strengths of CNNs and transformers to build ViT models for mobile vision tasks . Mobile vision tasks require light-weight , low latency , and accurate models that satisfy the device ’ s resource constraints , and are general-purpose so that they can be applied to different tasks ( e.g. , segmentation and detection ) . Note that floating-point operations ( FLOPs ) are not sufficient for low latency on mobile devices because FLOPs ignore several important inferencerelated factors such as memory access , degree of parallelism , and platform characteristics ( Ma et al. , 2018 ) . For example , the ViT-based method of Heo et al . ( 2021 ) , PiT , has 3× fewer FLOPs than DeIT ( Touvron et al. , 2021a ) but has a similar inference speed on a mobile device ( DeIT vs . PiT on iPhone-12 : 10.99 ms vs. 10.56 ms ) . Therefore , instead of optimizing for FLOPs1 , this paper focuses on designing a light-weight ( §3 ) , general-purpose ( §4.1 & §4.2 ) , and low latency ( §4.3 ) network for mobile vision tasks . We achieve this goal with MobileViT that combines the benefits of CNNs ( e.g. , spatial inductive biases and less sensitivity to data augmentation ) and ViTs ( e.g. , input-adaptive weighting and global processing ) . Specifically , we introduce the MobileViT block that encodes both local and global information in a tensor effectively ( Figure 1b ) . Unlike ViT and its variants ( with and without convolutions ) , MobileViT presents a different perspective to learn global representations . Standard convolution involves three operations : unfolding , local processing , and folding . MobileViT block replaces local processing in convolutions with global processing using 1MobileViT FLOPs can be further reduced using existing methods ( e.g. , DynamicViT ( Rao et al. , 2021 ) ) . transformers . This allows MobileViT block to have CNN- and ViT-like properties , which helps it learn better representations with fewer parameters and simple training recipes ( e.g. , basic augmentation ) . To the best of our knowledge , this is the first work that shows that light-weight ViTs can achieve light-weight CNN-level performance with simple training recipes across different mobile vision tasks . For a parameter budget of about 5-6 million , MobileViT achieves a top-1 accuracy of 78.4 % on the ImageNet-1k dataset ( Russakovsky et al. , 2015 ) , which is 3.2 % more accurate than MobileNetv3 and has a simple training recipe ( MobileViT vs. MobileNetv3 : 300 vs. 600 epochs ; 1024 vs. 4096 batch size ) . We also observe significant gains in performance when MobileViT is used as a feature backbone in highly optimized mobile vision task-specific architectures . Replacing MNASNet ( Tan et al. , 2019 ) with MobileViT as a feature backbone in SSDLite ( Sandler et al. , 2018 ) resulted in a better ( +1.8 % mAP ) and smaller ( 1.8× ) detection network ( Figure 2 ) . 2 RELATED WORK . Light-weight CNNs . The basic building layer in CNNs is a standard convolutional layer . Because this layer is computationally expensive , several factorization-based methods have been proposed to make it light-weight and mobile-friendly ( e.g. , Jin et al. , 2014 ; Chollet , 2017 ; Mehta et al. , 2020 ) . Of these , separable convolutions of Chollet ( 2017 ) have gained interest , and are widely used across state-of-the-art light-weight CNNs for mobile vision tasks , including MobileNets ( Howard et al. , 2017 ; Sandler et al. , 2018 ; Howard et al. , 2019 ) , ShuffleNetv2 ( Ma et al. , 2018 ) , ESPNetv2 ( Mehta et al. , 2019 ) , MixNet ( Tan & Le , 2019b ) , and MNASNet ( Tan et al. , 2019 ) . These light-weight CNNs are versatile and easy to train . For example , these networks can easily replace the heavyweight backbones ( e.g. , ResNet ( He et al. , 2016 ) ) in existing task-specific models ( e.g. , DeepLabv3 ) to reduce the network size and improve latency . Despite these benefits , one major drawback of these methods is that they are spatially local . This work views transformers as convolutions ; allowing to leverage the merits of both convolutions ( e.g. , versatile and simple training ) and transformers ( e.g. , global processing ) to build light-weight ( §3 ) and general-purpose ( §4.1 and §4.2 ) ViTs . Vision transformers . Dosovitskiy et al . ( 2021 ) apply transformers of Vaswani et al . ( 2017 ) for large-scale image recognition and showed that with extremely large-scale datasets ( e.g. , JFT-300M ) , ViTs can achieve CNN-level accuracy without image-specific inductive bias . With extensive data augmentation , heavy L2 regularization , and distillation , ViTs can be trained on the ImageNet dataset to achieve CNN-level performance ( Touvron et al. , 2021a ; b ; Zhou et al. , 2021 ) . However , unlike CNNs , ViTs show substandard optimizability and are difficult to train . Subsequent works ( e.g. , Graham et al. , 2021 ; Dai et al. , 2021 ; Liu et al. , 2021 ; Wang et al. , 2021 ; Yuan et al. , 2021b ; Chen et al. , 2021b ) shows that this substandard optimizability is due to the lack of spatial inductive biases in ViTs . Incorporating such biases using convolutions in ViTs improves their stability and performance . Different designs have been explored to reap the benefits of convolutions and transformers . For instance , ViT-C of Xiao et al . ( 2021 ) adds an early convolutional stem to ViT . CvT ( Wu et al. , 2021 ) modifies the multi-head attention in transformers and uses depth-wise separable convolutions instead of linear projections . BoTNet ( Srinivas et al. , 2021 ) replaces the standard 3×3 convolution in the bottleneck unit of ResNet with multi-head attention . ConViT ( d ’ Ascoli et al. , 2021 ) incorporates soft convolutional inductive biases using a gated positional self-attention . PiT ( Heo et al. , 2021 ) extends ViT with depth-wise convolution-based pooling layer . Though these models can achieve competitive performance to CNNs with extensive augmentation , the majority of these models are heavy-weight . For instance , PiT and CvT learns 6.1× and 1.7× more parameters than EfficientNet ( Tan & Le , 2019a ) and achieves similar performance ( top-1 accuracy of about 81.6 % ) on ImageNet1k dataset , respectively . Also , when these models are scaled down to build light-weight ViT models , their performance is significantly worse than light-weight CNNs . For a parameter budget of about 6 million , ImageNet-1k accuracy of PiT is 2.2 % less than MobileNetv3 . Discussion . Combining convolutions and transformers results in robust and high-performing ViTs as compared to vanilla ViTs . However , an open question here is : how to combine the strengths of convolutions and transformers to build light-weight networks for mobile vision tasks ? This paper focuses on designing light-weight ViT models that outperform state-of-the-art models with simple training recipes . Towards this end , we introduce MobileViT that combines the strengths of CNNs and ViTs to build a light-weight , general-purpose , and mobile-friendly network . MobileViT brings several novel observations . ( i ) Better performance : For a given parameter budget , MobileViT models achieve better performance as compared to existing light-weight CNNs across different mobile vision tasks ( §4.1 and §4.2 ) . ( ii ) Generalization capability : Generalization capability refers to the gap between training and evaluation metrics . For two models with similar training metrics , the model with better evaluation metrics is more generalizable because it can predict better on an unseen dataset . Unlike previous ViT variants ( with and without convolutions ) which show poor generalization capability even with extensive data augmentation as compared to CNNs ( Dai et al. , 2021 ) , MobileViT shows better generalization capability ( Figure 3 ) . ( iii ) Robust : A good model should be robust to hyper-parameters ( e.g. , data augmentation and L2 regularization ) because tuning these hyper-parameters is time- and resource-consuming . Unlike most ViT-based models , MobileViT models train with basic augmentation and are less sensitive to L2 regularization ( §C ) .

Summary: This work proposes a hybrid backbone which combines existing MV2 block and new MobileViT block for efficient classification, detection and segmentation. The core idea with MobileViT block is to use both convolution and self-attention block to aggregate local and global feature, respectively. Experiments on IN-1K, COCO and PASCAL VOC have shown quite promising results.