[论文解读] MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer

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

Advanced Neural Network Applications被引用 727

一句话总结

MobileViT 将 CNN 风格的局部归纳偏置与基于 Transformer 的全局处理相结合，打造一个轻量、适合移动端的视觉 Transformer，在参数量相近的情况下 outperform 对应的 CNN 和 ViT。它在 ImageNet、检测和分割任务上提供强劲的性能，训练方案简单。

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

研究动机与目标

• 激发对轻量、适合移动端的视觉 Transformer 的需求，这类模型将 CNN 的归纳偏置与全局处理相融合。
• 提出 MobileViT 块，将局部卷积处理与补丁级 Transformer 注意力融合。
• 证明 MobileViT 在分类、检测和分割任务中以更少的参数实现具有竞争力的准确性，且训练方案简单。

提出的方法

• 引入 MobileViT 块，先应用 n×n 卷积进行局部处理，再进行 1×1 投影到维度 d，将 X 展开为补丁序列，对每个补丁序列应用 Transformer 以建模补丁间关系，再折叠回高分辨率特征。
• 通过拼接与随后的融合卷积将变换后的特征与原始特征融合，以获得输出。
• 在三种网络尺寸（S、XS、XXS）中训练 MobileViT，使用步幅为 3×3 的干线（stem），用于下采样的 MV2 块，以及 Swish 激活，采用基本的数据增强和 AdamW 优化器。
• 在训练期间使用多尺度采样器，以改变输入分辨率和批量大小，提升效率和泛化能力。
• 在 ImageNet-1k、MS-COCO（SSD/SSDLite）和 PASCAL VOC 2012 上，使用标准评估指标，与轻量 CNN、ViT 变体和重量级 CNN 进行比较。

实验结果

研究问题

• RQ1 一个轻量、移动友好的 ViT 模型是否能通过将局部卷积与全局 Transformer 处理相结合，达到接近 CNN 的性能？
• RQ2 在相近的参数预算下，MobileViT 是否比纯 ViT 变体和可比 CNN 拥有更好的泛化和训练鲁棒性？
• RQ3 MobileViT 能否作为移动平台上用于分类、检测和分割任务的通用骨干网络？
• RQ4 多尺度训练和补丁大小选择对移动设备上的准确性和延迟有何影响？

主要发现

• 在 ImageNet-1k 上，MobileViT-S 以 5.6M 参数达到 78.4% 的 top-1，超越在相似预算下的 MobileNetv3 和 DeIT。
• MobileViT-XS（约 2.3M 参数）达到 74.8% top-1，XXS/XS/S 变体相较于轻量 CNN 显示出有利的参数-精度权衡。
• 在 MS-COCO 目标检测中，MobileViT-XS/S 骨干相比 MobileNetv3 在更小模型下将 mAP 提升约 1.8%。
• 在使用 DeepLabv3 的语义分割中，MobileViT 骨干分别取得 77.1% mIoU（MobileViT-XS）和 79.1%（MobileViT-S），参数远少于基于 ResNet-101 的骨干。
• MobileViT 骨干提供更好的泛化和训练鲁棒性，具有简单的数据增强，对 L2 正则化的敏感度较少，相较于许多 ViT 变体。
• 在移动设备上，MobileViT 实现实时推理，并在考虑硬件和内核优化时，表现出比典型 ViT 骨干更快的速度。