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.

  • Method

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

  • Results

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

  • 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.

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.
  • 1 Introduction: 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.
  • 1 Introduction: 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.
  • 1 Introduction: 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.
  • 1 Introduction: 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.

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.
  • Improving self-attention: 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.
  • Improving self-attention: 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.
  • 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.
  • 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.

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.
  • 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.
  • Separable self-attention: 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.
  • 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.
  • 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.

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.
  • Image classification: 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.
  • 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.
  • Efficiency comparison: 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.
  • 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.
  • Downstream tasks: MobileViTv2 delivers competitive MS-COCO object-detection performance across models with different capacities.The models use SSDLite with MobileViTv2 ImageNet-1k backbones.

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.
  • 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.
  • The context scores attend to semantically relevant image regions, supporting representation learning from those regions.

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.
  • Transformer-based vision models are slow on mobile devices because multi-headed self-attention is expensive on resource-constrained devices.
  • Experimental results on standard datasets and tasks demonstrate the effectiveness of separable self-attention over multi-headed self-attention.

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.

  • MobileViTv2 replaces multi-headed self-attention with separable self-attention to learn global representations.
  • MobileViTv2 removes the fusion block and skip-connection because they improve performance only marginally.
  • The architecture follows MobileViTv1, with an expansion factor of two in MobileNetv2 blocks and feed-forward layers.

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.

  • About 2%: MobileViTv2-0.50 with 1.4 million parameters improves when finetuned at 384 × 384 rather than the lower-resolution setting.The comparison is R1 versus R2 in Table 6.
  • 1.2%: ImageNet-21k-P pretraining improves MobileViTv2-2.0 performance.The comparison is R17 versus R18 in Table 6.
  • MobileViTv2 models pretrained on ImageNet-21k-P achieve similar performance with fewer FLOPs than models finetuned on ImageNet-1k at higher resolution.The passage cites R10 versus R11, R14 versus R15, and R18 versus R19.
  • MobileViTv2 performance on the ImageNet-21k-P validation dataset improves as model size increases.
  • 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.

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: MobileViTv2 outperforms MobileNetv1, MobileNetv2, ShuffleNetv2, ESPNetv2, and MobileNetv3 across different network sizes.
  • Light-weight ViTs: MobileViTv2 achieves better performance than DeiT, T2T, CrossViT, LocalViT, ConViT, and Mobile-former across different network sizes.
  • 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.
  • 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.
  • Context-score visualization: Context scores are computed for M pixels simultaneously across N patches, then folded into spatial context-score maps and min-max normalized for visualization.
  • Context-score visualization: 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.

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: Changing latent tokens from 1 to 8 produces negligible ImageNet-1k gains within ±0.1 top-1 accuracy, so experiments use one latent token.
  • Attention heads: Reducing MobileViTv1-S multi-headed self-attention heads from 4 to 1 drops top-1 accuracy by 0.7%.
  • Data augmentation: For models below 4.5 million parameters, standard augmentation works best, whereas models at least 4.5 million parameters benefit from advanced augmentation.
  • Sampling: The study evaluates random, top-magnitude, and uniform pixel- or patch-level sampling methods.Top-m and top-n selections use magnitude computed with the L2 norm.
  • Sampling: Pixel- and patch-sampling methods reduce FLOPs by 1.2× to 1.6× with little or no ImageNet-1k top-1 accuracy drop.The finding applies to both MobileViTv1 and MobileViTv2 models.
  • Sampling: Despite lower FLOPs, sampling methods are significantly slower on mobile devices because they increase tensor memory-access cost by changing memory order.Consequently, these methods are not used in MobileViTv2.

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.
  • Dataset configuration: The ImageNet-21k-P validation set does not overlap with the ImageNet-1k validation set and was created following Ridnik et al.
  • Weight initialization: Table 9 uses an ImageNet-21k-P model finetuned on ImageNet-1k at 256 × 256 resolution to initialize weights.
  • 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.
Loading 2206.02680v1…