[Transformer]On the Integration of Self-Attention and Convolution

Paper
Code待开源

Abstract

卷积和自注意力是特征学习的两种强大的技术，但是他们通常被看做相互独立的梁总方法。本文则认为二者之间潜在的联系十分密切，可以通过同样的操作完成计算。具体来说，本文首先证明了kxk大小的卷积核计算合一分解为k^2个单独的1x1卷积，然后计算他们的注意力权重将他们的值聚合起来。
因此卷积核自注意力第一阶段的操作是十分相似的，更重要的是与第二阶段相比，第一阶段占据了主要的计算复杂度，是通道数的平方。这一观察结果自然产生了将二者看似截然不同的操作如何优雅的继承在一起的探讨，因此本文提出一个混合模型（ACMixer）同时具备自注意力和卷积的优势，同时与纯卷积或纯自注意力相比又有更小的计算开销。大量的实验结果表明本文在图像识别和下游任务中取得了具有竞争力的结果。

Section I Introduction

近年来卷积和自注意力机制在计算机视觉方面取得了长足的发展，CNN被广泛用于图像识别、语义分割、目标检测等任务；另一方面自注意力最初用于自然语言处理，在视觉领域如超分辨率、图像生成等也展现出强大的性能。近期随着ViT的提出，基于注意力的模块可以获得与CNN相近甚至超出CNN的性能。

尽管卷积和自注意力均获得了优异的心梗但是卷积核自注意力模块通常遵循不同的设计范式，传统的卷积计算的是卷积核内内容的加权和，而卷积核是全局共享的，从而为特征图添加了重要的归纳偏执；与此相比，自注意力计算的是全局信息的加权平均，通过计算每对元祖之间的关联关系获得注意力权重，这种灵活性使得自注意力模块可以关注不同的区域，捕获更多的特征信息。
考虑到卷积和自注意力之间的差异和互补之处处，通过将二者结合可以充分发挥二者的优势。


之前也有工作探索将卷积与自注意力结合，如：
SENet,CBAM等证实了自注意力机制可以增强卷积模块；
SAN,BoTNet等则使用自注意力模块来代替卷积模块；
AA-ResNet，Container等工作则是同时使用两种机制，但依旧分成独立的路径，没有将二者真正结合。

本文试图解释卷积与自注意力之间的密切关系，通过将二者的操作分解，发现他们都严重依赖于1x1卷积，因此本文提出一种混合模型-ACMIx，可以优雅的将二者几个在一起。具体来说，本文会首先将输入经过1x1卷积映射然后获得丰富的特征表述，然后在复用中间特征进行不同方式的聚合。
ACMix可以有效发挥二者的优势同时还不用计算两次。
本文的贡献总结如下：

（1）揭露了卷积与自注意力之间的潜在关系，为理解二者之间的联系提供了全新的视角，也为设计新的学习范式提供了灵感；

（2）提供了一种优雅的继承方法，可以结合卷积和自注意力的好处。实验结果表明这一混合模型的性能优于纯卷积或纯注意力模型。

Section II Related Work

卷积神经网络通过卷积核提取局部特征，自注意力机制也在语言任务中表现出良好的性能，理论分析表明容量足够时自注意力可以表示任意卷积层的函数。近期也有工作探究如何将注意力机制引入视觉任务，目前有两种主流方法：一种是将自注意力作为网络的模块，另一种是将自注意力与卷积互补。

Part 1 Self-Attention only

基于自注意力强大的捕获长程依赖关系，大量工作致力于将自注意力作为基本模块来搭建视觉任务的模型，一些研究表明注意力可以完全替代卷积操作来构建网络模型，比如ViT已经在图像识别、目标检测、分割任务等显示出良好的性能。

Part 2 Attention enhanced Convolution

也有的研究表明，自注意力能够克服卷积的一些局限，因此许多研究人员探究使用注意力模块来增强卷积网络，比如SE操作，GE操作会重新为每个通道分配权重；BAM和CBAM则会分别计算空间上和通道上的权重。此外相关工作还有AA-ResNet，BoTNet等。

Part 3 Convolution enhanced Attention

自从ViT提出后也有诸多变体来进一步改进，有的研究就是使用卷积来增强Transformer模型，从而引入归纳偏执。比如CvT在tokenization过程中就使用了卷积，并且使用步长卷积来降低自注意力的计算复杂度；还有使用CNN stem加在ViT早期从而使得训练更稳定；CSWin则是采用了基于卷积的位置编码，显示出在下游任务中的良好迁移能力。Conformer则是将Transformer与一个CNN模型相结合，从而结合两类特征。

Section III Revisting Convolution and Self-Attention

首先回顾一下自注意力和卷积，为了更好的捕捉二者之间的联系，本文从全新的视角将他们的计算过程分解为不同阶段。




Part 1 Convolution

卷积是卷积神经网络中最核心的一部分，Fig 2(a)展示了卷积操作的详细过程，假设步长均为1.对于kxk的卷积核，标准的卷积操作可以表示为：


其中K表示某位置的卷积核权重，并且（1）式可以重写为不同位置的加权和：


为了进一步简化式子用fbar表示相对位移，因此可以进一步重写为：



因此标准卷积可以分为两个阶段：

Stage1计算特征图每一个位置与卷积核的线性投影，这与1x1卷积相同


Stage 2:投影后的特征根据卷积核的位置不断移动，最终聚合到一起。



可以看到绝大部分计算量集中在1x1卷积部分，后面的移位和聚合是很轻量级的。

Part 2 Self-Attention

注意力机制也广泛用于各种任务，与传统卷积相比，注意力机制允许模型关注更大上下文范围中更重要的区域，参考Fig 2(b)。

考虑包含N个头的自注意力模块，输入为CinxHxW,输出为CoutxHxW，则注意力的计算可以表示为：

||表示会将N个头的结果级联起来，其中注意力部分的计算表示为：



依旧可以分解为两个阶段：




Stage 1会进行1x1卷积将输入映射为k,v,q



Stage 2则是计算注意力权重矩阵，最后与value相乘得到最终的加权和。



与第一阶段相比，stage2的计算成本也是比较小的，与卷积的模式相同。

Part 3 Computational Cost

为了充分理解卷积核自注意力模块的计算瓶颈，首先分析一下FLOPs以及每一阶段的参数量。详情参见Table 1.


对于卷积来说，stage 1的计算复杂度是通道数的平方项，stage 2的则是C线性复杂度，没有额外的训练参数。
在自注意力模块也是类似。并且随着网络模型的增长这一复杂度越来越凸显。

为了进一步验证本文的分析，本文还总结了在ResNet50中实际的卷积核自注意力模块的计算成本。可以看到卷积操作99%的计算成本在stage1，自注意力操作83%的成本在stage 1.这与本文的理论分析是一致的。

Section IV Relating Self-Attention with Convolution

通过前一节的分析将卷积和SA分解成两阶段，第一阶段十分详细，主要就是用于特征学习，都可以通过1x1卷积实现，负责将特征投影到更深的空间；第二阶段则对应于特征聚合，虽然二者学习的范式有所不同。
从计算的角度来看，主要的计算量集中在stage1阶段。
综上所述，有以下结论：

（1）卷积和SA在特征映射部分的操作是相同的，均通过1x1卷积实现，也是主要的计算量所在；

（2）stage2的特征聚合虽然对捕获语义特征至关重要，但是这一阶段十分轻量级，不需要额外学习的参数。

Part 1 Integration of Self-Attention and Convolution

下面介绍如何将卷积与SA进行优雅整合。前面说到卷积和SA第一阶段都是1x1卷积，所以可以只映射一次，然后在对中间特征使用不同的聚合操作。Fig 2©展示了ACMixer的示意图，可以看到也被分成了两个阶段：

Stage 1:特征经过三个1x1卷积映射并且reshape到N pieces，一共获得3XN个feature map


Stage 2:使用不同的聚合操作。



对于SA，将特征分成N组，每一组包含3个特征图，分别作为q,k,v然后进行SA的计算。
对于卷积，则使用一个轻量级的全连接层产生k^2个特征图，通过移动聚合生成最终的输出，这样就像传统卷积会从局部感受野中收集信息。
最后将两条路径的输出加载一起，并且通过两个权重因子来调控二者的比重。











Part 2 Improved Shift and Summation

前面说回对卷积进行shift+summation操作，虽然这一操作是轻量级的但是向不同方向移动破坏了数据的局部性，很难做到矢量化实现，可能会极大降低模型的实际推理效率。作为补救措施，本文使用固定核尺寸的深度卷积。
如果某卷积核表示为：


对应的输出表示为：

通过精心设计特定方向上的卷积核权重，可以使得卷积输出等价于张量的唯一操作。为了进一步获得不同方向上向量的特征和，本文分别将所有的输入特征与卷积核连起来，将shift操作视作单个的组卷积。这样替换可以大大提升计算效率。

在此基础上本文还进行了一些调整来增强模型的灵活性。比如将卷积核的权重看做可学习的，这样一方面可以提高模型的容量一方面保证了模型原本的shift操作，本文还是用了多组卷积核来匹配卷积的路径数和SA的路径数。

Part 3 Computational Cost of ACMixer

Table 1也展示了ACMixer的FLOPs和参数量，可以看到stage1的计算成本和训练参数与SA相同，比卷积更加轻量级；


Stage2引入了额外的全连接层和组卷积，复杂度是C的线性次，在ResNet50上的参数量也与理论分析相近。

Part 4 Generalization to Other Attention Modes

许多研究也集中在探索SA操作变体来提升效率，如将SA的计算集中在窗口内的Swin Transformer，可以降低计算成本，加快推理速度，还有ViT,DeiT则是考虑在单层保留全局依赖。以上修改都证明是有效的，本文提出的ACMixer是独立于SA计算的，也可以适用于以上变体，简单来说，注意力权重的计算可以总结为：






Section V Experiments

本节测试ACMixer在图像分类、语义分割、目标检测上的性能。

Part 1 Results

Fig 4展示了分类任务的对比结果，ACMixer具有最小的FLOPs和参数量，同时精度也最高。


Table 2,3,4则分别展示了语义分割、目标检测的对比结果。可以看到在这些下游任务中ACMixer也优于基于CNN或基于Transformer的模型，同时依旧最优最少的FLOPs和参数量。



Part 2 Practical Inference Speed

本文进一步探究了ACMixer在Ascend910环境下的实际推理速度，这是一个用于移动和边缘计算的深度学习计算框架，Table 5是最终结果，与PVT-S相比，ACMixer达到了1.3x fps；模型规模也大这一表现越优秀。

Part 3 Ablation Study

Combining the output of both paths


Table 6展示的是消融实验的结果。主要 探究两条路径的组合权重，以及将Swin中的注意力替换为卷积操作等可以看到conv+SA的结合方式始终 优于单一计算方式；此外固定两条路径的权重也会导致性能变差，使用可学习的权重使得ACMixer灵活性更高，卷积核SA路径的强度可以根据网络进行自适应的调整。




Group Convolution Kernels

本文还验证了使用组卷积的好处，参见Table 7使用组卷积代替shift会大大提升推理速度，而使用可学习的卷积内核和精心进行初始化可以增强模型的灵活性，有助于最终性能的提升。


Part 4 Bias towards Different Paths

注意到本文还训练了两个参数，代表卷积路径和SA路径的权重，代表了模型在不同那个深度对卷积和SA的倾向性。

Fig5展示了不同层下二者权重的变化，可以看到模型早期更加倾向使用卷积，模型中间阶段倾向二者的结合并且逐渐偏向卷积；最后阶段则是更加倾向自注意力。

通过对卷积和自注意力的倾向性进行可视化，希望能激发我们设计更好的模型来继承二者。

Section VI Conclusion

本文探究了卷积和自注意力这两种强大技术之间的密切关系，通过将二者的计算过程分解，证明了其在特征映射阶段具有相同的计算开销。基于此本文提出一种混合模型，通过共享相同的映射操作来将二者继承。
在图像分类、目标检测等任务中证明了ACMixer的有效性。

本文希望这一工作能激发更多的思路，并引领集合卷积和自注意力的新方向。

Appendix

Part 1Model Architecture

Table 9-12分别展示了基于ResNet,SAN,pvt,Swin等的ACMix模型，为了公平对比本文只将原始的3x3卷积或SA模块替换成本文的ACMix





Part 2 Dataset and Training Setup

ImageNet：1000类图像 1.28million训练 50000test




SGD batch_size=256 8GPU




COCO：80000 training: 35000 validation





ADE20K：




ADE20K包含150类别 20000training:2000validation:3000testing

Part 3 Hyperparameters

heads=4

划分组数为4，每组包含C/4个通道
并且使用了较为流行的相对位置嵌入：








Part 4 Practical Costs for Other Models

本文还总结了基于卷积、SA和ACMix的其他模型的FLOPs和参数量，详情参见Table 8.可以看到超过60%的计算量集中在stage 1，这也证明了ACMix只引入了较小的计算代价就集成了卷积和SA。