MOBILEVIT: LIGHT-WEIGHT, GENERAL-PURPOSE, AND MOBILE-FRIENDLY VISION TRANSFORMER

摘要

轻量级卷积神经网络(CNN)是实现移动视觉任务的基础。他们的空间归纳偏向使他们能够在不同的视觉任务中以更少的参数学习表征。然而，这些网络在空间上是局部的。为了学习全局表征，采用了基于自我注意的视觉转换器(VITS)。与CNN不同，VITS是重量级的。在本文中，我们提出了以下问题：是否有可能将CNNS和VITS的优势结合起来，为移动视觉任务构建一个轻量级、低延迟的网络？为此，我们推出了MobileViT，这是一款适用于移动设备的轻量级通用视觉转换器。MobileViT为使用转换器进行全局信息处理提供了不同的视角。我们的结果表明，在不同的任务和数据集上，MobileViT的性能明显优于基于CNN和VIT的网络。在ImageNet-1k数据集上，MobileViT在大约600万个参数下实现了78.4%的TOP-1准确率，对于类似数量的参数，MobileViT分别比MobileNetv3(基于CNN)和Deit(基于VIT)的准确率高3.2%和6.2%。在MS-Coco对象检测任务中，对于相同数量的参数，MobileViT的准确率比MobileNetv3高5.7%。我们的源代码是开源的，可在以下位置获得：https://github.com/apple/ml-cvnets.

1 引言

基于自我注意的模型，特别是视觉转换器(VITS；图1a；Dosovitski等人，2021年)是学习视觉表征的卷积神经网络(CNN)的替代方案。简而言之，VIT将图像划分为一系列不重叠的图像块，然后使用变压器中的多头自我注意学习图像块间的表示(Vaswani等人，2017)。总的趋势是增加VIT网络中的参数数量以提高性能(例如，Touvron等人，2021a；Graham等人，2021年；Wu等人，2021年)。然而，这些性能改进是以模型大小(网络参数)和延迟为代价的。许多真实世界的应用(例如，增强现实和自动轮椅)需要视觉识别任务(例如，对象检测和语义分割)在资源受限的移动设备上及时运行。为了有效，这种任务的VIT模型应该是轻便和快速的。即使VIT模型的模型大小被缩小以匹配移动设备的资源限制，其性能也明显低于轻量级CNN。例如，对于大约500-600万的参数预算，Deit(Touvron等人，2021A)的精确度比MobileNetv3(Howard等人，2019年)低3%。因此，设计轻型VIT模型势在必行。

轻量级的CNN已经为许多移动视觉任务提供了动力。然而，基于VIT的网络还远未用于此类设备。与易于优化并与特定任务网络集成的轻量级CNN不同，VITS是重量级的(例如，VIT-B/16vs.MobileNetv3：86vs.750万个参数)，更难优化(肖等人，2021年)，需要广泛的数据增强和L2正则化以防止过度拟合(Touvron等人，2021A；Wang等人，2021年)，并且需要昂贵的解码器来执行下行任务，尤其是密集预测任务。例如，基于VIT的分段网络(Ranftl等人，2021年)学习约3.45亿个参数，并实现与基于CNN的网络DeepLabv3(Chen等人，2017年)相似的性能，具有5900万个参数。在基于VIT的模型中需要更多的参数可能是因为它们缺乏CNN固有的图像特定的感应偏差(肖等人，2021年)。为了构建健壮且高性能的VIT模型，将卷积和变压器相结合的混合方法正在获得兴趣(肖等人，2021年；达阿斯科利等人，2021年；Chen等人，2021年)。然而，这些混合模型仍然是重量级的，对数据增强敏感。例如，删除CutMix(钟等人，2020)和Deit风格(Touvron等人，2021A)数据增强会导致Heo等人的ImageNet准确率显著下降(78.1%至72.4%)。(2021年)。

将CNN和转换器的优点结合起来，为移动视觉任务建立VIT模型，仍然是一个悬而未决的问题。移动视觉任务需要轻量级、低延迟和满足设备资源限制的准确模型，并且是通用的，以便它们可以应用于不同的任务(例如，分割和检测)。请注意，浮点运算(Flop)不足以在移动设备上实现低延迟，因为Flop忽略了几个重要的推断相关因素，如内存访问、并行度和平台特征(Ma等人，2018年)。例如，Heo等人基于VIT的方法。(2021)，PIT的FLOPS次数比Deit少3倍(Touvron等人，2021A)，但在移动设备上的推理速度相似(Deit与iPhone-12上的PIT：10.99毫秒对10.56毫秒)。因此，本文没有针对FLOPs1进行优化，而是专注于为移动视觉任务设计一个轻量级(？3)、通用(？4.1和？4.2)和低延迟(？4.3)的网络。我们通过MobileViT实现了这一目标，它结合了CNN(例如，空间感应偏差和对数据增强不那么敏感)和VITS(例如，输入自适应加权和全局处理)的优点。具体地说，我们引入了MobileViT块，它将本地和全局信息有效地编码在一个张量中(图1b)。与VIT及其变体(带卷积和不带卷积)不同，MobileViT提供了一个不同的角度来学习全局表示。标准卷积包括三种操作：展开、局部处理和折叠。MobileViT block用变压器进行全局处理，取代卷积中的局部处理。这允许MobileViT块具有类似CNN和ViT的属性，这有助于它以更少的参数和简单的训练配方（例如，基本增强）学习更好的表示。据我们所知，这是第一项显示轻量VIT可以通过不同移动视觉任务中的简单培训配方实现轻量CNN级性能的工作。对于约500万至600万的参数预算，MobileViT在ImageNet-1k数据集（Russakovsky等，2015）上实现了78.4%的前1位准确度，比MobileNetv3准确度高3.2%，并且具有简单的训练配方（MobileViT与MobileNetv3:300对600个时代；1024对4096批量）。当MobileViT在高度优化的移动视觉任务特定体系结构中用作功能骨干时，我们还观察到性能的显着提高。用MobileViT作为SSDLite中的特征骨干（Sandler等，2018）替换MNASNet（Tan等，2019）导致更好（+1.8%mAP）和更小（1.8×）检测网络（图2）。

图 1 这里，MobileViT块中的Conv-n×n表示标准的n×n卷积，而MV2指的是MobileNetv2块。执行降采样的块标记为↓2。

图 2 与轻量级CNN模型相比，MobileViT具有更好的任务级泛化特性。列出了在MS-Coco数据集上具有不同特征提取器(MobileNetv1(Howard等人，2017)、MobileNetv2(Sandler等人，2018)、MobileNetv3(Howard等人，2019)、MNASNet(Tan等人，2019)、MixNet(Tan&Le，2019b)和MobileViT(我们的))的SSDLite网络的网络参数。

2 相关工作

轻量CNN。CNN中的基本构造层是标准卷积层。由于这一层的计算成本很高，已经提出了几种基于因式分解的方法来使其轻量级和对移动友好(例如，jin等人，2014；Chollet，2017；Mehta等人，2020)。其中，Chollet(2017)的可分离卷积已经引起了人们的兴趣，并被广泛用于移动视觉任务的最先进的轻量级CNN，包括MobileNet(Howard等人，2017；Sandler等人，2018；Howard等人，2019)、ShuffleNetv2(Ma等人，2018)、ESPNetv2(Mehta等人，2019)、MixNet(Tan&Le，2019b)和MNASNet(Tan等人，2019)。这些轻量级的CNN功能多样，易于训练。例如，这些网络可以很容易地替换现有任务特定模型(例如DeepLabv3)中的重量级主干(例如ResNet(他等人，2016))，以减小网络规模并改善延迟。尽管有这些好处，但这些方法的一个主要缺点是它们在空间上是局部的。这项工作将变压器视为卷积；允许利用卷积(例如，通用和简单的培训)和变压器(例如，全局处理)的优点来构建轻量级(§3)和通用(§4.1和§4.2)VITS。

视觉变形金刚。Dosovitski等人。(2021)应用Vaswani等人的变压器。(2017)用于大规模图像识别，并表明对于极大规模的数据集(例如JFT-300M)，VITS可以在没有图像特定感应偏差的情况下实现CNN级别的精度。通过广泛的数据扩充、大量的L2正则化和蒸馏，VITS可以在ImageNet数据集上进行训练，以达到CNN级别的性能(Touvron等人，2021a；b；周等人，2021)。然而，与CNN不同的是，VITS表现出不合标准的优化能力，而且很难训练。随后的工作(例如，Graham等人，2021；Dai等人，2021；Liu等人，2021；Wang等人，2021；袁等人，2021b；Chen等人，2021b)表明，这种不符合标准的可优化性是由于VITS中缺乏空间诱导偏向。在VITS中使用卷积结合这种偏置可以提高它们的稳定性和性能。为了获得卷积和变压器的好处，已经探索了不同的设计。例如，肖等人的VIT-C。(2021)在VIT中增加了一个早期卷积词干。CVT(Wu等人，2021)修改了变压器中的多头注意，并使用沿深度可分离的卷积而不是线性投影。僵尸网络(SRinivas等人，2021年)用多头注意取代了ResNet瓶颈单元中的标准3×3卷积。Convit(d‘Ascoli等人，2021年)利用门控位置自我注意结合了软卷积感应偏向。PIT(Heo等人，2021年)利用基于深度卷积的池层扩展了VIT。虽然这些模型可以通过广泛的扩展达到与CNN相当的性能，但这些模型大多是重量级的。例如，在ImageNet1k数据集上，PIT和CVT学习的参数分别比EfficientNet(Tan&Le，2019a)多6.1×和1.7倍，而性能相近(TOP-1准确率约为81.6%)。此外，当这些模型被缩小以建立轻型VIT模型时，它们的性能明显低于轻型CNN。对于约600万的参数预算，PIT的ImageNet-1k精度比MobileNetv3低2.2%。

讨论。与普通的VITS相比，卷积和变压器相结合产生了坚固和高性能的VITS。然而，这里一个悬而未决的问题是：如何结合卷积和变压器的优势，为移动视觉任务构建轻量级网络？本文的重点是设计轻量级的VIT模型，它的性能优于最先进的模型，具有简单的训练配方。为此，我们引入了MobileViT，它结合了CNNS和VITS的优点，以构建一个轻量级、通用和移动友好的网络。MobileViT带来了几个新的观察结果。(I)更好的性能：对于给定的参数预算，与现有轻量级CNN相比，MobileViT模型在不同的移动视觉任务中实现更好的性能(§4.1和§4.2)。(2)推广能力：推广能力是指培训和评价指标之间的差距。对于两个训练指标相近的模型，评估指标较好的模型具有更好的泛化能力，因为它可以在未知的数据集上进行更好的预测。与CNN相比(Dai等人，2021年)，以前的VIT变体(带卷积和不带卷积)不同，显示出较差的泛化能力，即使进行了大量的数据增强，而MobileViT显示出更好的泛化能力(图3)。(Iii)稳健性：一个好的模型应该对超参数(例如，数据增强和L2正则化)具有稳健性，因为调整这些超参数是耗时和耗费资源的。与大多数基于VIT的模型不同，MobileViT模型使用基本增强进行训练，并且对L2正则化不那么敏感(§C)。

图 3MobileViT显示出与CNN类似的泛化能力。

3 MOBILEVIT：一种轻型变压器

图1a所示的标准VIT模型将输入的X∈R^(H×W×C)重塑为一系列平面片(指patches) X_f∈R^{(N×PC)，将其投影到固定的d维空间X_p∈R}(N×d)，然后使用L个转换器块的堆栈来学习片间表示。视觉变形器自我注意的计算成本为O(N^2 d)。这里，C、H和W分别表示张量的通道、高度和宽度，P=wh是具有高度h和宽度w的面片中的像素数，N是面片的数量。由于这些模型忽略了CNN固有的空间归纳偏差，它们需要更多的参数来学习视觉表征。例如，与基于CNN的网络DeepLabv3(Chen等人，2017)相比，基于VIT的网络DPT(Dosovitski等人，2021)学习的参数多6倍，以提供类似的分段性能(DPT与DeepLabv3：3.45M比59M)。此外，与CNN相比，这些模型表现出不符合标准的优化性。这些模型对L2正则化很敏感，需要大量的数据增强以防止过度拟合(Touvron等人，2021A；肖等人，2021)。
本文介绍了一种轻量级VIT模型MobileViT。其核心思想是学习使用变压器作为卷积的全局表示法。这使我们能够在网络中隐含地合并类似卷积的属性(例如，空间偏差)，通过简单的训练配方(例如，基本增强)学习表示法，并轻松地将MobileViT与下游架构(例如，用于分段的DeepLabv3)集成。

3.1 MOBILEVIT架构

MobileViT block。如图1b所示，MobileViT块旨在用较少的参数在输入张量中对局部和全局信息进行建模。形式上，对于给定的输入张量X∈R^{(H×W×C)，MobileViT采用n×n标准卷积层，然后采用逐点（或1×1）卷积层生成X_L∈R}(H×W×d)。n×n卷积层编码局部空间信息，而逐点卷积通过学习输入通道的线性组合将张量投影到高维空间（或d维，其中d>C）。
使用MobileViT，我们希望在具有H×W的有效感受野的同时对远程非局部依赖性进行建模。对远程依赖性进行建模的广泛研究的方法之一就是膨胀的卷积。然而，这种方法需要仔细选择扩张率。否则，权重被应用于填充的零而不是有效的空间区域(Yu&Koltun，2016；Chen等人，2017；Mehta等人，2018)。另一个有希望的解决方案是自我关注(Wang等人，2018年；Ramachandran等人，2019年；Bello等人，2019年；Dosovitski等人，2021年)。在自我注意的方法中，多头自我注意的视觉转换器(VITS)被证明对视觉识别任务是有效的。然而，VITS重量很重，并且表现出不符合标准的优化能力。这是因为VITS缺乏空间诱导偏向(肖等人，2021；Graham等人，2021)。

为了使MobileViT能够学习具有空间感应偏差的全局表示，我们将X_L展开为N个不重叠的平坦面片X_U∈R^{(P×N×d)。这里，P=wh,N=HW/P是面片的数目，h≤n和w≤n分别是面片的高度和宽度。对于每个p∈{1，···，P}，通过应用transformers对面片间关系进行编码，以获得X_G∈R}(P×N×d)：
X_G §=" Transformer " (X_U §),1≤p≤P

与丢失像素空间顺序的VITS不同，MobileViT既不会丢失patch顺序，也不会丢失每个patch中像素的空间顺序(图1b)。因此，我们可以将X_G∈R^{(P×N×d)折叠，得到X_F∈R}(H×W×d)。然后使用逐点卷积将XF投影到低C维空间，并通过concat运算与X组合。然后，使用另一个n×n卷积层来融合这些级联特征。请注意，由于X_U §使用卷积对来自n×n区域的局部信息进行编码，而X_G §对第p个位置的P个面片上的全局信息进行编码，因此XG中的每个像素都可以对来自X中所有像素的信息进行编码，如图4所示。因此，MobileViT的整体有效感受野为H×W。

图 4 每个像素都可以看到MobileViT块中的其他像素。在本例中，红色像素使用转换器处理蓝色像素(位于其他面片中相应位置的像素)。由于蓝色像素已经使用卷积对有关相邻像素的信息进行了编码，因此允许红色像素对图像中所有像素的信息进行编码。在这里，黑色和灰色网格中的每个单元格分别表示一个面片（指patch）和一个像素。

与卷积的关系。标准卷积可以看作三个连续操作的堆栈：(1)展开、(2)矩阵乘法(学习局部表示)和(3)折叠。MobileViT块类似于卷积，因为它也利用相同的构建块。MobileViT块用更深层的全局处理(一堆转换器层)取代了卷积中的局部处理(矩阵乘法)。因此，MobileViT具有类似卷积的属性(例如，空间偏见)。因此，可以将MobileViT块视为卷积的转换器。我们有意设计的简单设计的一个优势是，可以开箱即用地使用卷积和转换器的低级别高效实现；允许我们在不同的设备上使用MobileViT，而不需要任何额外的努力。

重量轻。MobileViT块使用标准卷积和转换器分别学习局部和全局表示。由于以前的工作(例如，Howard等人，2017；Mehta等人，2021A)已经表明，使用这些层设计的网络是重量级的，因此自然会出现一个问题：为什么MobileViT是轻量级的？我们认为，问题主要在于学习具有变压器的全局表示法。对于给定的补丁（指patch），以前的工作(例如，Touvron等人，2021a；Graham等人，2021)通过学习像素的线性组合将空间信息转换为潜在的(图1a)。然后通过使用转换器学习块间信息来对全局信息进行编码。结果，这些模型失去了CNN固有的图像特定的感应偏差。因此，他们需要更多的能力来学习视觉表征。因此，它们又深又宽。与这些模型不同，MobileViT使用卷积和转换器的方式是，生成的MobileViT块具有类似卷积的属性，同时允许全局处理。这种建模功能使我们能够设计浅的和窄的MobileViT模型，而这些模型又是轻量级的。与使用L=12和d=192的基于VIT的模型Deit相比，MobileViT模型分别在32×32、16×16和8×8的空间级别使用L={2，4，3}和d={96,120,144}。由此产生的MobileViT网络比Deit网络更快(1.85×)、更小(2×)、更好(+1.8%)(表3；§4.3)。

计算成本。在MobileViT和VITS(图1a)中，多头自我注意的计算成本分别为O(N^2 Pd)和O(N^2 d)。理论上，与VITS相比，MobileViT的效率较低。然而，在实践中，MobileViT比VITS更有效率。在ImageNet-1K数据集上，MobileViT的FLOPS比Deit少2倍，准确度提高1.8%(表3；§4.3)。我们认为，这是因为与上面讨论的轻量化设计类似的原因。

MobileViT架构。我们的网络灵感来自于轻量级CNN的理念。我们针对通常用于移动视觉任务的三种不同网络规模(S：Small、XS：Extra Small和XXS：Extra Small)对MobileViT模型进行培训(图3c)。MobileViT的初始层是3×3标准卷积，然后是MobileNetv2(或MV2)块和MobileViT块(图1b和§A)。我们使用Swish(Elfwing等人，2018)作为激活函数。遵循CNN模型，我们在MobileViT块中使用n=3。特征图的空间维度通常是2和h，w≤n的倍数。因此，我们在所有空间级别设置h=w=2(有关更多结果，请参阅§C)。MobileViT网络中的MV2块主要负责下采样。因此，这些区块在MobileViT网络中是浅而窄的。图3D中MobileViT的空间级参数分布进一步表明，在不同的网络配置中，MV2块对总网络参数的贡献非常小。

3.2 多尺度采样器对训练效率的作用

在基于VIT的模型中，学习多尺度表示的标准方法是微调。例如，Touvron等人。(2021A)单独微调以224×224的空间分辨率训练的不同大小的Deit模型。这种学习多尺度表示的方法对于VITS是更好的，因为位置嵌入需要基于输入大小进行内插，并且网络的性能受内插方法的影响。与CNN类似，MobileViT不需要任何位置嵌入，它可以在训练期间受益于多尺度输入。
先前基于CNN的工作(例如，Redmon&Farhadi，2017；Mehta等人，2021年b)表明，多尺度培训是有效的。然而，这些工作中的大多数都是在固定次数的迭代后采样新的空间分辨率。例如，YOLOv2(Redmon&Farhadi，2017)在每10次迭代时从预定义的集合中采样新的空间分辨率，并在训练期间跨不同的GPU使用相同的分辨率。这会导致GPU利用率不足和训练速度变慢，因为所有分辨率都使用相同的批次大小(使用预定义集合中的最大空间分辨率确定)。为了便于MobileViT学习多尺度表示而无需精调，并进一步提高训练效率(即更少的优化更新)，我们将多尺度训练方法扩展到可变大小的批处理。给定空间分辨率S={(h1，w1)，···，(hn，wn)}的排序集合和最大空间分辨率(hn，wn)的批次大小b，我们在第t次训练迭代处随机采样空间分辨率(H_t,W_t )∈S，并计算用于第t次迭代的批次大小：b_t=(H_n W_n b)/(H_t W_t )。因此，较大的批次大小用于较小的空间分辨率。这减少了每个时期的优化器更新，并有助于更快的培训。

图5比较了标准采样器和多尺度采样器。在这里，我们将PyTorch中的DistributedDataParallel称为标准采样器。总体而言，多尺度采样器(I)减少了训练时间，因为它需要对可变大小的批次进行更少的优化器更新(图5b)，(Ii)将性能提高约0.5%(图10；§B)，以及(Iii)迫使网络学习更好的多尺度表示(§B)，即，当以不同的空间分辨率进行评估时，相同的网络产生更好的性能与用标准采样器训练的相比。在§B中，我们还证明了多尺度采样器是通用的，并改进了CNN的性能(例如，MobileNetv2)。