[Paper Review] MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer

Sachin Mehta, Mohammad Rastegari|arXiv (Cornell University)|Oct 5, 2021

Advanced Neural Network Applications727 citations

TL;DR

MobileViT combines CNN-like local inductive bias with Transformer-based global processing to create a light-weight, mobile-friendly vision transformer that outperforms comparable CNNs and ViTs at similar parameter counts. It delivers strong ImageNet, detection, and segmentation performance with simple training recipes.

ABSTRACT

Light-weight convolutional neural networks (CNNs) are the de-facto for mobile vision tasks. Their spatial inductive biases allow them to learn representations with fewer parameters across different vision tasks. However, these networks are spatially local. To learn global representations, self-attention-based vision trans-formers (ViTs) have been adopted. Unlike CNNs, ViTs are heavy-weight. In this paper, we ask the following question: is it possible to combine the strengths of CNNs and ViTs to build a light-weight and low latency network for mobile vision tasks? Towards this end, we introduce MobileViT, a light-weight and general-purpose vision transformer for mobile devices. MobileViT presents a different perspective for the global processing of information with transformers, i.e., transformers as convolutions. Our results show that MobileViT significantly outperforms CNN- and ViT-based networks across different tasks and datasets. On the ImageNet-1k dataset, MobileViT achieves top-1 accuracy of 78.4% with about 6 million parameters, which is 3.2% and 6.2% more accurate than MobileNetv3 (CNN-based) and DeIT (ViT-based) for a similar number of parameters. On the MS-COCO object detection task, MobileViT is 5.7% more accurate than MobileNetv3 for a similar number of parameters. Our source code is open-source and available at: https://github.com/apple/ml-cvnets

Motivation & Objective

• Motivate the need for light-weight, mobile-appropriate vision transformers that blend CNN inductive biases with global processing.
• Propose MobileViT blocks that fuse local convolutional processing with patch-level transformer attention.
• Demonstrate that MobileViT delivers competitive accuracy with fewer parameters and simple training recipes across classification, detection, and segmentation tasks.

Proposed method

• Introduce MobileViT block that applies an n×n convolution for local processing followed by a 1×1 projection to dimension d, unfold X into patches, apply a transformer to each patch sequence to model inter-patch relations, and fold back to high-resolution features.
• Fuse the transformed features with the original via concatenation and a subsequent fusion convolution to obtain the output.
• Train MobileViT in three network sizes (S, XS, XXS) with a strided 3×3 stem, MV2 blocks for down-sampling, and Swish activations, using basic augmentation and an AdamW optimizer.
• Utilize a multi-scale sampler during training to vary input resolutions and batch sizes, improving efficiency and generalization.
• Compare against light-weight CNNs, ViT variants, and heavy-weight CNNs on ImageNet-1k, MS-COCO (SSD/SSDLite), and PASCAL VOC 2012 with standard evaluation metrics.

Experimental results

Research questions

• RQ1 Can a light-weight, mobile-friendly ViT model achieve CNN-like performance by combining local convolutions with global transformer processing?
• RQ2 Does MobileViT offer better generalization and training robustness than pure ViT variants and comparable CNNs at similar parameter budgets?
• RQ3 Can MobileViT serve as a general-purpose backbone for classification, detection, and segmentation tasks on mobile platforms?
• RQ4 What is the impact of multi-scale training and patch size choices on accuracy and latency on mobile devices?

Key findings

• On ImageNet-1k, MobileViT-S achieves 78.4% top-1 with 5.6M parameters, outperforming MobileNetv3 and DeIT at similar budgets.
• MobileViT-XS (~2.3M params) reaches 74.8% top-1, and the XXS/XS/S variants show favorable parameter-accuracy trade-offs compared to light-weight CNNs.
• In MS-COCO object detection, MobileViT-XS/S backbones improve mAP over MobileNetv3 by up to ~1.8% with smaller models.
• For semantic segmentation with DeepLabv3, MobileViT backbones yield 77.1% mIoU (MobileViT-XS) and 79.1% (MobileViT-S) with substantially fewer parameters than ResNet-101-based backbones.
• MobileViT backbones deliver better generalization and training robustness, with simple augmentation and less sensitivity to L2 regularization compared to many ViT variants.
• On mobile devices, MobileViT achieves real-time inference and shows faster-than-typical ViT backbones when considering hardware and kernel optimizations.