MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer论文翻译理解

1. 背景

   目前CV领域主要使用轻量级的CNN，CNN带有的空间归纳偏置（局部性，平移不变性）使得它能在各种任务上以更少的参数学习表征。然而，CNN是局部的。为了学习全局特征，ViT被应用到了CV领域，但是比起CNN，ViT更复杂。那么有没有一种可能，将CNN和ViT的优点结合起来，去构建一个轻量级而且低延迟的网络呢。本文提出了MobileViT，轻量且通用的视觉transformer。MobileViT提出了一种不同的观点，使用transformer去处理全局信息，把transformer当卷积。

   基于自注意力的模型，如ViT，可以代替CNN去学习视觉特征。ViT是将图片切分成了一系列薄片，使用transformer的多头注意力机制去学习片间特征。为了提高ViT的性能，就需要去增加模型的参数，但是这带来的后果是模型更大，延迟更高。现实世界中的视觉识别任务需要在资源受限的移动设备上实时运行，为了更有效，ViT必须轻量而且更快，尽管ViT的模型大小已经减小到和移动设备匹配了，但是它的性能比起轻量级CNN还是差得很远。对于5-6百万级别的参数，DeIT准确率比MobileNet3低3%。

   轻量级CNN支持移动视觉任务，但ViT离支持移动设备还很遥远。轻量级CNN能很容易的优化，与特定任务的网络结合；而ViT太复杂了，很难优化，需要广泛的数据增强和L2正则化来阻止过拟合，对于下游任务还需要昂贵的decoder，尤其是密集型预测任务。例如对于同样的性能，一个基于ViT的分割网络需要3亿4千5百万个参数，而基于CNN的DeepLabv3只需要5千9百万个参数。ViT为什么需要更多参数，因为它缺少归纳偏置，而CNN自带归纳偏置。为了建立鲁棒性和高性能的ViT，将卷积和transformer结合的方式成为了研究热点。但是混合模型还是复杂，对数据增强依然敏感，去掉 CutMix 和 DeIT-style 数据增强，模型准确率从78.1%降到了72.4%。

   浮点操作不足以实现在移动设备上低延迟，因为浮点操作忽略了几个与推测相关的因素，如内存访问，并行程度，平台特征。同样基于ViT，PiT比DeIT少3倍浮点操作，但是二者在移动设备上性能差不多，所以本文不考虑浮点操作。我们提出了MobileViT，既有CNN的归纳偏置，对数据增强的不敏感；也有ViT的输入适应性权重和全局处理。我们创造性的提出了MobileViTBlock，在一个tensor中编码局部和全局信息。MobileViT以一种不同于ViT的方式学习全局特征。标准卷积包括展开、局部处理、折叠。MobileViTBlock使用transformer把局部处理换成了全局处理，这就让MobileViTBlock同时有了CNN和ViT的优点，使用更少的参数、更简单的方法，学到更好的特征。MobileViT也可以作为其他架构的骨架，例如把SSDLite中的MobileNetv3换成MobileViT，能让模型减小1.8+大小，mAP提升1.8%。

   目前工业级和学术界设计轻量化神经网络模型主要有4个方向：
   1）人工设计轻量化神经网络模型；
   MobileNet v1, MobileNet v2, ShuffleNet v1, ShuffleNet v2
   2）基于神经网络架构搜索（Neural Architecture Search,NAS）的自动化设计神经网络；
   NasNet, MNasNet
   3）CNN模型压缩；
   剪枝、权值共享和权值量化、哈夫曼编码
   压缩算法可分为四类：参数修剪和共享、低秩分解、迁移/压缩卷积滤波器和知识蒸馏等。
   4）基于AutoML的自动模型压缩。
   AutoML for Model Compres- sion，AMC
   PocketFlow——腾讯AI Lab机器学习中心
   TensorFlow Lite

   轻量级CNN优点：具有空间归纳配置特性（局部性，平移不变性），使得需要的参数更少；
   对数据增强和L2正则化不敏感
   缺点：只能学习局部特征，缺少全局特征

   ViT优点：可以通过transformer的多头注意力机制学习全局特征；
   缺点：缺少空间归纳配置特性，只能通过大量参数去学习局部特征；
   模型太复杂，需要很大的数据集，容易过拟合；
   对数据增强和L2正则化敏感

   目标：结合两者的优点，建立一个轻量级、低延迟、通用的网络。

2. 文献关键词
   CNN、ViT、MobileViT、轻量级、低延迟、通用

3. 相关工作
   轻量级CNN：标准卷积计算量大，目前流行的是可分离式卷积，如MobileNet等。这些网络可以替换其他网络的骨架以减小模型大小，降低延迟。但是缺点是只能编码局部信息。

   ViT：只有在数据集特别特别大的时候，ViT才能实现和CNN差不多的性能，即使它没有归纳偏置。使用广泛的数据增强、L2正则化、蒸馏，ViT才能达到CNN级准确率。但是ViT还是有不合标准的可优化性，使得模型难以训练，这是由于缺乏空间归纳偏置。在ViT中加入卷积，能提高稳定性和性能。ViT-C对ViT加入了卷积；CvT修改了transformer中的多头注意力机制，使用深度级可分离卷积代替线性映射。 BoTNet把ResNet瓶颈单元中的3*3标准卷积换成了多头注意力机制；ConViT使用 gated positional self-attention 去融合 soft convolutional 空间偏置。这些模型使用数据增强，也能达到CNN级性能，但是还是太复杂了。 PiT 和 CvT 分别学习了6.1倍和7.1倍参数，才能在image-1k上达到 EfficientNet 的性能。若要减小这些模型的规模，就导致性能比CNN差更多。

   讨论：1）对于差不多的参数数目，MobileViT比轻量级CNN性能更好，在各种任务上；
   2）MobileViT比ViT泛化能力更好，接近于CNN；
   3）MobileViT有更好的鲁棒性，比起ViT，对数据增强和L2正则化不敏感。

4. ViT网络结构
   ViT将输入图像切分成h×w的patch，那么一个通道就可以切成N个（H×W/h×w）patch，一共有N×C个patch，每一个patch的P=h×w。将N×P×C维度的输入映射成N×d维，使用L个transformer学习patch间特征，自注意力的计算成本为O(N^2 × d)。因为ViT忽略了归纳偏置，所以需要学习更多参数。DPT就要比DeepLabv3多学6倍的参数才能达到相同效果。另外ViT需要L2正则化和数据增强来防止过拟合。

5. MobileViTBlock结构
   为了建模全局信息，一种方法是使用空洞卷积，但是需要谨慎选择膨胀率。所以我们采用了第二种方法，就是自注意力机制。

   先给图像切片，使用33，11卷积核编码局部信息，不需要位置编码，然后使用transformer的多头注意力机制编码了全局信息，而且不会丢失patch内像素的顺序、patch的顺序。经过transformer之后，使用11卷积将输出降到C维，再和X连接，最后33卷积将局部信息和全局信息融合。注意transformer的输出图像，每一个像素都编码了输入图像的所有像素的信息，因为h=w=2

   这种卷积+transformer的有目的的简单设计可以开箱即用，让我们在不同设备上直接使用MobileViT。为啥在MobileViT上，这种卷积+transformer的结构突然就轻量级了呢？因为之前ViT只学习了像素的线性连接，丢失了归纳偏置，需要更多参数，因此模型宽而深。而MobileViT加入了卷积，就有了归纳偏置，可以建立一个窄而浅的网络。

   理论上，MobileViT的时间复杂度为O(N^2 × P × d)，ViT的时间复杂度为O(N^2 * × d)，前者更耗时，但实践上，前者效果更好，这归功于之前提到的轻量级设计。

6. MobileViT网络结构
   有三种不同规模的MobileViT：XXS, XS, S。对于XS, S，MV2的瓶颈扩大率为4，对于XXS，MV2的瓶颈扩大率为2。

   使用swish激活函数。

   MV2的参数量对于整个网络的参数来说微不足道。

   权重衰减：MobileViT对于权重衰减（正则化）具有强大的鲁棒性，将衰减率从0.1变化到0.0001，除了0.1时，模型准确率特别低之外，其他都差不多好。我们选择的是0.01，大多数CNN采用的是0.0001。

   跳跃连接：跳跃连接可以提升0.5%的性能，即使没有它，MobileViT也能达到和CNN差不多的水平。

   patch大小：当n=3时，尝试了2，3，4，{8,4,2}的方案。当h=w<=n时，能编码到所有像素信息，虽然3的准确率最高，但是因为分辨率都是2的倍数，需要额外的填充或者resize操作，更耗时。所以取2做一个平衡。

7. 多尺度采样器
   对于ViT来说，学习多尺度特征的方法是微调。例如DeIT在训练时分辨率为224*224，在微调时会换成各种尺寸。这种方式是适用于ViT的，因为ViT需要 positional embeddings ，positional embeddings 又需要根据输入大小进行插值，网络性能又受到插值方法的影响。而CNN和MobileViT不需要 positional embeddings ，在训练时就进行多尺度输入会很受益。

   之前的工作往往是在训练几个iteration后，采用新的分辨率，在不同的GPU上分辨率也一样，但是这会导致GPU利用率低，训练太慢，因为batch数由最大分辨率决定，一直不会改变。为了让MobileViT在训练时就能学习多尺度特征，不需要再微调，我们动态修改了batch大小。给定一组有序的分辨率集合S={(H1,W1),(H2,W2),…,(Hn,Wn)}，和根据最大分辨率确定的batch大小b，每经过t论iteration，就随机选择一个分辨率，重新计算batch大小bt=Hn×Wn×b/Ht×Wt。结果就是分辨率越小，batch大小越大，batch数目越少，单个epoch的更新次数更少，训练时间越短。

   好处有：更短的训练时间、更好的性能、泛化能力更强、学到更好的特征。而且多尺度采样器是通用的，不仅能提升MobileViT，还能提升轻量级和复杂CNN。

8. 在image-1k数据集上进行图像分类的实验过程
   image-1k数据集用于分类，训练集有128万张图片，验证集有5万张图片。在8个 NVIDIA GPU上，使用pytorch训练300个epoch，batch大小1024张图片，AdamW优化器，标签平滑交叉熵损失函数，多尺度采样器。最开始3千个iteration，学习率从0.0002增加到0.002，之后使用余弦调度退火到0.0002。使用L2权重衰减。使用基本的数据增强、用top-1准确率去评估性能。在测试时使用指移动平均值。

9. 在image-1k数据集上进行图像分类的实验结果
   1）与轻量级CNN相比（使用基本数据增强）
   参数大约两百五十万时
   2）与复杂CNN相比（使用基本数据增强）
   3）与ViT相比（不使用蒸馏）

   ViT使用强化数据增强的性能要比使用基础数据增强好很多，而MobileViT使用基础数据增强就能用更少的参数达到更好的性能。

10. 在MS-COCO数据集上进行目标检测的实验过程
    在MobileNetv2后面，SSD把标准卷积换成可分离式卷积，就变成了SSDLite。使用MobileViT提取特征，SSDLite目标检测。在image-1k数据集上训练好后，我们对MobileViT进行微调。
    MS-COCO数据集的训练集包括11万7千张图片，验证集包括5千张图片。在4 NVIDIA GPUs 上训练200epoch，输入分辨率320*320，batch大小128张图片，使用 AdamW 优化器，平滑L1损失函数来目标定位和分类，使用标准采样器。最开始500个iteration，学习率从0.00009增加到0.0009，使用余弦退火到0.00009。使用L2权重衰减，衰减率为0.01，在不同IoU阈值（从0.5到0.95，步长0.05）（0.5、0.55、0.6、0.65、0.7、0.75、0.8、0.85、0.9、0.95）上计算平均mAP（mAP是所有类别的平均精确度，AP是单个类别的平均精确度）。
    使用平滑L1损失函数来目标定位和分类

11. 在MS-COCO数据集上进行目标检测的实验结果
    1）MobileViT+SSDLite与轻量级CNN+SSDLite相比
    2）MobileViT+SSDLite与复杂CNN+SSD-300相比

12. 在 PASCAL VOC 2012 数据集上进行语义分割的实验过程

    使用MobileViT提取特征，DeepLabv3 进行语义分割。

    在4 NVIDIA GPUs 上训练50epoch，batch大小128张图片，使用 AdamW 优化器，交叉熵损失函数，使用标准采样器。最开始500个iteration，学习率从0.00009增加到0.0009，使用余弦退火到0.00009。经过标准训练过程后，又加入了其他注释和数据。使用L2权重衰减，衰减率为0.01，使用mIOU进行评估。

    为什么目标检测+语义分割不使用多尺度采样器了？因为这些特定任务都有确定的分辨率。DeepLabv3使用膨胀率6，12，18去学习多尺度特征，如果分辨率太低，空洞卷积核只能应用到填充的0上，影响多尺度学习效果。

13. 在 PASCAL VOC 2012 数据集上进行语义分割的实验结果
    1）MobileViT+DeepLabv3与轻量级CNN+DeepLabv3相比

14. 在移动设备上的测试
    针对8，16，31的输出步长，patch大小有两种配置，A={2,2,2}，B={8,4,2}。越小的patch性能越好，这是因为输出图像的每一个像素都能编码到输入图像全部像素的信息。B的速度更快，因为大的patch意味着patch数目少，理论上两种配置复杂度都是O(N^2 × P × d)，但是B的N更小，所以时间短。B的并行程度更高，因为自注意力针对大patch的更多像素可以同时计算。为了进一步降低延迟，使用了线性自注意力机制。除了MobileViT S在语义分割任务上不能实时之外，其他均能实时，因为语义分割的输入比分类和检测更大。

    讨论：在移动设备上，MobileViT和ViT比CNN慢，这和之前讨论的不相符。主要有两个原因：首先，GPU上的Transformer有专用的CUDA内核，在ViTs中开箱即用，以提高GPU上的可扩展性和效率。其次，CNN得益于多种设备级优化，包括与卷积层融合的批量标准化。然而，这些专用和可优化的操作在移动设备上都无法提供给transformer。如果将来能在移动设备上提供这些专用和可优化的操作，MobileViT和ViT的时间一定能赶上CNN。

15. 实验数据分析
    1）使用image-1k数据集，MobileViT有6百万参数，实现了78.4%的top1准确率，300epoch，1024batch size。
    而差不多的参数，MobileNet3实现了75.2%的top1准确率，600epoch，4096batch size；DeIT实现了72.2%的top1准确率。
    2）在MS-COCO目标检测任务上，差不多的参数，MobileViT比MobileNetv3高5.7%
    3）把SSDLite中的MobileNetv3换成MobileViT，能让模型减小1.8+大小，mAP提升1.8%。
    4）PiT实现了73%的准确率。