PaperlayerView exact revision on arXiv

Source-linked AI summary

Separable Self-attention for Mobile Vision Transformers

Sachin Mehta, Mohammad Rastegari

arXiv:2206.02680v1cs.CVcs.AIcs.LG

TL;DR

MobileViT models offer strong mobile-vision performance but suffer high latency because MHA has quadratic complexity and costly operations. The paper introduces separable self-attention with linear complexity and element-wise computation, producing MobileViTv2 models reported as faster and better across tasks, including 75.6% ImageNet top-1 accuracy with about three million parameters.

  • Problem

    Mobile vision transformers need lower-latency self-attention because MHA has O(k^2) complexity and costly operations that burden resource-constrained devices.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    Separable self-attention computes context relative to a latent token using two linear computations and element-wise operations, replacing MHA in MobileViT.

    Source: §3 · pp. 3–4, §3.2 · p. 4

  • Results

    MobileViTv2 is reported as faster and better than MobileViTv1 across tasks, with competitive segmentation performance and comparable inference speed to CNN-based models.

    Source: §4.1 · p. 7, §4.2 · p. 8

  • Takeaways & Limitations

    The method provides a fast, memory-efficient attention alternative for mobile vision transformers and can be integrated as a drop-in MHA replacement.

    Source: §3 · p. 3

Abstract

from arXiv · show

Mobile vision transformers (MobileViT) can achieve state-of-the-art performance across several mobile vision tasks, including classification and detection. Though these models have fewer parameters, they have high latency as compared to convolutional neural network-based models. The main efficiency bottleneck in MobileViT is the multi-headed self-attention (MHA) in transformers, which requires $O(k^2)$ time complexity with respect to the number of tokens (or patches) $k$. Moreover, MHA requires costly operations (e.g., batch-wise matrix multiplication) for computing self-attention, impacting latency on resource-constrained devices. This paper introduces a separable self-attention method with linear complexity, i.e. $O(k)$. A simple yet effective characteristic of the proposed method is that it uses element-wise operations for computing self-attention, making it a good choice for resource-constrained devices. The improved model, MobileViTv2, is state-of-the-art on several mobile vision tasks, including ImageNet object classification and MS-COCO object detection. With about three million parameters, MobileViTv2 achieves a top-1 accuracy of 75.6% on the ImageNet dataset, outperforming MobileViT by about 1% while running $3.2\times$ faster on a mobile device. Our source code is available at: \url{https://github.com/apple/ml-cvnets}

1 Introduction

Mobile vision transformers are limited by quadratic, operation-intensive self-attention on resource-constrained devices. The paper proposes separable self-attention with linear complexity and element-wise computation, and evaluates its effect on MobileViT models.

1 Introduction

  • MHA has O(k^2) complexity and uses costly batch-wise matrix multiplication, creating an inference bottleneck on resource-constrained devices.The bottleneck matters because such devices have reduced computational capabilities, restrictive memory constraints, and limited power budgets.Source: §1 · p. 1
  • Existing sparse-attention methods reduce complexity to O(k) or O(k log k) but incur performance drops.The supplied related-work passage identifies this trade-off for sparsity-based approaches.Source: §1 · p. 1
  • Linformer reduces self-attention complexity from O(k^2) to O(k) through low-rank projections but still uses costly operations.Its remaining batch-wise matrix multiplications may hinder deployment on resource-constrained devices.Source: §1 · pp. 1–2
  • Figure 1 compares attention-unit operations, complexity, and latency as token count changes, using CPU profiling with k = 256 in the single-layer operation comparison.The setup uses a single CPU core, d = 512, h = 8 for Transformer and Linformer, and p = 256 for Linformer.Source: Figure 1 · p. 2
  • Separable self-attention provides O(k) complexity and replaces computationally expensive attention operations with element-wise summation and multiplication.Experimental results on standard vision datasets and tasks are presented as evidence of effectiveness.Source: §1 · p. 2

2 Related work

Prior work improves transformer efficiency through locality, token grouping, token reduction, and training or optimization techniques. These approaches remain constrained by performance, sequence-length requirements, inference or memory costs, or their relationship to the proposed method.

Improving self-attention

  • Locality-based methods reduce complexity to O(k) or O(k log k) but can suffer large performance degradation with only moderate speed-ups.They restrict each token’s receptive field to a subset of tokens rather than attending globally.Source: §2 · p. 3
  • Similarity-based grouping methods such as Reformer reduce theoretical cost to O(k log k), but gains are noticeable mainly for sequences with k > 2048.Because ViTs use k < 1024, the passage says these approaches are not suitable for that setting.Source: §2 · p. 3
  • Existing methods remain limited in reducing inference time and memory consumption on resource-constrained devices.The paper motivates separable self-attention as fast and memory-efficient for this deployment setting.Source: §2 · p. 3

Improving transformer-based models

  • Token-reduction approaches improve transformer efficiency through methods including down-sampling and pyramidal structures.The proposed separable self-attention is described as a drop-in MHA replacement that can be integrated with transformer-based models.Source: §2 · p. 3

Other methods

  • Mixed-precision training, efficient optimizers, and knowledge distillation are orthogonal techniques, with mixed precision used by default during training.These methods address training or model performance rather than replacing the self-attention operation itself.Source: §2 · p. 3

3 MobileViTv2

MobileViTv2 replaces MobileViT's costly quadratic multi-headed self-attention with separable self-attention, which computes global information using linear operations and improves inference efficiency.

Motivation

  • MobileViT combines CNN and ViT strengths but has high latency, with MHA identified as its main efficiency bottleneck.MHA captures long-range token relationships but has quadratic complexity and costly matrix operations.Source: §3 · p. 3

Separable self-attention

  • Separable self-attention computes context scores relative to a latent token, then uses them to re-weight input tokens and produce a global context vector.The method replaces direct all-token interactions with two separate linear computations.Source: §3 · pp. 3–4, §3.2 · pp. 4–5, Figure 4 · p. 5
  • O(k) complexity replaces MHA's quadratic attention cost, while element-wise summation and multiplication replace costly batch-wise matrix multiplication.The design targets faster inference on resource-constrained devices.Source: §3 · p. 3, Figure 3 · p. 4, §3.2 · p. 5

MobileViTv2 architecture

  • MobileViTv2 is created by replacing MHA in MobileViT with separable self-attention.The resulting architecture serves as a drop-in integration of the proposed attention module.Source: §3 · p. 4, §3.3 · p. 6

Results

  • 3× inference-speed improvement is observed with similar ImageNet-1k performance when separable self-attention replaces MHA in MobileViT.The comparison uses models with similar parameter counts and FLOPs, with latency measured on iPhone12.Source: §3.2 · pp. 5–6, Table 1 · p. 6

4 Experimental results

Experiments evaluate MobileViTv2 across classification, segmentation, and detection settings, comparing accuracy, efficiency, and latency under different training and hardware conditions.

Image classification

  • Replacing MobileViTv1’s MHA with separable self-attention makes MobileViTv2 faster and better.Source: §4.1 · p. 7
  • MobileViTv2 is about 8× faster on a mobile device and delivers 2.5% better ImageNet-1k performance than MobileFormer.On GPU, both models run at a similar speed.Source: §4.1 · p. 7

Efficiency comparison

  • MobileViTv2 bridges the latency gap between CNN- and ViT-based models while maintaining performance with similar or fewer parameters.ConvNeXT remains 2× faster than MobileViTv2 for similar performance on a mobile device.Source: §4.1 · pp. 7–8
  • 1.6× to 1.3×: the GPU speed gap between ConvNeXT and MobileViTv2 decreases as input resolution increases from 224 × 224 or 256 × 256 to 384 × 384.The paper attributes this pattern to greater parallelism in ViT-based models at larger scales.Source: §4.1 · p. 8

Downstream tasks

  • MobileViTv2 delivers competitive semantic-segmentation performance with significantly fewer parameters and FLOPs, while inference speed is comparable to CNN-based models.The comparison covers ADE20k and PASCAL VOC 2012 using mIoU.Source: §4.2 · p. 8, Table 3 · p. 8
  • MobileViTv2 delivers competitive MS-COCO object-detection performance across models with different capacities.The models use SSDLite with MobileViTv2 ImageNet-1k backbones.Source: §4.2 · p. 8

5 Visualizations of self-separable attention scores

Context score maps visualize which image regions MobileViTv2 attends to at multiple output strides, highlighting semantically relevant regions across feature levels.

  • Separable self-attention layers attend to low-, mid-, and high-level features across output strides.Source: §5 · p. 8
  • Context score maps are shown at output strides 8, 16, and 32 alongside the input image.Output stride is the ratio of the input spatial dimension to the feature-map spatial dimension.Source: §5 · p. 8, Figure 5 · p. 9
  • The context scores attend to semantically relevant image regions, supporting representation learning from those regions.Source: §5 · p. 8, Figure 5 · p. 9

6 Conclusions

The paper concludes that separable self-attention provides a linear-complexity, hardware-friendly alternative to expensive MHA for transformer-based vision models on mobile devices.

  • Separable self-attention has linear complexity and uses hardware-friendly element-wise operations instead of expensive MHA computations.Experimental results on standard datasets and tasks demonstrate effectiveness over multi-headed self-attention.Source: §6 · p. 9
  • Transformer-based vision models are slow on mobile devices because multi-headed self-attention is expensive on resource-constrained devices.Source: §6 · p. 9
  • Experimental results on standard datasets and tasks demonstrate the effectiveness of separable self-attention over multi-headed self-attention.Source: §6 · p. 9

A Detailed architecture of MobileViTv2

MobileViTv2 follows MobileViTv1’s architecture but replaces MHA with separable self-attention and removes components whose performance improvements were marginal.

B MobileViTv2’s classification performance

MobileViTv2 classification performance improves with higher-resolution finetuning, ImageNet-21k-P pretraining, and increasing model size. It is presented as smaller and better than comparable lightweight models.

Architecture

  • The MobileViTv2 block uses 3 × 3 depth-wise convolution, 2 × 2 patch dimensions, and B× repetitions of separable self-attention and feed-forward layers.Table 5 identifies B as the repetition count for the transformer block inside the MobileViTv2 block.Source: Figure 6 · p. 13, Table 5 · p. 13

C Comparisons with light-weight networks on the ImageNet-1k dataset

MobileViTv2 combines convolutional and separable-attention processing while comparing favorably with lightweight CNN- and ViT-based networks across network sizes. Its context-score visualizations indicate aggregation across entire images and multiple representation levels.

Light-weight CNNs

Light-weight ViTs

MobileViTv2 block

  • The MobileViTv2 block unfolds the input feature map into patches before feeding it to separable self-attention to learn non-local representations.The unfolded feature map is denoted xu.Source: D Visualizations of separable… · p. 14

Context-score visualization

  • At output strides 8, 16, and 32, MobileViTv2 uses B = 2, B = 4, and B = 3 layers, producing 2, 4, and 8 context-score maps respectively.Source: Figure 8 · p. 15
  • Context scores are computed for M pixels simultaneously across N patches, then folded into spatial context-score maps and min-max normalized for visualization.Source: D Visualizations of separable… · p. 16
  • The visualizations show information aggregation from the entire image and learning of high-, mid-, and low-level representations.The reported settings include complex backgrounds, illumination and viewpoint changes, and different objects.Source: D Visualizations of separable… · p. 16

E MobileViTv2’s ablation studies on the ImageNet-1k dataset

Ablations examine augmentation, loss, latent-token count, and sampling choices for MobileViTv2 on ImageNet-1k. The reported findings favor configuration simplicity when alternatives add little accuracy or worsen mobile latency.

Latent tokens

Attention heads

Data augmentation

Sampling

F MobileViTv2 training configurations

The training configuration section specifies dataset-specific training and finetuning setups for MobileViTv2, including ImageNet models and downstream segmentation tasks. It also records a validation-set construction condition and an optimizer choice for Ade20k.

Training and finetuning

  • Training and finetuning configurations are provided for MobileViTv2-2.0 on ImageNet-1k and ImageNet-21k-P, and for downstream tasks.These configurations are listed in Tables 8, 9, and 10.Source: F MobileViTv2 training configurations · p. 17

Dataset configuration

  • The ImageNet-21k-P validation set does not overlap with the ImageNet-1k validation set and was created following Ridnik et al.Source: Table 8 · p. 17

Weight initialization

  • Table 9 uses an ImageNet-21k-P model finetuned on ImageNet-1k at 256 × 256 resolution to initialize weights.Source: Table 9 · p. 18

Downstream finetuning

  • SGD was used for Ade20k finetuning because it was more stable than AdamW across MobileViTv2 configurations.The PSPNet configurations match Deeplabv3 on PASCAL VOC and Ade20k.Source: Table 10 · p. 18

Related papers

Loading 2206.02680v1…