MobileVIT：轻量级视觉Transformer+移动端部署

方向提出

轻量级卷积神经网络的spatial inductive biases（ 空间归纳偏置）使他们能够在不同的视觉任务中学习参数较少的表征，然而这些网络在空间上是局部的。反观重量级的基于self-attention的VIT可以学习到全局表征，所以作者等人针对轻量化提出一个问题： 是否有可能结合CNN 和 ViT 的优势，为移动视觉任务建立一个轻量级和低延迟的网络？因此本文提出了一个新的视角：使用基于Transformer的卷积来处理全局信息。精度提升可见下图：

归纳偏置可以理解为：从现实生活中观察到的现象中归纳出一定的规则，然后对模型做一定的约束，从而可以起到“模型选择”的作用，即从假设空间中选择出更符合现实规则的模型。

引言

ViT是学习视觉表征的卷积神经网络的替代品，但其性能提升是以模型大小(网络参数)和延迟为代价，许多生活中的应用（如增强现实技术和自动轮椅技术）的资源受限，即使ViT模型大小减小到与移动设备的约束资源相匹配，也会出现精度大幅下降，且性能不如传统CNN模型的情况，故设计轻量化ViT不可避免。

传统CNN易于优化且可根据特定任务整合不同网络，ViT则需要大规模的数据且更难优化，学习量大且计算量大，这是因为ViT缺乏图像固有的归纳偏差。故本文结合CNN和ViT的好处来实现设计一个针对于移动视觉任务的轻量级、通用级和低延迟级的网络。

总的来说，MobileViT的优势有：

• 更好的性能：对于给定的参数预算，MobileViT 在不同的移动视觉任务（图像分类、物体检测、语义分割）中取得了比现有的轻量级 CNN 更好的性能。
• 更好的泛化能力：泛化能力是指训练和评价指标之间的差距。对于具有相似训练指标的2个模型，具有更好评价指标的模型更具有通用性，因为它可以更好地预测未知数据集。与CNN相比，即使有广泛的数据增强，其泛化能力也很差，MobileViT显示出更好的泛化能力。

• 更好的鲁棒性：一个好的模型应该对超参数具有鲁棒性，因为调优这些超参数会消耗时间和资源。与大多数基于ViT的模型不同，MobileViT模型使用基本增强训练，对L2正则化不太敏感。

MobileViT

传统ViT预热

如图，我们将输入图片切块后进行Flatten展平，后经过一个Linear层将维度缩放到d，在添加位置编码后通过L个Transformer Block学习到不同特征，最后经过Linear层糅合输出预测。需要注意此处Self-Attention的计算成本是 O() 。

由于VIT忽略了空间归纳偏差，所以它们需要更多的参数来学习视觉表征。此外，与CNN相比，VIT及其多种变体的优化性能不佳，这些模型对L2正则化很敏感，需要大量的数据增强以防止过拟合。

MobileVit-Block

如图为MobileVit整体结构，我们可以看到输入图片(HxWxC)经过一个普通卷积层（Conv3*3）输入到连续的五个MV2中，当（H，W）为32*32时来到了网络的核心部分：MobileVit-Block，接着Block与MV2交叉堆叠，最后一个Block的输出通过一个Conv-1*1 + 全局池化来到了全连接层，紧接着得到最终的输出。

图片中凡注有 ↓2标记的皆为下采样块，其中MV2指MobileNet-v2块，具体将在后文介绍。

接下来我们一起剖析网络的核心部分：MobileVit-Block（如下图所示）

图6

Block的输入X (shape: HxWxC)我们接着采用图5的数据，即H=32，W=32。

从左向右看，Block的输入经过一个连续的Conv-n*n和Conv-1*1，得到了维度为XL(H x W x d)的输出，接着来到了Block最核心的部分：Transformer as Convolutions（global representations），如下图：

n*n卷积层编码局部空间信息，而点卷积层通过输入通道的线性组合将张量投影到一个d维空间（d>c） 核心部分主要用以解决对于远距离非局部依赖性进行建模，而另一种被广泛研究的建模远程依赖关系的方法是扩张卷积（空洞卷积），但是扩张卷积需要谨慎选择膨胀率；否则权重将应用于填充的 0 而不是有效的空间区域。

图7文中将标准卷积操作描述为：unfolding —> local processing —> fold，此处我们可以看到 local processing 部分使用L个堆叠的Transformer块来进行替换。

因此MobileViT具有类似于卷积的属性(如空间偏差)。Block可以看作是Transformer的卷积。其优势在于，卷积和Transformer的低层次高效实现可以开箱即用；并允许在不同的设备上使用MobileViT而不需要任何额外的负担。

输入的通过Unfold操作得到了维度为(shape:PNd)的输出（即N个non-overlapping flattened patches）（此处的unfold相当于ViT中的输入数据处理过程），其中 P = h*w 为所切图片块的大小（h,w < H,W），N = HW/P为图片块总数，接着我们通过L个堆叠的Transformer对进行全局信息关注，其输出为，紧接着我们通过Fold操作得到与维度一致的，至此Transformer as Convolutions部分结束，我们回到图6。

经过一个11的卷积来降低维度，得到了与原始输入X维度一致的 ，此处的因经过Transformer as Convolutions部分故含有全局信息，我们将其与Block的原始输入X进行一个concat的操作，之后我们再进行nn的卷积，就相当于同时获取图片的局部和全局信息，如下图所示：

因为我们的

每个图片块都含有来自其他所有图片块的信息，故可以毫不夸张的说，我们每进行一次nn的卷积，其真实有效感受野为整张图片的大小，即HW。

图8

在图中，红色像素通过Transformer处理蓝色像素(其他图片块中相应位置的像素)。因为蓝色像素已经使用卷积对邻近像素的信息进行了编码，这就允许红色像素对图像中所有像素的信息进行编码。在这里，黑色和灰色网格中的每个单元分别表示一个图片块和一个像素。

至此Block部分讲解完毕，针对上述产生的问题：

1. Block中的Unfold操作是什么？以及最后的Fold是什么？
2. 如何实现轻量化？
3. 什么是MV2？

下面我们逐一进行分析。

问题解决

Unfold & Fold

前文有提到：unfold 相当于ViT中输入数据处理过程。故在这里我们一起来回顾ViT——数据处理：

首先来看ViT的部分结构图：

图9

从图片的左下角开始看起，我们看到的是一个个被切分好的图片块，这里需要对输入作出解释：

假设原始输入的图片数据是 H x W x C，我们需要对图片进行块切割，假设图片块大小为P1 x P2，则最终的块数量N为：N = （H/P1）x（W/P2）。

这里需要注意H和W必须是能够被P整除的

接下来到了图一正中间的最下面，我们看到图片块被拉成一个线性排列的序列，也就是“一维”的存在（以此来模拟transformer中输入的词序列，即我们可以把一个图片块看做一个词），即将切分好的图片块进行一个展平操作，那么每一个向量的长度为：Patch_dim = P1 x P2 x C。

经过上述两步操作后，我们得到了一个N x Patch_dim的输入序列。

与MobileViT进行对比，即下图：

图10

对比分析可得，上图中的P*C对应着上文公式中的Patch_dim，N即为图片块的数量。具体的代码实现也极其简单（将在后文代码部分介绍中详细阐述）

同样的Fold操作可以看做Unfold的反向操作，即将图片块重新放置回原来的位置原来的维度。

Light-Weight分析

MobileViT Block使用标准卷积和Transformer分别学习局部和全局表示。因为之前的工作表明，使用这些层设计的网络是heavy-weight，一个自然的问题出现了:为什么MobileViT是轻量级的?

作者认为问题主要在于通过Transformer学习全局表示。对于给定的patch，之前的工作是通过学习像素的线性组合将空间信息转化为潜在信息。然后通过使用Transformer对全局信息进行编码学习patches之间信息。因此，这些模型失去了图像特定的归纳偏差。

因此需要更多的参数来学习视觉表征，所以那些模型又深又宽。与这些模型不同的是，MobileViT使用卷积和Transformer的方式使生成的MobileViT Block具有类似卷积的属性，同时允许全局处理。这种建模能力能够设计出浅而窄的MobileViT模型，从而使其weight更轻。

MobileViT和vit中multi-head self-attention的计算成本分别为和。理论上，与ViT相比，MobileViT效率较低。然而在实践中，MobileViT比ViT更有效率。

什么是MV2？

MV2即MobileNetV2，其具体模块结构如下所示：

图14

下图更为详细的表明了模块结构：

图15

从高维向低维转换，使用ReLU激活函数可能会造成信息丢失或破坏（故此处不使用非线性激活函数）

因为Expansion layer 和 Projection layer都有可学习参数，故整个网络结构可以学习到如何更好的扩展数据和重新压缩数据。

下图为MobileNetV2整体结构，但在MobileViT中我们只使用其子块。

结语

本文针对ViT的轻量化研究创新性的提出Transformer as Convolutions，糅合CNN与ViT优势，又弃其糟粕。同时在论文中作者还对MobileViT的移动设备应用（如iPhone）以及不同视觉任务应用（如图像分类、物体检测、语义分割）进行了实验研究，对于训练方法作者采用多尺度采样训练，见下图，并使用通用化实验说明其有效性。