MobileViT | Apple提出更小、更快、更高性能的端侧实时Transformer类轻量型MobileViT

arXiv: https://arxiv.org/pdf/2110.02178.pdf

Abstract

轻量型CNN是端侧视觉任务的首选，其具有的空域归纳偏置使其能够以较少的参数进行强表达能能力学习。然而这些方案是空域局部性，而不具备全局表达能力，而Transformer则通过自注意力机制构建了全局表达建模能力。但是，Transformer是一种重参数方案。是否可以取CNN与ViTs各自之长为端侧视觉任务构建一种轻量型、低延迟网络呢 ？

为达到上述目的，本文提出一种用于端侧设备的轻量型通用ViT架构MobileViT，它提供了一种不同的全局处理的视角，即transformer as convolutions 。实验结果表明：在不同任务与数据集上，MobileViT显著优于CNN与ViT等方案。比如，MobileViT仅需6M参数量即可取得78.4%的top1精度 @ImageNet，比MobileNetV3高3.2%，比DeIT高6.2%；在COCO目标检测任务上，MobileViT取得了比MobileNetV3高5.7%的指标。

Method

下图a给出了标准Transformer结构示意图，由于这类方案忽视了空域归纳偏置，故他们需要更多的参数学习视觉表达能力。比如，相比DeepLabV3，DPT需要6倍参数方可取得相似性能。与此同时，相比CNN，这类方案优化难度大，对L2正则比较敏感，需要更多的数据增广以防止过拟合。

上图b则给出了本文所提MobileViT架构示意图，其核心思想在于通过Transformer as Convolution学习全局表达能力，进而使得我们可以同时集成卷积与Transformer两者各自的优势。

MobileViT Architecture

上面图示给出了MobileViT Block结构示意图，它旨在以更少的参数进行全局与局部信息建模。给定输入，MobileViT采用生成。为使MobileViT能够学习带空域归纳偏置的全局表达能力，MobileViT在上述两个卷积中间插入了一个Transformer操作：

首先，将展开成N个非重叠平展块，然后通过Transformer学习块内相关性，最后将所得到的合并为。

如上图所示，由于分别编码了局部与全局信息，故MobileViT的感受野为。

Relationship to convolutions 标准卷积可视作unfolding→matrix multiplication→folding等操作的堆叠，MobileViT采用了类似的操作，区别在于：MobileViT采用更深的全局处理(Transformer)替换了卷积中的局部处理(Matrix Multiplication)。因此，MobileViT具有与类卷积属性(比如空域偏置)，故将其视作Transformer as Convolutions 。

Light-weight 为什么说MobileViT轻量呢？MobileViT的结构设计方式使其具有类卷积属性，同时可以进行全局处理，这种处理机制使得我们可以设计浅而窄的MobileViT结构，故而更轻量。相比DeIT中的，MobileViT中，这就使得MobileViT更快()、更小()、性能更高()。

Computational  cost  MobileViT与ViT的计算复杂度分别为，理论上MobileViT效率更低。但实际上，MobileViT计算效率更高。我们认为原因与Light-weight部分相似。

MobileViT Architecture  本文结构基于MobileNetV2改进而来，而激活函数则为Swish。参考CNN，MobileViT中的，MobileViT中的MV2模块主要用于下采样。不同大小的MobileViT架构配置见下表。

Multi-scale Sampler for Training Efficiency

标准的ViT方案是通过微调学习多尺度表达，这是因为ViT中的位置嵌入需要根据输入尺寸进行插值，故模型性能还会与插值方法相关。类似CNN，MobileViT无需任何位置嵌入，故它可以在训练过程中从多尺度输入中受益。

为使得MobileViT能够学习多尺度表达且无需微调，同时提升训练效率，我们将多尺度训练方法扩展为可变尺寸batchsize机制，见上图。给定分辨率集合、batch b以及最大分辨率，在第t次迭代中我们随机采样并计算对应的batch:。因此，更大的batch对应更小的分辨率，这种处理机制可以促进更快训练。

相比Pytorch中的DDP，所提多尺度采样机制(具体实现建议查看原文)具有以下几个优势：

• 降低训练耗时

• 提升约0.5%性能

• 促使网络学习更好的多尺度表达能力

Experiments

上图给出了MobileViT与CNN的性能对比，从中可以看到：

• 在2.5M参数量大小方面，MobileViT性能比MobileNetV2gao5%，比ShuffleNetNetV2高5.4%，比MobileNetV3高7.4%；

• MobileViT具有比重参数量的模型(如ResNet、DenseNet、EfficientNet)更高的性能。

上图给出了MobileViT与ViTs的性能对比，从中可以看到：MobileViT仅需常规增广与更少的参数量即可取得更高的性能。比如：相比DeIT，MobileViT参数量少2.5倍，性能高2.6%。

上表给出了COCO检测任务上的性能对比，可以看到：

• SSDLite+MobileViT具有比SSDLite搭配其他轻量型模型更高的性能；

• SSDLite+MobileViT具有比SSD-300更高的性能，同时参数量大幅减少。

上表给出了VOC2021分割任务上的性能对比，可以看到

• DeepLabV3+MobileViT模型更小、性能更高 ；

• DeepLabV3+MobielViT具有与ResNet101相当性能，同时参数量少9倍。

上图给出不同任务、不同输入时模型的推理耗时-性能对比，可以看到：

• 相比更大块尺寸，更小块尺寸(即Config(2,2,2))模型具有更高精度；

• 在限定参数负载下，Config(8、4、2)的模型推理速度更快；

• 除了MobileViT-S外，其他所有模型均能够以30fps速度在端侧运行。

上表给出了不同方案的速度、参数量以及性能对比，可以看到：ViT方案要比MobileNetV2更慢 。

上面的图与表对多尺度采样的作用进行了对比说明，从中可以看到：

• 多尺度采样有助于降低泛化性能差异，多尺度采样训练可以提升0.5%性能；

• 多尺度采样是一种通用技巧，它不仅可以提升MobileViT，还可以提升ResNet50的性能。

个人思考

Transformer的优异性能已经在CV的各个领域得到了验证，同时GPU端侧的推理速度也并不慢(与CNN基本相当水平吧)，但轻量型、端侧推理方面一直尚未有非常好的进展。尽管在轻量型方面，Transformer+CNN的组合也有一些研究，比如MobileFormer，但其对比仍然是参数量，而不是实际推理速度，速度与参数量有关，但又不完全相关 ，或者说：参数量少并不意味着速度快 。

本文应该是首次将CNN+Transformer的组合做到了端侧实时，而且不局限于分类任务，同时还将检测与分割等纳入考虑范畴，实乃近期为数不多的好文之一。当然，速度比起来CNN来肯定还是要慢一些的。比如：3.5M参数量的MobileNetV2推理速度仅需0.92ms，精度为73.3；而MobileViT的参数量为2.3M，推理速度为7.28ms，精度为74.8。比较遗憾的是，MobileViT尚未开源，无法亲自体验其高性能、高推理速度。

从实现上来看，MobileViT属于Transformer向CNN的嵌入，而之前的paper大多侧重于CNN向Transformer嵌入。MobileViT在利用Transformer进行全局建模能力的同时还需要尽可能的降低计算量达到高推理速度目的，故在不同阶段仅使用1个MobileViT，避免过多的计算量。由于现有推理框架的底层优化机制问题，MobileViT的速度比起MobileNetV2来还是存在不小的差距。所以，在期待有更轻量的Transformer架构的同时，也期待有更好的适用于Transformer类架构的推理框架。