[Transformer]MobileViT：Light-weight, General-purpose, and Mobile-friendly Vision Transformer

from Apple

Abstract

轻量级的卷积神经网络是移动端视觉任务的de-facto，得益于其空间上的归纳偏执使得他们可以通过较少的参数就学习特征表述。但是这些网络都是基于局部空间的，为了学习全局的特征表示提出了基于SA的ViT模型。但是ViT并不轻量级。本文提出这样一个问题：是否有可能结合CNN和ViT各自的长处搭建一款轻量级低延迟的适合移动端视觉任务的框架呢？
本文提出MobileViT，从全新的视角来处理全局信息，即transformer作为卷积。实验结果显示MobileViT比基于CNN和基于ViT的框架取得了更有益的性能。比如在ImageNet-1K达到了78.4%的top-1精度，比MibileNetv3和DeiT分别提升了3.2和6.2%。在COCO目标检测任务则比MobileNet提升了5.7%。

Section I Introduction

基于SA的Transformer已经成为CNN的一种替代方案，ViT通过将输入图像切分一系列互不重叠的Patch然后使用多头注意力特征表述，也有一系列工作来进一步提升ViT的性能，但都是以模型大小、网络参数和延迟为代价。但是现实世界许多场景都是需要在资源有限的平台上进行视觉识别任务，这就要求设计的ViT模型更加轻量级低延迟。比如当参数量5-6M时DeiT比MobileNetv3的性能低3%,因此轻量级ViT的模型设计志在必得。




轻量级CNN已经应用于多种移动端的视觉任务中，但是基于Transformer的ViT模型在轻量级方面任重而道远，首先ViT是非常“重量级”的，并且需要预训练不好优化，往往需要数据增强，如果是密集预测任务还需要昂贵的decoder。比如相同的分割任务ViT参数量达到了345M而同样的CNN（DeepLab v3）只需要59M。



为了搭建更加鲁邦的模型，有的工作致力于将卷积与Transformer混合，但是这种hybrid的模型仍然十分“重量级”，并且对数据增强十分敏感。如果移除CutMix会导致准确性显著下降。




如何将卷积与Transformer结合仍然值得我们探索，尤其是轻量级探索，那就要求 设计 更加轻量级、低延迟、准确率高的满足一定硬件资源约束的模型，并且最好是通用的才能适合不同的任务。

需要注意的是FLOPs低并不一定意味着设备的低延迟，因为FLOPs忽略了其他因素，比如内存访问次数、并行成都和硬件平台的特性。比如Heo等人提出的基于ViT的方法比DeiT的FLOPs少3x但是在实际硬件部署二者差不多；因此本文没有选择FLOPs作为优化目标，而是重点设计一个轻量级、通用和低延迟的网络。

本文最终设计的网络叫做MobileViT，较好的结合了CNN的空间归纳偏执和对数据增强的敏感性，以及VIT全局依赖和输入自适应的优点。尤其是提出的MobileViT模块可以有效的计算张量的局部和全局信息，参见Fig 1(b)。



与其他ViT变体不同之处在于，MobileViT提供了学习全局表示的不同视角。
标准卷积包含3类操作：展开、局部处理、折叠。

MobileViT将卷积中的局部操作替换为Transformer的全局计算，这样使得MobileViT同时具有卷积和Transformer的性能共，可以以更少的参数就学习更好的特征表述。
据我们所知这是第一个表明轻量级ViT可以达到与轻量级CNN相近的性能，并且不需要复杂的训练。
参数量在5-6M时轻量级的MobileViT在ImageNet-1K上达到了78.4%的精度，比MobileNet高3.2%;同时本文还观察到将MobileViT作为特征提取器用于其他移动端视觉任务时也对性能有显著提升，比如用MobileViT替换MNASNet可以获得1.8%mAP的提升，网络减少了1.8x。

Section II Related Work

Light-weight CNN


卷积层是卷积神经网络的基础构建，为了减轻卷积层的计算成本，有的尝试了如何将卷积分解，比如MobileNet，ShuffleNet，ESPNet，MixNet，MNASNet等这些轻量级网络，也比较好训练，因此可以替换一些重量级网络的backbone，但是他们都有一个问题：就是在空间上是局部的。


本文则致力于利用卷积和Transformer二者的优势来搭建轻量级的通用网络。



Vision Transformer



Transformer的优势在于可以捕捉全局联系，但是却难以训练，主要就是因为ViT缺乏空间上的归纳偏执信息，因此需要借助在大规模数据集上预训练的结果。


也有的工作尝试借助卷积来弥补归纳偏执信息，比如使用深度可分离卷积替换原来的linear操作，BoTNet则是将ResNet中3x3的卷积替换为了MHSA操作。虽然这些改进都一定程度上提升了性能，但是一般都是“重量级”的网络。比如PiT,CvT分别是EfficientNet的参数量的6.1x和1.7x；如果把参数量降下来那这些模型性能就比轻量级CNN差了，比如同等参量的PiT精度会比MobileNet下降2.2%.



Discussion




一个悬而未决的问题是：如何有效的结合卷积和Transformer来搭建轻量级的网络模型。为此本文将ViT与CNN结合搭建了MobileViT，有以下优势：




（1）更好的性能。
在相同的参数下本文的MobileViT获得了最佳的性能；





（2）更好的泛化性能。
泛化能力指的是在测试集和验证集上的差距，与ViT相比MobileViT在没见过的数据集上有更好的性能






（3）鲁棒性。






一个好的模型应该对超参数具有足够的鲁棒性，因为人工调参非常耗时也耗费资源，与ViT相比，MobileViT只需要在简单的数据增强就能训练，并且对L2正则化不敏感。

Section III MobileViT：A light-weight Transformer

Fig 1（a）展示了标准的ViT模型，可以看到对输入需要将其切分为一系列patch，然后将patch投射到某一维度，随后送入Transformer block进行处理。


其中SA的计算复杂度为O(N2d)；但由于缺失空间信息，因此ViT需要更多的参数量来进行学习才能达到CNN相近的性能。


本文MobileViT的核心是transformer as convolution(Transformer即卷积)来学习全局特征表述。优势是借助卷积学习空间偏置信息，简单的训练就可以学习到特征表述，非常适合继承到其他下游任务中。




Part 1 MobileViT

Fig 1(b)展示了MobileViT block的基本结构，主要就是如何以较少的参数量同时建模全局和局部信息。


对于输入向量X 首先会进行nxn卷积，然后进行point-wise卷积（或1x1卷积）来产生特征，nxn卷积负责编码局部空间信息，point-wise conv负责将张量通过通道的线性组合投影到更高位空间。



此外还需要对全局依赖进行建模，一种候选操作是膨胀卷积，但是需要仔细的设置膨胀率，不然不会作用到有用的区域。




另一种候选方案就是self-attention。但是SA一方面是heavy-weight，另一方面缺少空间信息。
因此为了使得MobileViT能够学习到带有空间偏执的全局特征表述，本文将XL拆成N个互不重叠的展平的patch，其中patch_size = hw,N=HW/P,对于每一个小的patch进行Transformer的计算。这样MobileViT既不会丢失每个patch的顺序也不会丢失patch内部像素的空间信息。然后将么每一个patch级联重组。
将重组之后的结果再经过卷积操作获得最终输出，这样就将全局特征和局部特征融合在了一起。



其中Xup负责的是局部编码，Xgp负责的是全局编码，Xg中每一个patch对应全局的一块区域，因此最终感受野是HxW.




Relationship to Convolutions




传统的卷积可以分解成3步：
展开-矩阵乘法-折叠




MobileViT与之类似，也分成了这3不，只不过将卷积里的矩阵乘法替换成了堆叠的transformer layer。因此可以看做是一种transformer as convolution(Transformer即卷积)。这种设计的好处一是开箱即用，而是允许在不同的设备上引入MobileViT。




Light-weight 
MobileViT




分别用卷积核Transofmer来学习局部和全局特征，但之前的研究基于Transformer模块搭建的网络都很“heavy”，那为什么MobileViT很轻量级呢？本文认为主要是在用Transformer学习全局表征这里。之前的工作大多将输入经过线性映射到隐空间，然后在使用Transformer学习patch之间的关系，对全局信息进行编码，但是这就丢失了图像特有的归纳偏执，这是CNN所具有的。




而MobileViT同时结合了卷积和Transformer使得block虽然进行Transformer的操作但是具有卷积的特性，这样可以设计较为浅、窄的网络，是一种轻量级的网络。
比如ViT,DeiT设置为L=12，d=192维度时，MobileViT只要设置L={2,4,3}，d={96,120,144}即可，比DeiT推理速度块1.85x，规模小2倍，同时性能提升了1.8%。






Computational cost

MobileVit:O(N2Pd)


ViT:O(N2d)




MobileViT网络架构
分为3种规格：S,XS,XXS 模型参量分别是5.6,2.3,1.3M


具体结构参见Fig 1（c）.



MobileViT最开始是一个标准的3x3卷积，后面就是MobileNetv2(MV2)和MobileViT模块，从Fig 3(d)也能看到不同的网络组合中，MV2对参数量的贡献量很小。





Part 2 Multi-scale Sampler for Training Efficiency

在ViT模型中学习多尺度的一种方法就是进行微调，比如不断使用不同的分辨率训练DeiT，这种方法适合ViT的原因是位置嵌入需要根据输入大小进行插值，而MobileViT不需要任何位置嵌入，在训练期间就接受了多尺度输入。
早期CNN的网络也显示出使用多尺度训练时有效的，但是大多数都会在一个新的空间分辨率迭代一定次数，这会影响GPU利用率、训练也会变慢。
为了提升效率，本文对多尺度训练方法进行了优化来使用不同大小的批处理。对于一定空间分辨率，本文会在每个GPU的第t次迭代中随机采样一个空间分辨率进行训练，batch_size较大就使用较小的空间分辨率这样减少每次迭代优化器的更新，也可以加速训练。

Fig 5展示了标准和本文改进的multi-scale sampler的对比。可以看到优化后显著减少了训练时间，并且也获得了0.5%的精度提升。

Section IV Experiment Results

Part 1 Image Classification

Part 2 MobileViT as backbone

COCO Detection




Segmentation






Part 3 Performance on Mobile Devices

轻量级和低延迟的网络对于移动端的部署十分重要，本文先训练全精度的MobileViT然后使用公开的CoreMLTools将其转换为在iPhone12上的推理时间，来衡量其效率。
Mobile friendly
Fig 8展示了三种任务下不同patch_size的MobileViT的推理时间。可以看到Config-B比Config-A更快，因为patchsize更大，并行效果更好。




如果想要进一步降低延迟，还可以使用线性复杂度的SA。
并且这两种设定都可以达到实时性的要求（30FPS）。






Discussion






从Table 3可以观察到，MobileViT以及其他Transformer模型还是比MobileNet慢，主要是因为GPU上有专门的CUDA核用于卷积运算，并且CNN网络还有其他设备级别的优化，比如BN融合等可以减少延迟和方才能次数。本文相信未来针对MobileViT进行专门的设备级优化也可以进一步提升MobileViT的速度。