京东AI：用于视觉识别的上下文Transformer网络-Contextual Transformer Networks for Visual Recognition

论文下载：https://arxiv.org/pdf/2107.12292.pdf

源码下载：https://github.com/JDAI-CV/CoTNet

Abstract

       使用自注意力机制的Transformer已经导致了自然语言处理领域的革命，最近出现了Transformer-style的架构设计，并在许多计算机视觉任务中取得了有竞争力的结果。然而，大多数现有的设计直接在二维特征图上使用自注意力，以获得基于每个空间位置的孤立的query-key对的注意力矩阵，但没有充分利用相邻键之间的丰富上下文。在这项工作中，我们设计了一种新的transformer-style模块，即上下文Transformer块（CoT），用于视觉识别。这种方法充分利用输入键之间的上下文信息来指导动态注意力矩阵的学习，增强了视觉表征能力。技术上，CoT块首先通过一个3×3的卷积对输入键进行上下文编码，从而得到输入的静态上下文表示。我们进一步将编码的键keys与输入查询queries拼接起来，通过两个连续的1×1卷积学习动态多头注意力矩阵。学习到的注意力矩阵乘以输入的values，以实现输入的动态上下文表示。静态和动态上下文表征的融合最终被视为输出。我们的CoT块很吸引人，因为它可以很容易地取代ResNet架构中的每一个3×3卷积，产生一个被称为上下文Transformer网络(CoTNet)的transformer-style backbone。通过在广泛应用(例如图像识别、目标检测和即实例分割)上的大量实验，我们验证了CoTNet作为更强主干的优越性。

1. Introduction

       卷积神经网络(CNN)展示了学习辨别性视觉表征的高能力，并令人信服地很好地泛化到一系列计算机视觉(CV)任务，例如图像识别、目标检测、语义分割，CNN架构设计的实际方法是基于离散卷积算子(例如，3×3或5×5卷积)，它有效地利用了空间局部性和平移等变性。然而，卷积的有限感受域对全局/长程依赖性的建模产生了不利影响，这种长程交互为许多CV任务提供了子服务。最近，自然语言处理(NLP)领域在强大的语言建模架构[14,49]中兴起了具有自注意力的transformer，它以可伸缩的方式触发长距离交互。受此启发，通过整合基于CNN的架构和transformer-style的模块，CV任务的极限得到了稳定的突破。例如，ViT[15]和DETR[5]使用Transformer中self-attention直接处理图像patches或CNN输出。[40,58]提出了一种局部自注意力模块的独立设计，它可以完全取代ResNet架构中的空间卷积。然而，以前的设计像传统的self-attention块(如图1a)主要依赖于独立的query-key对来计算注意力矩阵，从而忽略相邻键之间的丰富上下文。

图1 比较传统的自注意力和我们的上下文Transformer(CoT)块。(a) 传统的自注意力仅利用孤立的query-key对来测量注意力矩阵，但未利用keys之间丰富的上下文。相反，(b)CoT block首先通过3×3卷积挖掘keys之间的静态上下文。接下来，基于query和上下文key，利用两个连续的1×1卷积进行自注意力，生成动态上下文。静态和动态上下文的融合作为最终输出

       在这项工作中，我们提出了一个简单的问题--有没有一种很好的方法通过利用二维特征图上输入键之间的丰富上下文来增强Transformer-style的体系结构？为此，我们提出了一种独特的Transformer-style块设计，命名为上下文Transformer(CoT)，如图1 (b)所示。这种设计将键间的上下文挖掘和二维特征图上的自注意力学习结合在一个单一的体系结构中，从而避免了为上下文挖掘引入额外的分支。从技术上讲，在CoT块中，我们首先通过对3×3网格中的所有相邻键执行3×3卷积来将keys的表示上下文化。上下文化的键特征key feature可以被视为输入的静态表示，它反映了局部相邻键之间的静态上下文。在此基础上，我们将上下文化的键特征和输入查询拼接后送入到两个连续的1×1卷积，以生成注意力矩阵。这个过程自然地利用了每个query和所有keys之间的相互关系，在静态上下文的指导下进行自注意力学习。进一步利用学习到的注意力矩阵聚合所有输入的values，从而实现输入的动态上下文表示，以描述动态上下文。我们将静态和动态上下文表示的组合作为CoT块的最终输出。总之，我们的出发点是在输入键上同时捕捉上述两种空间上下文，即通过3×3卷积的静态上下文和基于上下文化自注意力的动态上下文，以促进视觉表征学习。

       我们的CoT可以看作是一个统一的构建块，是现有ResNet结构中标准卷积的替代方案，无需增加参数和计算负担。通过直接用CoT块替换ResNet结构中的每个3×3卷积，我们提出了一种用于图像表示学习的新的上下文transformer网络(称为CoTNet)。通过一系列CV任务的广泛实验，我们证明我们的CoTNet优于一些最先进的backbones。值得注意的是，对于ImageNet上的图像识别，CoTNet相对于ResNeSt(101层)的top-1错误率绝对降低了0.9%。对于COCO数据集上的目标检测和实例分割，CoTNet绝对提高了ResNeSt的1.5%和0.7%的mAP。

2. Related Work

2.1 Convolutional Networks

       由于通过AlexNet[29]在ImageNet数据集上的突破性性能，卷积网络(ConvNet)已成为CV领域的主导架构。ConvNet设计的一种主流遵循LeNet[30]的主要规则，即向更深的方向连续堆叠由低到高的卷积：8层的AlexNet、16层的VGG[43]、22层的GoogleNet[46]、152层的ResNet [22]。此后，人们提出了一系列convnet架构设计的创新，以增强convnet可视化表现能力。例如，受Inception模块中split-transform-merge策略的启发，ResNeXt[53]在同一拓扑中使用聚合的residual reansformations对ResNet进行升级。DenseNet[27]还支持跨层连接，以提高ConvNet的能力。SENet[26,25]没有利用ConvNet[28,37]中的空间依赖性，而是捕获通道之间的相互依赖性，以执行通道特征重新校准。[47]进一步扩展自动搜索的ConvNet，以获得一系列EfficientNet的网络，从而实现更高的准确性和效率。

2.2. Self-attention in Vision

       从Transformer中的self-attention中得到启发，在各种NLP任务中不断获得令人印象深刻的表现，研究界开始更加关注视觉场景中的self-attention。NLP域[49]中原始的的self-attention机制被用来捕获序列模型中的长程依赖。在视觉领域，self-attention机制从NLP到CV的简单迁移是在图像中不同空间位置的特征向量上直接进行self-attention。特别是，在ConvNet中探索self-attention的早期尝试之一是非局部操作[51]，它作为一个额外的块，将self-attention应用于卷积的输出。[3] 通过全局多头自注意力机制进一步增强卷积算子，以便于图像分类和目标检测。[24,40,58]在局部区域local patch（例如，3×3网格）中使用自注意力，而不是在整个特征图[3,51]中使用全局自注意力。这种局部自注意力的设计有效地限制了网络消耗的参数和计算，因此可以完全取代整个深层架构中的卷积。最近，通过将原始图像重新塑造成一维序列，采用序列transformer[7]来自动回归预测像素，以进行自监督表征学习。接下来，[5,15]直接将纯transformer应用于局部特征序列或图像块，用于目标检测和图像识别。最近，[44]设计了一个强大的backbone，用全局自注意力层替换ResNet中最后三个3×3的卷积。

2.3. Summary

       在这里，我们还将重点探讨视觉backbone结构设计中的自注意力。大多数现有的技术直接利用传统的self-attention，因此忽略了相邻键之间丰富的上下文的显式建模。相比之下，我们的上下文transformer块在一个简单的结构中统一了keys之间的上下文挖掘和基于特征图的自注意力学习，并具有良好的参数预算。

3. Our Approach

       在本节中，我们首先简要回顾了在视觉backbones中广泛采用的传统自注意力。接下来，介绍了一种新的transformer-style块，称为上下文transformer(CoT)，用于图像表示学习。这种设计超越了传统的自注意力机制，通过额外利用输入键之间的上下文信息来促进自注意力学习，最终提升了深层网络的表征特性。在整个深层体系结构中使用CoT块替换3×3卷积后，进一步阐述了两种上下文transformer网络，即分别源自ResNet[22]和ResNeXt[53]的CoTNet和CoTNeXt。

3.1 Multi-head Self-attention in Vision Backbones

图2  (a)传统的自注意力块和(b)我们的上下文化Transformer块的详细结构  和分别表示元素求和和局部矩阵乘法。 

       在这里，我们提出了视觉backbones中可缩放局部多头自注意力的一般公式[24,40,58]，如图2(a)所示。形式上，给定一个大小为H×W×C(H:高度，W:宽度，C:通道数)的输入二维特征图X，我们通过embedding矩阵(Wq，Wk，Wv)分别将X转换为queries Q=XWq，keys K=XWk和values V=XWv。值得注意的是，每个embedding矩阵在空间中实现为1×1卷积。然后，我们得到keys K和queries Q之间的局部关系矩阵:

这里Ch是head的数量，表示局部矩阵乘法运算，该运算测量每个查询query与空间中局部k×k网格内的对应关键字keys之间的配对关系。因此，在R的第i个空间位置的每个特征是一个k×k×Ch维向量，它由所有heads的Ch个局部query-key关系映射(size：k×k)组成。局部关系矩阵R进一步丰富了每个k×k网格的位置信息：

其中表示在每个k×k网格内的二维相对位置embeddings，并且在所有Ch heads之间共享。接下来，通过将增强的空间感知局部关系矩阵与沿通道维度每个head的Softmax运算进行归一化来获得注意力矩阵A：A=Softmax()。在对A的每个空间位置特征向量进行reshape为Ch个局部注意力矩阵(size：k×k)后，最终的输出特征图被计算为每个k×k网格内的所有values与学习的局部注意力矩阵的聚合：

值得注意的是，每个head的局部注意力矩阵仅用于聚集通道维度的V的均匀划分的特征图，最终输出Y是所有heads聚集的特征图的拼接。

3.2. Contextual Transformer Block

       传统的自注意力很好地触发了不同空间位置之间的互动，这取决于输入本身。然而，在传统的自注意力机制中，所有成对的query-key关系都是在独立的query-key对上独立学习的，没有探索其间丰富的上下文。这严重限制了视觉表征学习的二维特征图上的自注意力学习能力。为了解决这个问题，我们构建了一个新的transformer-style构建块，即图2(b)中的上下文transformer(CoT)块，它将上下文信息挖掘和自注意力学习集成到一个统一的体系结构中。我们的出发点是充分利用相邻keys之间的上下文信息，以有效的方式促进自注意力学习，并增强输出聚合特征图的代表性。

       特别地，假设我们有相同的输入二维特征图。keys，queries和values分别被定义为K=X，Q=X和V=XWv。CoT块不是像典型的自注意力那样通过1×1卷积对每个key进行编码，而是首先在所有相邻keys上采用k×k群卷积，在空间上采用k×k网格，以便将每个key表示上下文化。学习到的的上下文化keys 自然地反映局部相邻键之间的静态上下文信息，我们将作为输入X的静态上下文表示。然后，在将上下文化的键和queries Q拼接的条件下，通过两个连续的1×1卷积（Wθ带有ReLU激活函数，Wδ没有激活函数）实现注意力矩阵： 

       换句话说，对于每个head，A的每个空间位置的局部注意力矩阵是基于查询特征和上下文化的键特征学习的，而不是基于单独的query-key对。这种方式通过挖掘的静态上下文的额外指导来增强自注意力学习。接下来，根据上下文化的注意力矩阵A，我们将典型的self-attention中的所有values V聚集，计算出注意力特征图： 

       鉴于参与特征图捕捉了输入之间的动态特征交互，我们将其命名为输入的动态上下文表示。因此我们的CoT 块的最终输出(Y)为静态上下文和通过注意力机制获得的动态上下文的融合。

 3.3. Contextual Transformer Networks

表1 ResNet-50(左)和CoTNet-50(右)的详细结构 括号内显示了残差块的形状和操作，外部列出了每个stage中堆叠blocks的数量。CoTNet-50的参数和FLOPs比ResNet-50略少 

表2 带有32×4d模板的ResNeXt-50(左)和带有2×48d模板的CoTNeXt-50(右)的详细结构。括号内显示了残差块的形状和操作，外部列出了每个stage中堆叠blocks的数量。C表示了分组卷积的组数。与ResNeXt-50相比，CoTNeXt-50的参数数量稍多，但FLOPs相似

       我们的COT的设计是一个统一的自注意力块，并作为ConvNet中标准卷积的替代。因此，用CoT替代卷积来增强上下文化自注意力的视觉backbones是可行的。在这里，我们展示了如何在不显著增加参数预算的情况下将COT块集成到现有的最先进的ResNet架构中(例如，ResNet[22]和ResNeXt[53])。表1和表2显示了基于ResNet-50/ResNeXt-50主干的两种不同的上下文transformer网络(CoTNet)结构，分别称为CoTNet-50和CoTNeXt-50。请注意，我们的CoTNet可以灵活地泛化到更深层次的网络(例如，ResNet-101)。

       CoTNet-50:具体来说，CoTNet-50是通过直接用CoT块替换ResNet-50中的所有3×3卷积(在res2、res3、res4和res5阶段)构建的。由于我们的CoT块在计算上与典型的卷积相似，CoTNet-50与ResNet-50具有相似(甚至略小)的参数数和FLOPs。

       ​CoTNeXt-50:类似地，对于CoTNeXt-50的构造，我们首先用CoT块替换ResNeXt-50组卷积中的所有3×3 卷积核。与典型卷积相比，当组数(即Table 2中的C)增加时，组卷积内的kernels深度显著降低。在ResNeXt-50中，组卷积的计算代价因此减少了一个C因子。因此，为了实现与ResNeXt-50相似的参数个数和FLOPS，我们还将CoTNeXt-50的输入特征图的比例从32×4d减小到2×48d。最后，CoTNeXt-50只需要比ResNeXt-50多1.2倍的参数和1.01倍的FLOPs。

3.4. Connections with Previous Vision Backbones

       在本节中，我们将详细讨论我们的上下文Transformer和之前最相关的视觉backbones之间的关系和区别。

       Blueprint Separable Convolution可分离卷积[18]使用1×1逐点卷积加上一个k×k深度卷积来近似传统卷积，旨在减少沿深度维度的冗余。一般来说，这种设计与transformer-style模块（例如，典型的self-attention和我们的CoT模块）有一些共同之处。这是因为transformer-style块还利用1×1点卷积将输入转换为values，然后在类似的深度管理器中使用kxk的局部注意力矩阵进行聚合计算。此外，对于每个head，transformer-style block中的聚合计算采用通道共享策略，在没有任何显著精度的情况下高效实现。这里，所使用的通道共享策略也可以解释为绑定块卷积[52]，它在相同的通道块上共享相同的filters。

       Dynamic Region-Aware Convolution动态区域感知卷积[6]引入了一个滤波生成器模块（由两个连续的1×1卷积组成），用于学习不同空间位置的区域特征的专用滤波器。因此，它与CoT块中的注意力矩阵生成器具有相似的作用，该生成器可以为每个空间位置实现动态的局部注意力矩阵。然而，[6]中的滤波器生成器模块根据主要输入特征图生成专门的滤波器。相比之下，我们的注意力矩阵生成器充分利用了上下文化keys和queries之间的复杂特征交互，实现了自注意力学习。

       Bottleneck Transformer[44]是同时代的工作，它也旨在通过transformer-style模块取代3×3卷积，用自注意力机制增强ConvNet。具体来说，它采用了全局多头自注意力层，这在计算上比我们的CoT块中的局部自注意力更昂贵。因此，对于同一个ResNet backbone，[44]中的BoT50只使用Bottleneck Transformer块替换最后三个3×3卷积，而我们的CoT block可以完全替换整个深层架构中的3×3卷积。此外，我们的CoT block在[24,40,58]中通过在输入keys之间探索丰富的上下文来加强自注意力学习，从而超越了典型的局部自注意力。

5. Conclusions

       在这项工作中，我们提出了一种新的Transformer-style体系结构，称为上下文Transformer（CoT）块，它利用输入keys中的上下文信息来引导自注意力学习。CoT block首先捕获相邻键之间的静态上下文，进一步利用该上下文触发挖掘动态上下文的自注意力。这种方式将上下文挖掘和自注意力学习巧妙地结合到一个单一的体系结构中，从而增强了视觉表征能力。在现有的ResNet架构中，我们的CoT块可以很容易地取代标准卷积，同时保留有利的参数预算。为了验证我们的说法，我们通过在ResNet架构（例如ResNet或 ResNeXt）中替换3×3卷积来构建上下文Transformer网络（CoTNet）。在ImageNet上学习的CoTNet架构验证了我们的建议和分析。在COCO上进行的目标检测和实例分割实验也证明了经过我们的CoTNet预先训练的视觉表示的泛化。