南京大学开源！ResT：高效Transformer架构！
共 3700字，需浏览 8分钟
·
2021-06-18 10:59
↑↑↑点击上方蓝字，回复资料，10个G的惊喜
paper: https://arxiv.org/abs/2105.13677
code: https://github.com/wofmanaf/ResT
本文是南京大学提出的一种高效Transformer架构:ResT，它采用了类似ResNet的设计思想：stem提取底层特征信息、stages捕获多尺度特征信息。与此同时，为解决MSA存在的计算量与内存占用问题，提出了EMSA模块进一步降低计算量与内存消耗。所提ResT在图像分类、目标检测以及实例分割等任务均取得了显著的性能提升，比如在ImageNet数据上，在同等计算量前提下，所提方法取得了优于PVT、Swin的优异性能，实乃一种强力骨干网络。
Abstract
本文提出一种高效多尺度Vision Transformer：ResT，它可作为图像中识别的通用骨干架构。不同于现有采用固定分辨率+标准Transformer模块的Transformer模型，它有这样几个优势：
(1) 提出了一种内容高效的多头自注意力模块，它采用简单的深度卷积进行内存压缩，并跨注意力头维度进行投影交互，同时保持多头的灵活性； (2) 将位置编码构建为空域注意力，它可以更灵活的处理任意分辨率输入，且无需插值或者微调； (3) 并未在每个阶段的开始部分进行序列化，我们把块嵌入设计成重叠卷积堆叠方式。
我们在图像分类与下游任务上对所提ResT进行了性能验证，实验结果表明：所提ResT大幅优于现有骨干架构，比如，ResNet18(69.7%)、PVT-Tiny(75.1%)相似大小的模型下，所提方法取得了79.5%的top1精度，这说明它是一种强有力的骨干网络。
Method
上图给出了ResT的架构示意图，它具有与ResNet相似的结构方案，比如采用stem模块提取底层特征，后接四个stage捕获多尺度特征。每个stage包含三个成分：一个块嵌入模块，一个位置编码模块以及L个高效Transformer模块。具体来说，在每个stge开始前，块嵌入模块用于降低输入的分辨率并扩展通道维度；位置编码用于约束位置信息提升块嵌入的特征提取能力；然后将所得送入到后续高效Transformer模块中。
Rethinking of Transformer Block
标准的Transformer模块由MSA与FFN以及残差链接构成，在MSA与FFN之前还采用LN。对于输入token，每个Transformer模块的输出表示如下：
其中，MSA的单头SA与FFN的定义分别如下：
MSA与FFN的计算复杂度分别为, 。
Efficient Transformer Block
如前所述，MSA有两个缺点：(1) 计算量随平方增长，这会导致较大的训练与推理负载；(2) MSA的每个头仅负责输入的部分子集，这会影响模型的性能，尤其当通道维度非常小时。
为解决上述问题，我们提出了上图所示的高效多头自注意力模块。可以看到：
类似MSA，EMSA首先采用投影集合得到Q； 为压缩内存，2D输入将被reshap为3D形式，然后送入深度卷积以因子降低空域维度; 将上述所得特征reshape为2D形式并送入后两个投影集合得到K与V； 然后采用下面公式计算注意力，注：Conv为卷积，用于对不同头进行信息交互。为补偿Conv导致的多样性素食，我们在Softmax之后添加了IN。
最后，每个头的输出进行拼接并线性 投影构成最终的输出。
EMSA的计算复杂度为，具有比MSA更低的计算量。此时，高效Transformer模块定义如下：
Patch Embedding
标准的Transformer采用一序列token作为输入，以ViT为例，3D图像需要拆分为的块，这些块再平展为2D形式并映射为隐嵌入。然而，这种直接的tokenization难以捕获底层特征信息(比如边缘、角点)。此外，ViT中的tokens长度是固定的，这使其难以进行下游任务(比如目标检测、实例分割)适配。
为解决上述问题，我们构建了一种高效多尺度骨干ResT用于稠密预测。正如前面所提到的，每个阶段的高效Transformer模块在同尺度同分辨率上跨通道、空域维度进行处理。因此，块嵌入模块同样需要渐进的扩展通道维度，同时降低空域分辨率。
类似于ResNet，我们采用stem模块以倍率4收缩宽高维度。为高效捕获底层特征信息，我们引入了一种简单而有效的方式：堆叠三个卷积，stride分别为212，前两个后接BatchNorm与ReLU。在234阶段，采用块嵌入模块下采样空间分辨并提升通道维度，这与stride=2的卷积作用类似。
Position Encoding
位置编码对于序列顺序的探索非常关键，ViT一文将可学习参数加到输入tokens中编码位置信息。假设为输入，表示位置参数，那么编码后输入表示如下：
然而，此时要求位置长度与输入tokens长度相同，这无疑会限制了其应用。
为解决上述问题，我们需要设计一种新的变长位置编码，我们将上式修改为如下：
其中表示组线性操作，组数为c。
除了上述形式外，我们还可以采用更灵活的注意力机制得到像素级权值。因此，我们提出了一种简单且高效的像素注意力(Pixel-wise Attention,PA)模块进行位置编码。具体来说，PA采用采用深度卷积计算像素权值，然后采用sigmoid激活，那么带PA的位置编码可以描述如下：
由于每个stage的输入token通过卷积得到，我们可以将位置编码嵌入到块嵌入模块中，整体结果见上图。注：这里的PA可以采用任意空域注意力替换，这使得ResT中的PE极为灵活。
Linear Head
分类头采用全局均值池化+线性分类器的方式，ResT的架构配置信息见下表。
Experiments
接下来，我们在常用基准任务上进行所提方案验证，包含ImageNet数据上的图像分类、COCO数据上的目标检测与实例分割等。
Image Classification
上表给出了图像分类任务上不同方案的性能对比，从中可以看到：
在小模型方面，ResT-small凭借相似的复杂度以79.6%精度大幅超过PVT-T的75.1%； 在中等模型方面，ResT-base凭借相似复杂度以81.6%超过Swin-T的81.3%； 在大模型方面，ResT-Large凭借相似复杂度以83.6%精度超过Swin-S的83.3%； 相比ConvNet，如RegNet，所提ResT凭借相似复杂度取得了更佳的性能； 总而 言之，在不同复杂度模型方面，ResT均显著优于现有模型。
Object Detection and Instance Segmentation
上表给出了RetinaNet架构下的不同骨干模型在目标检测上的性能对比，可以看到：
在小模型方面，相比PVT-T，ResT-Small取得了2.8的指标提升； 在大模型方面，相比PVT-S，ResT-Base取得了0.8的指标提升。
上表给出了实例分割任务上的性能对比，可以看到：
在小模型方面，相比PVT-T，ResT-Small取得了1.8boxAP指标提升，1.0MaskAP指标提升； 在大模型方面，相比PVT-S，ResT-Base分别取得了2.1与1.9的指标提升。
Ablation Study
接下来，我们对所提ResT进行消融实验分析，主要从stem、EMSA、PE三个角度进行对比分析。
从下图的Table5可以看到：ResT中的stem比PVT、ResNet中的Stem更加高效，分别取得了0.92%、0.64%的性能提升。
从上图的Table6可以看到：
当移除卷积操作与IN后，模型性能下降1.16%，这说明长序列与灵活性的组合对于注意力非常重要； 当移除IN后，模型同样出现了大幅性能下降，我们将其归因于不同头之间的多样性遭到了破坏。
上表对比了不同PE的性能对比，从中可以看到：
当移除PA编码后，模型性能从72.88%下降到71.54%，这说明位置编码对于ResT非常重要； LE与GL具有相似性能，而PA以0.84%精度优于GL，这说明：空域注意力可用于进行位置编码建模。
也可以加一下老胡的微信 围观朋友圈~~~
推荐阅读
（点击标题可跳转阅读）
老铁，三连支持一下，好吗？↓↓↓