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

作者单位：Apple

论文：

https://arxiv.org/abs/2110.02178

代码：

GitHub - apple/ml-cvnets: CVNets: A library for training computer vision networks

1 Introduction

传统CNN易于优化且可根据特定任务整合不同网络，ViT则需要大规模的数据且更难优化，学习量大且计算量大，这是因为ViT缺乏图像固有的归纳偏差。

结合CNN和ViT的优势，为移动视觉任务建立一个轻量级、低延迟、精确率高的网络，能够满足设备的资源限制，并能在不同的任务上有很好的泛化效果。

CNN的优点：空间归纳偏置，对数据增强手段的依赖性较低

ViT的优点：全局处理

更好的性能：对于给定的参数预算，MobileViT 在不同的移动视觉任务（图像分类、物体检测、语义分割）中取得了比现有的轻量级 CNN 更好的性能。

更好的泛化能力：泛化能力是指训练和评价指标之间的差距。对于具有相似训练指标的2个模型，具有更好评价指标的模型更具有通用性，因为它可以更好地预测未知数据集。与CNN相比，即使有广泛的数据增强，其泛化能力也很差，MobileViT显示出更好的泛化能力。

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

 

​

 

2 Architecture

2.1 网络整体架构

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

带有↓2标记的皆为下采样块。

n=3，h=w=2

​

Output stride: Ratio of the spatial dimension of the input to the feature map

2.2 MV2（MobileNetV2）：Inverted Residuals and Linear Bottlenecks

首先，我们将V1与V2进行对比。

MobileNet V1用深度可分离卷积提取特征，然后用1*1卷积来扩张通道。减少了运算量和参数量。在网络结构上没有采用shortcut的方式。

由于ReLU运算容易造成信息丢失，因此在MobileNet V2中将最后1*1卷积后的RELU改为linear。这就是论文标题中的“Linear bottlenecks”。

并在做深度卷积之前使用1*1卷积先进行升维。

​

然后我们再将v2与resnet进行对比，

ResNet中的residuals block是：“压缩”→“卷积提特征”→“扩张”，

MobileNetV2则是“Inverted residuals”,“扩张”→“卷积提特征”→ “压缩”

​

针对stride=1 和stride=2，在block上有稍微不同，主要是为了与shortcut的维度匹配，因此，stride=2时，不采用shortcut。

​

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

本文中，用MV2 block从高维向低维转换，expansion factor为4，用Swish激活函数。

​

​

 

 

2.3 MobileVit-Block

在MobileViT中，提出MobileViT block来同时编码局部和全局的信息。

Block的输入首先经过一个n*n卷积层编码局部空间信息，再经过1*1卷积层将维度变为d（d>c），得到输出XL(H x W x d)。

接着就是MobileViT block的核心：Transformer as Convolutions（可以理解为用transformer来做卷积），也就是图中中间框里的部分。

Transformer as Convolutions 模块的输出XF（H x W x d）经过一个1*1的卷积来降低维度到C，由concat与原始输入X相连。然后进行n*n卷积用于结合局部和全局特征。

 

2.3.1 Transformer as Convolutions（global representations）

用于对远距离非局部依赖性进行建模。

目前常用的建模远程依赖关系的方法是空洞卷积，但是需要谨慎选择空洞率；否则权重将应用于填充的 0 而不是有效的空间区域。

另一种常用的是方法是self-attention，但是重量级，缺乏空间归纳偏置（spatial inductive bias）。

因此，本文中使用的方法是transformer as convolutions。

 

文中将标准卷积操作描述为：unfolding —> local processing （矩阵乘法）—> folding。

MobileViT block将卷积中的local processing用global processing（L个堆叠的transformer块）来替换。因此MobileViT具有类似于卷积的属性(如空间偏差)。

Block可以看作是利用Transformer来做卷积。因此在可以在较小参数量、简单训练策略的情况下，在不同的视觉任务上达到与轻量CNN相同的性能。

 

具体流程为：

输入的XL通过Unfold操作得到了维度为Xu(shape:P*N*d)的输出（即N个non-overlapping flattened patches），其中 P = h*w 为所切图片块的大小（h,w < H,W），N = HW/P为图片块总数，接着我们通过L个堆叠的Transformer对Xu进行全局信息提取，其输出为XG(P*N*d)，紧接着我们通过Fold操作得到与XL维度一致的XF(H*W*d)。

​

因为XU通过卷积编码了n*n区域的局部信息，XG中的每个像素点可以编码X中所有像素点的信息，因此整体的等效感受野为整张图片的大小，即H*W。

​

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

2.3.2 Unfold & Fold

ViT：

​

vit输入的图片数据是 H x W x C，reshape为N*PC，再project为N*d。（其中，P=hw是每个patch的h*w，N是patch的数量）送入transformer层中。

MobileViT：

unfold 相当于ViT中输入数据处理过程。

输入的XL通过Unfold操作得到了维度为Xu(shape:P*N*d)的输出（即N个non-overlapping flattened patches），其中 P = h*w 为所切图片块的大小（h,w < H,W），N = HW/P为图片块总数，接着我们通过L个堆叠的Transformer对Xu进行全局信息提取，其输出为XG(P*N*d)，紧接着我们通过Fold操作得到与XL维度一致的XF(H*W*d)。

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

2.4 Light-Weight分析

MobileViT Block使用标准卷积和Transformer分别学习局部和全局表示。因为之前的工作表明，使用这些层设计的网络是重量级的，为什么MobileViT是轻量级的?

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

MobileViT使用卷积和Transformer的方式使生成的MobileViT Block具有类似卷积的属性，同时允许全局处理。这种建模能力能够设计出shallow and narrow的MobileViT模型，从而使其weight更轻。

MobileViT  L= {2, 4, 3}  d={96, 120, 144} at spatial levels 32 × 32, 16 × 16, and 8 × 8.

The resulting MobileViT network is faster (1.85×), smaller (2×), and better (+1.8%) than DeIT network.

3 Experiment results

3.1 Image classification

​

​

3.2 Mobile object detection

​

3.3 Mobile Semantic Segmentation

​

3.4 Performance on mobile devices

​