MPViT : Multi-Path Vision Transformer for Dense Prediction

论文地址: https://arxiv.org/pdf/2112.11010.pdf
code: https://git.io/MPViT

一、引言

密集的计算机视觉任务，如目标检测和分割需要有效的多尺度特征表示来检测或分类不同大小的目标或区域。虽然卷积神经网络(CNN)一直是这类任务的主要架构，但最近引入的Vision Transformer(ViTs)旨在取代它们成为骨干。与cnn类似，vit构建了一个简单的多阶段结构(即，从细到粗)，用单尺度的补丁进行多尺度表示。在本研究中，从不同于现有Vision Transformer的视角，探索了多尺度patch嵌入和多路径结构，构建了多路径Vision Transformer(MPViT)。MPViT采用重叠卷积块嵌入的方法，将相同大小(即序列长度)的特征与不同尺度的块同时嵌入。然后，通过多条路径将不同尺度的标记独立地送入Transformer编码器，并聚合生成的特征，从而在同一特征级别上实现细特征和粗特征表示。得益于多样化、多尺度的特征表示，MPViTs从微小(5M)到基本(73M)的缩放，在ImageNet分类、对象检测、实例分割和语义分割方面，始终比最先进的Vision Transformer取得更好的性能。这些广泛的结果表明，MPViT可以作为多种视觉任务的多功能骨干网。
在目标检测和分割等密集预测任务中，在多个尺度上表示特征以区分不同大小的对象或区域是至关重要的。现代CNN骨干在密集预测方面表现出更好的性能，它利用了卷积核级别或特征级别的多个尺度。Inception Network或VoVNet利用同一特征级别的多粒度卷积核，产生不同的感受野进而提高检测性能。HRNet通过在整个卷积层中同时聚合细和粗特征来表示多尺度特征。
尽管CNN模型被广泛用作密度预测的特征提取器，但目前最先进的(SOTA)Vision Transformer的性能已经超过了CNN。vit变体关注的是在应用于高分辨率密集预测时如何解决自我注意的二次复杂度，而不太关注构建有效的多尺度表示。例如，继传统的CNN之后，最近的Vision Transformer主干使用单一尺度的补丁(即令牌)构建了一个简单的多级结构(如从细到粗结构)。CoaT通过使用协同尺度机制同时表示细和粗特征，允许跨层注意并行进行，从而提高检测性能。然而，协同尺度机制需要大量的计算和内存开销，因为它为基本模型(例如，CoaT-Lite)增加了额外的跨层注意。因此，ViT体系结构的多尺度特征表示仍有改进的空间。
尽管CNN模型被广泛用作密集预测的特征提取器，但目前最先进的(SOTA)Vision Transformer的性能已经超过了CNN。vit变体关注的是在应用于高分辨率密集预测时如何解决自注意的二次复杂度，而不太关注构建有效的多尺度表示。例如，继传统的CNN之后，最近的Vision Transformer主干使用单一尺度的补丁(即令牌)构建了一个简单的多级结构(如从细到粗结构)。CoaT通过使用协同尺度机制同时表示细和粗特征，允许跨层注意并行进行，从而提高检测性能。然而，协同尺度机制需要大量的计算和内存开销，因为它为基本模型(例如，CoaT-Lite)增加了额外的跨层注意。因此，ViT体系结构的多尺度特征表示仍有改进的空间。
这项工作重点是如何有效地表示多尺度的特征与Vision Transformer的密集预测任务。受CNN模型的启发，利用多粒度卷积核用于多个感受野，本文提出了一种Vision Transformer的多尺度贴片嵌入和多路径结构方案，称为多路径Vision Transformer(MPViT)。如下图所示，多尺度贴片嵌入通过重叠卷积操作将不同大小的视觉贴片同时令牌化，适当调整卷积的填充/步幅，得到具有相同序列长度(即特征分辨率)的特征。然后，来自不同尺度的令牌被独立地并行地送入Transformer编码器。每个具有不同大小补丁的Transformer编码器执行全局自注意。然后，生成的特征被聚合，在同一特征级别上实现精细和粗糙特征表示。在特征聚合步骤中，引入了一个全局到局部特征交互(GLI)过程，该过程利用卷积的局部连通性和Transformer的全局上下文，将卷积的局部特征与Transformer的全局特征连接起来。

综上所述，本文的主要贡献如下:
1、提出了一种多径结构的多尺度嵌入算法，用于同时表示密集预测任务的细、粗特征。
2、引入全局到局部特征交互(GLI)来利用卷积的局部连通性和转换器的全局上下文。
3、提供消融研究和定性分析，分析不同路径尺寸和块规模的影响，发现高效和有效的配置。

（一）、Vision Transformers for dense predictions

目前SOTA Vision Transformer的重点是在应用于高分辨率密集预测时，降低自注意的二次复杂度。有些工作用局部区域的细粒度斑块来限制注意范围，并将其与滑动窗口或稀疏的全局注意相结合。也可以通过减少序列长度利用粗粒度的全局自我注意，通过空间缩减(例如，池化)。通过跨特征通道而不是符号操作自注意来实现线性复杂性。具有从细到粗的简单金字塔结构，XCiT具有与ViT一样的单级结构。当应用于密集的预测任务时，XCiT添加向下/向上采样层，在ImageNet上进行预训练后提取多尺度特征。Xu等人介绍了具有简单金字塔结构的CoaT- lite和在CoaT- lite之上具有跨层注意的CoaT。跨层注意使CoaT的性能优于CoaT- lite，但需要大量的内存和计算开销，这限制了模型的扩展性。

（二）、Comparison to Concurrent work

CrossViT也使用不同的补丁大小(例如，小和大)和单级结构中的双路径，如ViT和XCiT。然而，CrossViT的分支之间的交互只通过令牌进行，而MPViT允许不同规模的所有补丁进行交互。此外，与CrossViT(仅分类)不同，MPViT更广泛地探索更大的路径维度(例如，超过两个)，并采用多级结构进行密集预测。

二、实现细节

（一）、Multi-Path Vision Transformer

如上图显示了多路径Vision Transformer(MPViT)架构。构建了一个四阶段特征层次结构，用于生成不同尺度的特征映射。由于多级体系结构具有更高分辨率的特性，因此必然需要更多的计算。因此，由于整个模型的线性复杂性，使用了包括CoaT中所做的因式自我注意的Transformer编码器。在LeViT中，卷积的stem比非重叠的块嵌入表现出更好的底层表示(即，不丢失显著信息)。受到LeViT的启发，给定一个大小为H × W × 3的输入图像，也采用一个由两个3 × 3卷积层组成的stem块，其中通道为$C_2/2$, C2为步幅为2，生成大小为$H/4\times W/4\times C_2$的特征，其中C_2为第2阶段的通道大小。每个卷积之后是Batch Normalization和Hardswish激活函数。从阶段2到阶段5，在每个阶段中堆叠拟议的多尺度补丁嵌入(MS-PatchEmbed)和多路径Transformer(MP-Transformer)块。

(二)、Multi-Scale Patch Embedding

本文设计了一个多尺度的补丁嵌入(MS Patch Embed)层，该层利用了同一特征级别上的细粒度和粗粒度可视标记。为此，使用带有重叠补丁的卷积操作，类似于CNN和CvT。给定来自前一阶段的2d重塑输出特征映射(即令牌映射)作为阶段i的输入:

我们学习一个函数$F_{k\times k}(\cdot )$的map $x_i$变成一个新的带有通道尺度$C_i$的token$F_{k\times k}。(x_i)$其中F(·)是核大小为k × k, stride为s padding为p的二维卷积运算。输出的2D令牌映射$F_{k\times k}(x_i)\in R^{Hi\times Wi\times Ci}$有高度和宽度如下:

卷积补丁嵌入层使我们能够通过改变步幅和填充来调整令牌的序列长度。也就是说，可以用不同的补丁大小输出相同大小(即分辨率)的特征。因此，我们并行地形成了几个具有不同内核大小的卷积补丁嵌入层。例如，如图1所示，可以用3 × 3,5 × 5,7 × 7的patch大小生成序列长度相同的各种大小的视觉标记。
由于在通道和滤波器尺寸相同的情况下叠加连续卷积运算会增大感受野(例如，两个3 × 3等价于5 × 5)，所需的参数更少(例如，$2\times 3^2<5^2$)，因此在实践中我们选择连续的3 × 3卷积层。对于三路径结构，当降低空间分辨率时，我们使用三个连续的3 × 3卷积，通道大小为$C^{{}^{\prime }}$，填充为1，步幅为s，其中s为2，否则为1。由于MPViT的多路径结构具有更多的嵌入层，我们通过采用3 × 3深度可分离卷积来减少模型参数和计算开销，该卷积由3 × 3深度卷积和嵌入层中的1 × 1点卷积组成。所有卷积层后面都是批处理规范化和Hardswish激活函数。最后，将不同大小的令牌嵌入特征分别送入每个Transformer编码器。

(三)、Global-to-Local Feature Interaction

虽然Transformer中的自注意可以捕捉到长期依赖关系(即全局上下文)，但它很可能会忽略每个补丁中的结构信息和局部关系。此外，Transformer受益于形状偏置，允许他们聚焦在图像的重要部分。相反，CNN可以从尺度不变性中利用局部连通性——图像中的每个补丁都用相同的权值处理。这种归纳偏差促使CNN在对视觉对象进行分类时更依赖于纹理而不是形状。因此，MPViT以一种互补的方式将CNN的局部连通性与全局上下文Transformer结合起来。为此，我们引入了一个全局到局部特征交互模块，该模块学习了局部和全局特征之间的交互，以丰富特征表示。具体来说，在第i阶段,代表局部特征的$L_i\in R^{H_i\times W_i\times C_i}$，我们采用一个深度残差瓶颈块，由1 × 1卷积、3 × 3深度卷积和1 × 1卷积组成，通道大小$C_i$相同。利用每个Transformer 的$G_{i,j}\in R^{H_i\times W_i\times C_i}$的二维重塑全局特征。局部和全局特征的聚合是通过串联来实现的：

j是路径的索引,$A_i\in R^{H_i\times W_i\times (j+1)C_i}$为聚合特征,H(·)是一个学习与特征交互的函数。产生最后的特征$X_{i+1}\in R^{H_i\times W_i\times (j+1)C_{i+1}}$带有下一级的维度特征$C_i$。对于H(·)，我们使用通道为$C_{i+1}$的1 × 1卷积。最后的特征$X_{i+1}$作为下一阶段多尺度嵌入层的输入。

(四)、Model Configuration

为了减轻多路径结构的计算负担，使用CoaT中提出的高效因子化自注意:

其中$Q,K,V\in R^{N\times C}$为线性投影查询，键，值，N, C分别为令牌个数和嵌入维数。为了保持相似的参数和flop，增加路径的数量需要减少通道C或层数L(即Transformer编码器的数量)。L个因子化自注意层，N个令牌和h个Transformer编码器头，总时间复杂度为$O(LhN{C}^2)$，显存复杂度为$O(Lh{C}^2+LhNC)$。其复杂度为通道C的二次方，而层数L是线性。因此，从单路径(即，在消融研究中证实，降低C比降低L能获得更好的性能。由于具有较高的特征分辨率，对于三路径模型，阶段2的计算成本相对较高。在第2阶段也将路径数设置为2。因此，从阶段3开始，三路径模型有3条路径。
虽然三路径和双路径在ImageNet分类上的准确率相似，但三路径模型在密集预测任务中表现出更好的性能。这表明扩展路径维的多样化特征对于密集预测任务是有用的。因此，建立了基于三路径结构的MPViT模型。将MPViT模型从对应CoaT-Lite Tiny(5M)或DeiT-Tiny(5.7M)的小尺度MPViT-Tiny (5M)放大到对应SwinBase (88M)的大尺度MPViT- base (74M)。所有MPViT模型使用8个Transformer编码器头，mlp的膨胀为Tiny和其他模型分别设置为2和4。表1描述了MPViTs的详细信息。

三、实验

（一）、ImageNet Classification

遵循DeiT中和其他基线transformer的训练方式。使用AdamW优化器训练300个周期，批大小为1024，权重衰减为0.05,5个热身周期，初始学习率为0.001，该学习率由余弦衰减学习率调度程序缩放。将每张图像裁剪为224 × 224。随机深度下降只在Small和Base大小模型中使用，我们分别将其速率设置为0.05和0.3。更多详情见附录。

(二)、Object Detection and Instance Segmentation

分别用RetinaNet和Mask R-CNN验证MPViT作为对象检测和实例分割的有效特征提取器。我们在COCO数据集上对模型进行基准测试。我们在ImageNet-1K上预训练主干，并将预训练的主干插入到RetinaNet和Mask R-CNN中。按照常用设置和swing - transformer的训练配方，我们使用多尺度训练策略训练3×(36轮次)的模型。使用AdamW优化器，初始学习率为0.0001，权值衰减为0.05。

(三)、消融实验

Exploring path dimension：

为了与基线进行公平的比较，不使用stem、随机深度下降路径和第3.3节中介绍的卷积局部特征。对于双路径，在第2阶段，更高的特征分辨率需要更多的计算，因此我们减少层数L(即编码器的数量)。在第5阶段，更高的嵌入维数导致更大的模型尺寸，因此我们也降低了L和嵌入维数C，在第4阶段增加了L。由于多路径导致更高的计算成本，在第3和第4阶段减少C来补偿。因此，表5中的双路径(a)在具有相似的模型大小和略高的FLOPs的情况下，优于单路径。
将双路径扩展为三路径时，我们分别消融了嵌入维数C和层数L。对于表5中(b)的嵌入维数，我们保持C，但降低L，以保持模型尺寸和FLOPs相似，这导致精度低于双路径。反之，当我们降低C并保持L时，©的分类精度与双路算法相似，但检测性能高于双路算法。最后，我们进一步将路径扩展为四路径(d)，保留L，减少c。四路径的分类精度接近，但检测性能并不优于©的三路径。这些结果告诉我们三个教训:
(1)层数(即更深)比嵌入维度(即更宽)更重要，这意味着更深更薄的结构在性能上更好。
(2)多粒度标记嵌入和多路径结构可以为目标检测器提供更丰富的特征表示。
(3)在模型尺寸和FLOPs相同的约束下，三路径是最佳选择。
我们注意到，扩展路径维度的策略不会增加内存负担，如表5所示。双路径(a)和三路径(b,c)比单路径占用更少的内存。此外，(a)和(b)比©消耗更多的内存，因为(a)和(b)在阶段3和4有更大的c。这是因为C(二次型)在内存使用中比L(线性)更重要。因此，我们降低嵌入维数和扩展路径维数和层(更深)的策略导致了一个内存高效模型。但是，由于多路径结构导致的总层数的增加降低了推理速度。
Multi-Scale Embedding:

在上表中，我们研究了patch大小和结构在多尺度嵌入中的影响。使用三个卷积层并行地使用相同的步幅2和分别为3、5和7的补丁大小。即，每个路径嵌入层使用前面的输入特征独立运行。为了提高参数效率，我们还使用三个卷积层串联，其核大小分别为3，步长分别为2,1,1。我们注意到后者具有等效的感受野(例如，3、5、7)。串联版本在减少模型尺寸和flop的同时，比并行版本改进了。直观地说，这种性能提升可能来自于这样一个事实，即系列版本实际上包含具有非线性的小3层CNN，这允许更复杂的表示。
Global-to-Local feature Interaction:
GLI模块中实验了不同的聚合方案，聚合了卷积局部特征和全局变压器特征，测试了两种类型的操作:相加和拼接。如表6所示，求和运算没有显示性能提升，而拼接则显示分类和检测任务都有改进。直观地说，在1×1卷积之前对特征求和会很天真地混合这些特征，而拼接则保留了它们，允许1×1卷积学习特征之间更复杂的相互作用。这一结果表明，GLI模块有效地学习了本地和全局特征之间的交互，以丰富表示。

四、附录