MOBILEVIT: 轻量级、通用且适用于移动设备的视觉Transformer

摘要

轻量级卷积神经网络（CNN）是移动视觉任务的事实标准。它们的空间归纳偏置使得它们能够在不同的视觉任务中以较少的参数学习表示。然而，这些网络在空间上是局部的。为了学习全局表示，基于自注意力的视觉Transformer（ViT）被采用。与CNN不同，ViT是重量级的。本文提出了以下问题：是否有可能将CNN和ViT的优势结合起来，构建一个适用于移动视觉任务的轻量级低延迟网络？为此，我们介绍了MobileViT，一个适用于移动设备的轻量级通用视觉Transformer。MobileViT为信息的全局处理提供了一个不同的视角。我们的结果表明，MobileViT在不同任务和数据集上显著超越了基于CNN和ViT的网络。在ImageNet-1k数据集上，MobileViT在约600万参数下达到了78.4%的Top-1准确率，比MobileNetv3（基于CNN）和DeIT（基于ViT）提高了3.2%和6.2%，且参数数量相似。在MS-COCO目标检测任务上，MobileViT比MobileNetv3高出5.7%的准确率，且参数数量相似。我们的源代码是开源的，链接：https://github.com/apple/ml-cvnets。

1. 引言

基于自注意力的模型，特别是视觉Transformer（ViT；图1a；Dosovitskiy等，2021），是卷积神经网络（CNN）的替代方案，用于学习视觉表示。简而言之，ViT将图像分割成一系列不重叠的块，然后使用多头自注意力机制在Transformer中学习这些块之间的表示（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不同，ViT模型较为重量级（例如，ViT-B/16与MobileNetv3：86百万参数与7.5百万参数），优化难度较大（Xiao等，2021），需要广泛的数据增强和L2正则化以防止过拟合（Touvron等，2021a；Wang等，2021），并且在下游任务中尤其是密集预测任务中需要昂贵的解码器。例如，基于ViT的分割网络（Ranftl等，2021）学习了约345百万个参数，并且在性能上与基于CNN的网络DeepLabv3（Chen等，2017）相似，而后者的参数量为59百万。ViT模型需要更多参数的原因可能是它们缺乏CNN固有的图像特定归纳偏置（Xiao等，2021）。为了构建强大且高效的ViT模型，结合卷积和Transformer的混合方法正在受到关注（Xiao等，2021；d’Ascoli等，2021；Chen等，2021b）。然而，这些混合模型仍然是重量级的，并且对数据增强非常敏感。例如，去除CutMix（Zhong等，2020）和DeIT风格（Touvron等，2021a）的数据增强会导致Heo等（2021）在ImageNet上的准确率从78.1%下降到72.4%。

将 CNN 和 transformer 的优势结合起来，构建适用于移动视觉任务的 ViT 模型仍然是一个开放问题。移动视觉任务需要轻量、低延迟且精确的模型，以满足设备的资源限制，并且要具备通用性，能够应用于不同的任务（例如，分割和检测）。需要注意的是，浮点运算（FLOPs）对于移动设备上的低延迟并不充分，因为 FLOPs 忽略了几个与推理相关的重要因素，如内存访问、并行度和平台特性（Ma 等人，2018）。例如，Heo 等人（2021）提出的基于 ViT 的方法 PiT，其 FLOPs 比 DeIT（Touvron 等人，2021a）少 3 倍，但在移动设备上的推理速度相似（DeIT vs. PiT 在 iPhone-12 上：10.99 毫秒 vs. 10.56 毫秒）。因此，本文的重点并非优化 FLOPs，而是专注于设计一个轻量、通用且低延迟的网络，用于移动视觉任务。我们通过 MobileViT 实现了这一目标，MobileViT 结合了 CNN（例如，空间归纳偏置和对数据增强的较低敏感性）和 ViT（例如，输入自适应加权和全局处理）的优势。具体来说，我们引入了 MobileViT 块，它能有效地在张量中编码局部和全局信息（图 1b）。与 ViT 及其变种（有卷积和无卷积的情况）不同，MobileViT 提出了一个不同的视角来学习全局表示。标准卷积包括三个操作：展开、局部处理和折叠。MobileViT 块将卷积中的局部处理替换为使用 transformer 进行的全局处理。这使得 MobileViT 块具有类似于 CNN 和 ViT 的特性，从而帮助其以更少的参数和简单的训练方案（例如，基本的数据增强）学习更好的表示。根据我们的最佳知识，这是首个展示轻量级 ViT 可以通过简单的训练方案，在不同移动视觉任务中实现轻量级 CNN 水平性能的工作。在大约 500-600 万个参数预算下，MobileViT 在 ImageNet-1k 数据集（Russakovsky 等人，2015）上的 top-1 准确率为 78.4%，比 MobileNetv3 高出 3.2%，并且使用了简单的训练方案（MobileViT vs. MobileNetv3：300 vs. 600 轮；1024 vs. 4096 批量大小）。我们还观察到，当 MobileViT 被用作高度优化的移动视觉任务特定架构中的特征骨干时，性能显著提升。将 MNASNet（Tan 等人，2019）替换为 MobileViT 作为 SSDLite（Sandler 等人，2018）中的特征骨干，结果是更好的（+1.8% mAP）且更小（1:8×）的检测网络（图 2）。

2 相关工作

轻量级 CNN

CNN 的基本构建层是标准卷积层。由于该层计算开销大，提出了几种基于因式分解的方法来使其轻量化并适应移动设备（例如，Jin 等人，2014；Chollet，2017；Mehta 等人，2020）。其中，Chollet（2017）的可分离卷积引起了广泛关注，并被广泛应用于移动视觉任务的最先进轻量级 CNN 中，包括 MobileNets（Howard 等人，2017；Sandler 等人，2018；Howard 等人，2019），ShuffleNetv2（Ma 等人，2018），ESPNetv2（Mehta 等人，2019），MixNet（Tan & Le，2019b），和 MNASNet（Tan 等人，2019）。这些轻量级 CNN 灵活且易于训练。例如，这些网络可以轻松替代现有任务特定模型中的重量级骨干（例如，ResNet（He 等人，2016）），从而减少网络大小并提高延迟性能。尽管这些方法有许多优点，但其中一个主要缺点是它们在空间上是局部的。本文将 transformer 视为卷积，从而能够利用卷积（例如，灵活且简单的训练）和 transformer（例如，全局处理）的优点来构建轻量级（§3）和通用（§4.1 和 §4.2）ViT。

视觉 transformer

Dosovitskiy 等人（2021）将 Vaswani 等人（2017）的 transformer 应用于大规模图像识别，并展示了在极大规模数据集（例如，JFT-300M）上，ViT 能够在没有图像特定归纳偏置的情况下实现 CNN 水平的精度。通过广泛的数据增强、重 L2 正则化和蒸馏，ViT 可以在 ImageNet 数据集上训练，并实现 CNN 水平的性能（Touvron 等人，2021a；b；Zhou 等人，2021）。然而，与 CNN 不同，ViT 显示出次优的可优化性，并且训练困难。随后的一些工作（例如，Graham 等人，2021；Dai 等人，2021；Liu 等人，2021；Wang 等人，2021；Yuan 等人，2021b；Chen 等人，2021b）表明，这种次优的可优化性是由于 ViT 缺乏空间归纳偏置。通过在 ViT 中结合卷积，能够提高其稳定性和性能。为了充分利用卷积和 transformer 的优势，已经探索了不同的设计。例如，Xiao 等人（2021）的 ViT-C 在 ViT 中加入了早期卷积层。CvT（Wu 等人，2021）修改了 transformer 中的多头注意力，并使用深度可分离卷积替代线性投影。BoTNet（Srinivas 等人，2021）在 ResNet 的瓶颈单元中用多头注意力替代了标准的 3×3 卷积。ConViT（d’Ascoli 等人，2021）通过门控位置自注意力结合了软卷积归纳偏置。PiT（Heo 等人，2021）通过深度可分离卷积池化层扩展了 ViT。尽管这些模型能够通过广泛的增强达到与 CNN 相当的性能，但大多数模型仍为重量级。例如，PiT 和 CvT 的参数量分别是 EfficientNet（Tan & Le，2019a） 6:1 和 1:7 的倍数，并且在 ImageNet-1k 数据集上的表现相似（top-1 准确率约为 81.6%）。此外，当这些模型被缩小以构建轻量级 ViT 模型时，它们的性能明显不如轻量级 CNN。在大约 600 万个参数预算下，PiT 在 ImageNet-1k 数据集上的准确率比 MobileNetv3 少了 2.2%。

讨论

将卷积和transformer结合起来，能够得到比原始ViT更强大和高效的ViT模型。然而，这里存在一个开放性问题：如何将卷积和transformer的优势结合起来，构建轻量化的网络，以适应移动视觉任务？本文专注于设计轻量化的ViT模型，这些模型能够在简单的训练方案下超越现有的最先进模型。为此，我们提出了MobileViT，结合了CNN和ViT的优势，构建一个轻量化、通用且适用于移动设备的网络。MobileViT带来了几个新颖的观察结果：(i) 更好的性能：在给定参数预算的情况下，MobileViT模型在不同的移动视觉任务中表现优于现有的轻量级CNN (§4.1和§4.2)。(ii) 泛化能力：泛化能力是指训练和评估指标之间的差距。对于两个训练指标相似的模型，具有更好评估指标的模型具有更强的泛化能力，因为它能在未见过的数据集上做出更好的预测。与先前的ViT变体（有和没有卷积）相比，尽管进行了大量的数据增强，但它们的泛化能力较差，MobileViT表现出了更好的泛化能力（图3）。(iii) 鲁棒性：一个好的模型应该对超参数（例如数据增强和L2正则化）具有鲁棒性，因为调整这些超参数既费时又费资源。与大多数基于ViT的模型不同，MobileViT模型使用基础的数据增强进行训练，并且对L2正则化的敏感度较低（§C）。

3 MobileViT: 一种轻量化的Transformer模型

一个标准的ViT模型，如图1a所示，将输入 $\mathbf{X}\in {\mathbb{R}}^{H\times W\times C}$ 重新组织成一个展开的补丁序列 ${\mathbf{X}}_f\in {\mathbb{R}}^{N\times PC}$，并将其投影到一个固定的 $d$ 维空间 ${\mathbf{X}}_p\in {\mathbb{R}}^{N\times d}$，然后通过堆叠L个Transformer块来学习补丁间的表示。Vision Transformers (ViT) 中自注意力的计算成本为 $O(N^{2}d)$。这里，$C$、$H$ 和 $W$ 分别表示张量的通道数、高度和宽度，$P=wh$ 是补丁中像素的数量（补丁的高度为 $h$，宽度为 $w$），$N$ 是补丁的数量。由于这些模型忽略了CNN中固有的空间归纳偏置，它们需要更多的参数来学习视觉表示。例如，基于ViT的网络DPT（Dosovitskiy等，2021）相比于基于CNN的DeepLabv3（Chen等，2017），需要学习6倍更多的参数，才能提供相似的分割性能（DPT vs. DeepLabv3: 345M vs. 59M）。与CNN相比，这些模型在优化上表现较差，它们对L2正则化敏感，并且需要大量的数据增强来防止过拟合（Touvron等，2021a；Xiao等，2021）。

本文介绍了一种轻量化的ViT模型——MobileViT。其核心思想是通过Transformer学习全局表示，同时借助卷积的优势。这样，我们能够在网络中隐式地加入类似卷积的特性（例如空间偏置），通过简单的训练策略（例如基础的数据增强）学习表示，并且可以轻松地将MobileViT与下游架构（例如用于分割的DeepLabv3）集成。

3.1 MobileViT架构

MobileViT块
MobileViT块，如图1b所示，旨在通过更少的参数对输入张量中的局部和全局信息进行建模。形式上，对于给定的输入张量 $\mathbf{X}\in {\mathbb{R}}^{H\times W\times C}$，MobileViT首先通过1x1卷积层产生 ${\mathbf{X}}_L\in {\mathbb{R}}^{H\times W\times d}$。该n×n卷积层对局部空间信息进行编码，而逐点卷积通过学习输入通道的线性组合将张量投影到一个高维空间（或d维空间，其中 $d>C$）。

通过MobileViT，我们希望在具有 $H\times W$ 的有效感受野的同时建模长程的非局部依赖。为了建模长程依赖性，一种常见的方法是使用扩张卷积。然而，这种方法需要小心选择扩张率，否则权重将应用于填充的零，而不是有效的空间区域（Yu & Koltun, 2016；Chen等，2017；Mehta等，2018）。另一种有前景的解决方案是自注意力（Wang等，2018；Ramachandran等，2019；Bello等，2019；Dosovitskiy等，2021）。在所有自注意力方法中，带有多头自注意力的Vision Transformers（ViTs）已被证明在视觉识别任务中有效。然而，ViTs模型通常比较重，并且在优化性上表现不佳，因为它们缺少空间归纳偏置（Xiao等，2021；Graham等，2021）。

为了使MobileViT能够在具有空间归纳偏置的情况下学习全局表示，我们将 ${\mathbf{X}}_L$ 展开为N个不重叠的展开补丁 ${\mathbf{X}}_U\in {\mathbb{R}}^{P\times N\times d}$。这里，$P=wh$，$N=\frac{HW}{P}$ 是补丁的数量，$h\le n$ 和 $w\le n$ 分别是补丁的高度和宽度。对于 $p\in \{1,\cdots ,P\}$，有：

$${\mathbf{X}}_G(p)=\text{Transformer}({\mathbf{X}}_U(p)),1\le p\le P（1）$$

温馨提示：
阅读全文请访问"AI深语解构" MOBILEVIT: 轻量级、通用且适用于移动设备的视觉Transformer