MobileViT——论文简述

一、目的

如今，自注意力模型（如vision transformer）成为了CNN的一种替代，要想提高其性能，需要增加网络中参数的数量，但这势必导致模型太贵和训练缓慢的问题。ViT不像CNN那样容易优化，它需要大量的数据增强以及正则化来减少过拟合；同时对于密集的预测任务而言花费又太大。

为了将自注意力模型应用到移动端，作者指出ViT模型需要大量参数的原因是缺少空间归纳偏置，而这在CNN中是本来就存在的，故而提出将CNN与transformer模型相结合的方法，同时又能减少模型参数。

在移动端的模型需要具备低权重、低时延和准确性。由于移动端的特点，FLOPs的优化（DynamicViT）显得没那么重要，作者关注 light-weight、general-purpose和low latency三个方向进行模型设计。

MobileViT能用更少的参数以及更简单的训练方法来进行更好的特征学习。

二、模型结构

如上图所示是一个标准的ViT模型。P=wh（patch的宽和高）。

如上图所示是MobileViT的架构图。其核心思想是用transformer学习全局特征，这使得在网络中隐含地纳入类似卷积的属性（如空间归纳偏置），用简单的训练方法（如基本的增强），并容易将MobileViT与下游架构（如用于语义分割的DeepLabv3）整合。

2.1 MobileViT block

MobileViT block的目的是用更少的参数对全局和局部的信息进行建模。其中n = 3，h = w = 2。

（1）对于一个 H x W x C 的输入张量，先用 n x n 大小的卷积编码局部空间信息，再用点卷积通过学习输入通道的线性组合将张量投射到 d 维，其中 d > C。

（2）对于 H x W x d 的张量，将其展平成 N 个 P x d 的patch，其中 P = wh ， N = HW / P，再经过 L 个transformer。由于Mobile即不丢失patch的顺序，也不丢失每个patch中各个像素之间的空间位置，所以可以将其 fold 成 H x W x d 的张量。

（3）接着用点卷积将张量映射到 C 维，并且与最初的输入进行拼接，再用 n x n 大小的卷积进行特征融合。

unfold 成的 P x N x d 通过卷积编码了 n x n 大小的局部空间信息，通过transformer后的张量的每一个位置 p(1 ≤ p ≤ P) 又编码了 P 个patch的信息（每个红色的像素与对应蓝色像素做attention），如下图所示。

标准的卷积操作可以看作：(1) unfolding，(2) matrix multiplication (学习局部特征)， (3) folding 三种操作之和，MobileViT block也是类似，只不过将matrix multiplication替换为一系列的transformer层（全局处理），结果就是MobileViT也拥有卷积的类似属性。

res = x
fm = self.local_rep(x)

# convert feature map to patches
patches, info_dict = self.unfolding(fm)

# learn global representations
for transformer_layer in self.global_rep:
    patches = transformer_layer(patches)

# [B x Patch x Patches x C] -> [B x C x Patches x Patch]
fm = self.folding(x=patches, info_dict=info_dict)

fm = self.conv_proj(fm)

fm = self.fusion(torch.cat((res, fm), dim=1))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

2.2 多尺度采样

对于ViT模型，一个标准的方法就是微调。同CNN一样，MobileViT不需要进行位置编码，多尺度的输入对训练也是有帮助的。论文中描述了对于不同分辨率图像 batch size 的选择。

实验中采用S = { (160; 160); (192; 192); (256; 256); (288; 288); (320; 320) }

三、结果

有一点可以确定的是，MobileViT比MobileNetv2要运行的慢。

（1）在GPU上存在专门的CUDA内核用于transformer，以提高其在GPU上的可扩展性和效率；

（2）CNN受益于设备层面的优化，包括batch normalization与卷积层的融合，这些优化改善了延迟和内存访问。

然而，为transformer设计的优化操作目前还不能用于移动设备。因此，移动设备上MobileViT和基于ViT的网络是次优的。