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

TL;DR

MobileViT addresses the need for mobile-friendly vision models by fusing CNN-style spatial inductive biases with transformer-based global reasoning in a light-weight block. The MobileViT block performs local convolutions and global patch-level self-attention, enabling CNN-like efficiency alongside ViT-level global representations. The authors introduce a Multi-Scale Sampler to improve training efficiency and generalization across resolutions. Empirical results show MobileViT achieves state-of-the-art accuracy among lightweight models on ImageNet-1k and strong performance as a general-purpose backbone for object detection and semantic segmentation, with favorable real-time performance on mobile devices. The work highlights hardware-aware design considerations and points to future hardware optimizations to further close the gap with conventional CNNs on mobile platforms.

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

Figures (18)

• Figure 1: Visual transformers vs. MobileViT
• Figure 2: MobileViT shows better task-level generalization properties as compared to light-weight CNN models. The network parameters are listed for SSDLite network with different feature extractors (MobileNetv1 howard2017mobilenets, MobileNetv2 sandler2018mobilenetv2, MobileNetv3 howard2019searching, MNASNet tan2019mnasnet, MixNet tan2019mixconv, and MobileViT (Ours)) on the MS-COCO dataset.
• Figure 3: MobileViT shows similar generalization capabilities as CNNs. Final training and validation errors of MobileNetv2 and ResNet-50 are marked with $\star$ and $\circ$, respectively (§ \ref{['sec:appendix_msc']}).
• Figure 4: Every pixel sees every other pixel in the MobileViT block. In this example, the red pixel attends to blue pixels (pixels at the corresponding location in other patches) using transformers. Because blue pixels have already encoded information about the neighboring pixels using convolutions, this allows the red pixel to encode information from all pixels in an image. Here, each cell in black and gray grids represents a patch and a pixel, respectively.
• Figure 5: Multi-scale vs. standard sampler.