PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sachin Mehta, Mohammad Rastegari

arXiv:2110.02178v2cs.CVcs.AIcs.LG

TL;DR

Mobile vision needs models that combine CNN efficiency and spatial inductive biases with transformer global processing without ViT-scale cost and complexity. MobileViT implements this combination by using transformers as convolutions, and it reports stronger performance across mobile vision tasks and datasets. Its reported limitation is that it remains slower than MobileNetv2 on mobile devices.

  • Problem

    Mobile vision requires light-weight, low-latency, accurate, general-purpose models, but CNNs are spatially local while ViTs are heavy-weight and costly to deploy.

  • Method

    MobileViT learns global representations with transformers as convolutions, implicitly incorporating convolution-like spatial bias while integrating with downstream architectures.

  • Results

    MobileViT significantly outperforms CNN- and ViT-based networks across tasks and datasets, including 78.4% ImageNet-1k top-1 accuracy with about 5–6 million parameters and +1.8% mAP with a 1.8× smaller SSDLite detection network.

  • Takeaways & Limitations

    MobileViT provides a light-weight, general-purpose vision transformer that combines CNN-like properties with global transformer processing and simple training recipes.

  • Takeaways & Limitations

    MobileViT and other ViT-based networks are slower than MobileNetv2 on mobile devices.

Abstract

from arXiv · show

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

1 INTRODUCTION

MobileViT addresses the challenge of combining CNN spatial inductive biases with transformer global processing in a light-weight, general-purpose, low-latency model for mobile vision. It reports stronger accuracy with modest parameter budgets and simpler training recipes across classification and detection.

  • Motivation: ViTs improve global representation learning but are heavy-weight, difficult to optimize, and costly for downstream dense prediction compared with light-weight CNNs.ViT-B/16 versus MobileNetv3 has 86 versus 7.5 million parameters, while a ViT-based segmentation network has about 345 million parameters versus 59 million for DeepLabv3.
  • Design goal: Mobile vision models must be light-weight, low latency, accurate, general-purpose, and compatible with device resource constraints.FLOPs alone do not determine mobile latency because memory access, parallelism, and platform characteristics also matter.
  • Approach: MobileViT replaces local convolutional processing with global transformer processing, combining CNN-like spatial biases with ViT-like global representations.The approach is presented as transformers as convolutions and supports simple training recipes such as basic augmentation.
  • Results: 78.4% top-1 accuracy is achieved on ImageNet-1k with about 5–6 million parameters, 3.2% higher than MobileNetv3.MobileViT uses 300 epochs and batch size 1024 versus 600 epochs and batch size 4096 for MobileNetv3.
  • Results: +1.8% mAP and a 1.8× smaller detection network result when MobileViT replaces MNASNet as the SSDLite feature backbone.The comparison is reported on the MS-COCO dataset.

2 RELATED WORK

Prior work shows that light-weight CNNs are versatile and easy to train but spatially local, while ViTs capture global representations yet are difficult to optimize. MobileViT combines these properties and is reported to improve performance and generalization with simple training recipes.

  • Light-weight CNNs: Light-weight CNNs use efficient convolutional designs that are versatile, easy to train, and suitable for replacing heavy backbones in task-specific models.Their major drawback is spatial locality.
  • Vision transformers: ViTs can reach CNN-level performance with large datasets or extensive augmentation, regularization, and distillation, but remain difficult to train.The related work characterizes ViTs as having substandard optimizability compared with CNNs.
  • Hybrid models: Hybrid convolution-transformer models are robust and high-performing, but combining their strengths into light-weight mobile networks remains an open question.MobileViT is introduced to target light-weight, general-purpose, mobile-friendly vision transformers.
  • MobileViT observations: MobileViT models achieve better performance than existing light-weight CNNs across different mobile vision tasks for a given parameter budget.This claim is listed as one of MobileViT’s reported observations.
  • MobileViT observations: MobileViT shows better generalization capability than previous ViT variants and is described as robust to training hyper-parameters such as data augmentation and L2 regularization.Generalization is defined using the gap between training and evaluation metrics.

3 MOBILEVIT: A LIGHT-WEIGHT TRANSFORMER

MobileViT combines convolutional local processing with transformer-based global processing while preserving spatial structure. Its lightweight design targets mobile deployment and supports multi-scale training.

  • MobileViT architecture: MobileViT learns global representations with transformers as convolutions while retaining convolution-like spatial bias.The block replaces convolutional local processing with deeper global transformer processing.
  • MobileViT architecture: MobileViT unfolds convolutional features into patches, applies transformers across corresponding patch locations, then folds them back without losing spatial order.Local information is encoded before transformer-based inter-patch processing.
  • Light-weight design: MobileViT uses standard convolutions and transformers to learn local and global representations, respectively, with a lightweight architecture across three network sizes.The networks use S, XS, and XXS variants built from MobileNetv2 and MobileViT blocks.
  • Multi-scale sampler: The multi-scale sampler reduces training time, improves performance by about 0.5%, and yields better multi-scale representations than standard sampling.It uses variably sized batches and fewer optimizer updates.

4 EXPERIMENTAL RESULTS

Experiments evaluate MobileViT on classification, detection, segmentation, and mobile-device inference. Across these settings, it outperforms comparable lightweight backbones, while mobile latency remains affected by device-level transformer optimization.

  • Image classification: MobileViT outperforms lightweight CNNs across network sizes and also exceeds several heavier CNNs on ImageNet-1k.At about 2.5 million parameters, it outperforms MobileNetv2 by 5%, ShuffleNetv2 by 5.4%, and MobileNetv3 by 7.4%.
  • Image classification: MobileViT models are easy and robust to optimize, supporting application to new tasks and datasets.The reported classification comparisons use models trained from scratch without distillation.
  • Mobile object detection: 1.8% higher performance and 1.8× smaller model size are achieved when SSDLite uses MobileViT instead of MNASNet on MS-COCO detection.MobileViT also outperforms SSDLite with other lightweight CNN backbones at 320 × 320 input resolution.
  • Mobile semantic segmentation: 1.4% better performance and 1.6× smaller size are reported when DeepLabv3 uses MobileViT instead of MobileNetv2 for segmentation.MobileViT is competitive with ResNet-101 while requiring 9× fewer parameters.
  • Mobile devices: MobileViT inference is measured on an iPhone 12, where larger patch configurations are faster but less accurate than smaller patch configurations.The comparison attributes the speed difference to fewer patches and greater parallelism.
  • Mobile devices: MobileViT and other ViT-based networks are slower than MobileNetv2 on mobile devices because optimized transformer operations are unavailable.The resulting mobile inference graph is described as sub-optimal.

A MOBILEVIT ARCHITECTURE

MobileViT combines MobileNetv2 blocks with MobileViT transformer blocks in a lightweight CNN-inspired architecture. Its design uses local down-sampling and reduced transformer dimensions for mobile-oriented models.

  • A MOBILEVIT ARCHITECTURE: MobileViT begins with a strided 3 × 3 convolution, followed by MobileNetv2 blocks and MobileViT blocks.MobileNetv2 blocks mainly perform down-sampling, while the architecture uses Swish activations.
  • A MOBILEVIT ARCHITECTURE: MobileViT blocks use convolution kernel size n = 3 and patch dimensions h = w = 2 at all spatial levels.The patch dimensions are chosen because feature-map dimensions are usually multiples of 2 and satisfy h, w ≤ n.
  • A MOBILEVIT ARCHITECTURE: The first feed-forward layer in each transformer layer has output dimension 2d instead of the standard 4d.Here, d is the transformer layer’s input dimension.

B MULTI-SCALE SAMPLER

The multi-scale sampler improves MobileViT’s generalization and robustness across input resolutions and model sizes. Its reported benefits also extend to CNNs, including improved performance and training efficiency.

  • B MULTI-SCALE SAMPLER: Multi-scale sampling lowers the validation error and nearly closes the training–validation error gap for MobileViT-S.The model trained with multi-scale sampling also shows greater robustness when evaluated at different input resolutions.
  • B MULTI-SCALE SAMPLER: About 0.5%: multi-scale sampling improves MobileViT performance across different model sizes.The improvement is reported from comparisons using standard and multi-scale samplers.
  • B MULTI-SCALE SAMPLER: Multi-scale sampling improves both performance and training efficiency for ResNet-50 and MobileNetv2-1.0.The experiments use basic augmentation on ImageNet-1k.
  • B MULTI-SCALE SAMPLER: About 1.4% and 14%: MobileNetv2-1.0 gains performance while reducing training time by 14%.The passage reports these changes for MobileNetv2-1.0 trained with the multi-scale sampler.

C ABLATIONS

The ablations examine regularization, skip connections, and patch sizing in MobileViT. They report robustness to weight decay, a measurable skip-connection benefit, and a latency–accuracy rationale for small patches.

  • C ABLATIONS: MobileViT-S remains robust across weight-decay values from 0.1 to 0.0001, except at 0.1.The experiments use 0.01, and at 0.0001 MobileViT reports 77.4 top-1 accuracy with 5.7 M parameters versus DenseNet’s 76.2 with 14 M parameters.
  • C ABLATIONS: 0.5%: adding the MobileViT block’s skip connection improves MobileViT-S performance on ImageNet.Without the connection, MobileViT-S still delivers similar or better performance than the cited CNN- and ViT-based models with basic augmentation.
  • C ABLATIONS: When h, w ≤ n, MobileViT aggregates information more effectively, which improves performance.The patch-size study uses n = 3 and compares configurations across spatial levels.

LS EMA Top-1

The supplied passages describe the effects of label smoothing and exponential moving average on MobileViT-S top-1 performance. They also note that patch resizing can add latency.

  • LS EMA Top-1: Patch dimensions that do not divide feature-map dimensions require bilinear interpolation during folding and unfolding.The added resizing operations are associated with latency in the patch-size experiments.
  • LS EMA Top-1: Label smoothing marginally improves MobileViT-S performance on ImageNet-1k, while exponential moving average has little or no effect.The comparison is reported in Table 7 for MobileViT-S.

D TRAINING DETAILS FOR SSDLITE AND DEEPLABV3

SSLite-MobileViT and DeepLabv3-MobileViT use specified training schedules, optimization settings, and task-specific resolution handling. DeepLabv3 models avoid multi-scale sampling because their dilation-based representations are resolution-dependent.

  • SSDLite-MobileViT and DeepLabv3-MobileViT are trained for 200 and 50 epochs, respectively, using a standard sampler and effective batch size of 128.Training uses four NVIDIA GPUs.
  • The learning rate rises from 0.00009 to 0.0009 during the first 500 iterations, then follows cosine annealing back to 0.00009 with L2 weight decay of 0.01.
  • Multi-scale sampling is omitted because DeepLabv3’s dilation rates operate on padded zeros at lower resolutions, making multi-scale learning ineffective.The cited example contrasts 512 × 512 with 256 × 256 inputs.

E EXTENDED DISCUSSION

The extended discussion evaluates MobileViT through memory, FLOPs, cross-task efficiency, device latency, and implementation effects. Results show strong efficiency and accuracy, but performance depends on hardware and accelerator optimization.

  • Memory footprint: At output strides 8, 16, and 32, MobileViT blocks require memory that is lower than or comparable to lightweight CNNs.The authors therefore characterize MobileViT networks as memory efficient.
  • FLOPs: For equal FLOPs, MobileViT is 1.45× faster, 2.1× smaller, and 1.8% more accurate than PiT.The comparison uses R2 versus R4 in Table 9.
  • Downstream tasks: Across classification, detection, and segmentation, MobileNetv2 is faster but less accurate, while the FLOPs gap narrows as input resolution increases.On semantic segmentation, MobileNetv2 and MobileViT have similar FLOPs.
  • Inference time on different devices: On iPhone CPU and neural engine, MobileViT outperforms DeIT and PiT, whereas DeIT and PiT are faster on GPU.The authors associate this difference with architecture, input resolution, and unavailable GPU-accelerated folding and unfolding operations.
  • Implementation effects: Replacing unoptimized folding and unfolding with PyTorch operations reduces MobileViT latency from 0.62 ms to 0.47 ms.
  • Discussion: The findings identify opportunities to optimize ViT-based models, including MobileViT, for different accelerators.

F QUALITATIVE RESULTS ON THE TASK OF OBJECT DETECTION

Qualitative detection results indicate that SSDLite-MobileViT-S detects objects across varied visual conditions. The reported settings include changes in illumination, viewpoint, background, and nonrigid deformation.

  • SSDLite-MobileViT-S detects objects under changes in illumination and viewpoint.
  • SSDLite-MobileViT-S detects objects across different backgrounds and nonrigid deformations.
  • Figures 14, 15, and 16 present object detection results of SSDLite-MobileViT-S on the MS-COCO validation set.

G SEMANTIC SEGMENTATION RESULTS ON AN UNSEEN DATASET

DeepLabv3-MobileViT is evaluated on an unseen MS-COCO validation set after mapping MS-COCO classes to PASCAL VOC classes. It outperforms the comparable DeepLabv3-MobileNetv3-Large backbone while using fewer parameters.

  • The evaluation uses 5k unseen MS-COCO validation images, maps MS-COCO classes to PASCAL VOC classes, and reports mIOU.
  • DeepLabv3-MobileViT-S is 1.7× smaller and 5.1% more accurate than DeepLabv3-MobileNetv3-Large.The comparison uses the same DeepLabv3 segmentation model.
  • Figures 17 and 18 show qualitative DeepLabv3-MobileViT-S segmentation results with input images, predicted masks, and mask overlays.The figures also show color encoding for object classes in the PASCAL VOC dataset.
Loading 2110.02178v2…