[Transformer]CAT: Cross Attention in Vision Transformer

from 快手，北邮

Paper
Code

Abstract

Transformer广泛应用于NLP后，在CV领域也引起了广泛关注，凡是将单词token替换为图像的patch使得Transformer计算量大大增加。本文提出一种新的注意力机制-Cross Attention，不再计算全局注意力而是将注意力的计算局限在patch内部来捕获局部信息，在单通道特征图之间的patch应用attention来捕获全局信息。这两种操作都比标准SA所需的计算量要小，替换为这两种操作后本文可以在降低计算成本的同时保持较好的性能，并且为解决视觉任务建立也一种称之为Ctrss Attention Transformer(CAT)的层次网络。

在COCO数据集和ADE20K数据集上的实验结果证明了CAT可以作为一种通用的backbone。

Section I Introduction

随着深度学习的发展和CNN的广泛应用，计算机视觉任务获得了长足的发展，2012年后CNN已经成为主流框架，作为各种视觉任务的关键特征提取器，也有各种CNN变体用于各种下游任务，如目标检测、语义分割等。



Transformer作为一种新的网络结构在NLP领域中取得了显著的成果，鉴于其强大的提取全局信息的能力有效的解决了RNN在处理序列模型时难以并行的问题，这是一个重大的进步，也启发了计算机视觉任务的发展。






近期越来越多的研究将Transformer用于计算机视觉任务，但是由于NLP中文本序列的长度一般都是固定的，因此Transformer迁移到图像任务时能力会受限，因为不同任务的分辨率是不同的。如果将每一个像素点都作为1个token将产生巨大的计算量，因此在ViT,iGPT等将一块区域内的像素作为一个token来处理，这样可以一定程度上减少计算量。但是还是会随着输入分辨率指数增加，但是产生的特征图大小却是一样的，参见Fig 1,使得前述方法无法作为后续任务的backbone。




本文受CNN中提取局部特征的启发，通过计算patch内的注意力来模仿CNN的特性，这样也可以减少SA的计算量。同时为了获得全局信息，使得patch之间有所交流，本文还计算了单通道feature map之间的注意力。与计算所有通道上的注意力相比，计算量也大大减少。





因此原始的注意力计算被替换为patch内注意力计算和单通道特征图的注意力计算这两种操作；基于交叉注意力可以构建一个强大的backbone来生成不同尺度的特征图，满足下游任务对不同粒度特征的需求。本文在不增加计算量的情况下完成了全局注意力的计算，这是一种处理CNN和Transformer联合特征的更合理的方法。
本文的base模型在ImageNet-1K数据集上达到了82.8%的top1精度，可以媲美当前CNN和Transformer的SOTA；同时还作为目标检测和语义分割的backbone可以提升他们的性能。因为CNN和Transformer是互补的，本文希望未来可以将两者更好的结合，CAT就是朝着这个方向迈出的一步，本文希望在这个方向上有更好的发展。

Section II Related Work

CNN/CNN-Based network

CNN具有权重共享、平移不变、局部性等特点，比MLP在CB领域更具优势，已经成为近十年视觉任务的主流框架。代表性网络有AlexNet,ResNet，Mobilenet，ShuffleNet等。
Global attention in

Transformer-based network 


Transformer的核心是MHSA多头注意力机制，可以有效提取不同单词之前的关系。ViT，DeiT是最初视觉Transformer的骨干框架，会将图像切分成patch后嵌入并展平作为输出，patch-size=16；随后借助cls-token进行分类；在CvT和CeiT中使用卷积操作取代了QKV的线性投影；CrossViT则使用两个分支，输入不同大小的patch，从而获得不同尺度的特征聚合得到的全局信息。

但是上述方法会使用所有patch来计算MSA，只关注了不同patch之间的关系；并且随着输入的增加，计算量显著增加，很难用于大分辨率输入的视觉任务。


Local Attention in Transformer-based network 



在视觉任务中patch内部的关系也十分重要。TNT将每一个patch分成更小的patch，这样可以捕获patch内部的信息，作为全局信息的补充；Swin Transformer将每一个patch作为窗口来计算patch内部的相关性，再通过窗口的移动捕获更多的特性。但是两种方法都存在各自的问题，首先将全局信息与局部信息结合起来提取，会使得计算量无法估计；其次局部信息与相邻patch之间的信息交互 这里面缺乏全局信息。



本文提出一种跨patch的自注意力模块来保证全局信息的交互，同时避免了随着输入分辨率的增加而使得计算量爆炸的问题。



Hierarchy networks and downstream tasks




Transformer已成功用于视觉任务和NLP任务，但是因为Transformer的输入与输出大小一样，很难实现类似CNN的层次结构，而这种层次结构在下游任务中至关重要。FPN+ResNet已经成为目标检测任务的标准范式；金字塔特征则是语义分割的标准范式，近期PVT和Swin会在不同stage降低分辨率，这类似ResNet,也是本文采用的方法。

Section III Method

Part 1 Overall architecture

本文的方法旨在将patch内和patch间的注意力结构，通过基础模块的堆叠搭建层次化网络。整体结构参见Fig 2(a).

Step 1:图像切分为patch，进行patch embedding

Step 2:使用多个CAT层获得不同尺度的特征。

patch_size = N

因此patch内包含NNC1的原始信息，patch数目为H1 / N * W1 / N


stage1 处理后输出F1=H1 * W1 * C1



在送入stage2之前会进行patch projection，将2x2xC的像素块reshape到1x1x4C，再通过线性映射到1X1X2C维度
stage2的输出F2= H1/2 * W1/2 * C2 即长宽减半，通道数double
经过4个stage处理后可以得到{F1,F2,F3,F4}四种不同尺度和维度的特征图，与典型的CNN层次结构一样，可以为下游任务提供不同粒度的特征图。





Part 2 Inner-Patch Self-Attention Block

在计算机视觉任务中，不同的通道代表不同的语义特征，token与此类似，但是如果将图像每一个像素都作为一个token则计算量太大了。





参考公式1计算量会随着输入尺寸的平方增长，在RCNN中输入最少也要800像素，YOLO序列最少的图像也超过500像素点，大多数语义分割的边长也有512像素。
计算量是224像素的5倍。
受CNN局部特征的启发，本文引入SA的局部计算-补丁内自注意力（Inner-Patch Self-Attention,IPSA）,如Fig 3(a)所示，将注意力范围局限在patch内，而不是整张图；同时Transforemer会根据输入生成不同的注意力图，这与参数固定的CNN相比有显著优势，类似卷积中的动态参数。这种方法可以显著降低计算量：

将（HW)^ 2 降为N^2HW,N为patch_size。

比如H=W=56，C=96,N=7的设定下，原始SA的计算量约为2.0G,而IPSA只有0.15G。

Part 3 Cross-Patch Self-Attention

在patch内计算SA只能捕获patch内部像素之间的关系，缺失了整张图的信息。CNN网络通过堆叠卷积模块来扩大感受野，或者通过空洞卷积来扩大感受野，Transformer天然可以捕获全局关系，但是诸如ViT，DeiT并不是最优方案。
每个通道的特征图天然具有全局信息，因此本文提出CPSA，将每个通道的特征图分成H/N * W/N个patch，然后计算SA获得全局信息，这有点类似深度可分离卷积，计算复杂度为：

IPSA是将所有通道的输入切成2x2分在IPSA处理后叠加重塑为原始形状，CPSA则是将单通道的输入切成2x2，CPSA处理后恢复为原始形状，示意图参见Fig 3。
通过IPSA和CPSA就可以获得patch内和全局的特征。而Swin Transformer中的移动窗口是手动设计的，难以实现，并且捕获全局信息的能力有限；本文的窗口更加合理。FLOPS约为同样设定下只有0.1G，远远小于2.0G的MSA。


MHSA中每个头部可以注意到单词之间不同的语义关联，在CNN领域不同的头则可以关注patch之间的语义信息，与CNN中的channel类似。在CPSA中 将head设置为patch_size,但实验结果发现并没有带来性能提升，因此本文默认使用单头注意力。

Position encoding


IPSA使用相对位置编码，CPSA使用绝对位置编码

Part 4 Cross Attention based Transformer

Cross Attention Block(CAB)结构如Fig 2(b)所示，包含两个IPSA和一个CPSA。计算流程如下：





本文还比较了卷积做patchembedding层和slice切片的效果，卷积核和步长设置为Patch_size，对比结果参见Table 5，二者性能一样，因此本文默认使用前者。根据每个stage的维度不同搭建了三种规模的CAT，参见Table 1.







Section IV Experiment

进行了COCO数据集目标检测任务和ADE20K语义分割任务。

Part 1 Image Classification

图像分类任务中 细节是CAB使用了dropout以防过拟合

Table2展示了CAT与其他SOTA网络的对比结果。


可以看到FLOPs比ResNet少65%时CAT-T可以到80.3%的top-1精度；CAT-S,CAT-B的精度分别是81.8%和82.8%；而与目前较为先进的SwinT相比，CAT-S提升了0.5%，说明了比Swin的shift window操作具有更强的捕捉patch之间关系的能力。




Part 2 Object Detection

Table 3展示了目标检测的对比结果，可以看到选择CAT-S,CAT-B作为backbone可以使得FCOS提升3.4%，MASK-RCNN提升4.2%,展示了CAT具有更好的特征提取能力。

Part 3 Semantic Segmentation

Table 4展示的是语义分割的实验结果。选择CAT-S,CAT-B作为backbone，FPN进行分割任务，可以看到比ResNet作为backbone提升了4.2%.

Part 4 消融实验

Patch Embedding function


本文首先比较的是patche embedding的方法，分别是卷积、slice两种方法，Table 5的实验结果显示二者具有相同的性能因此本文使用的是步长卷积的方式。


Multi-head和shifted window


前文提到多头会关注单词之间不同的语言特征，本文在CPSA中将头数设置为patch_size，发现对性能没什么提升（参见Table 5的实验结果），本文还测试了shifted-window的方法，发现对CAT也没什么提升。


Table 6展示的CPSA使用不同位置编码的影响，可以看到采用绝对位置编码+一定随机失活效果是最好的。




Section V Conclusion

本文提出的CAT框架可以更好的结合局部特征和全局信息，是一种更为鲁邦的Transformer视觉框架。他可以有效的提供多尺度的特征信息，适用于下游任务。

CAT的核心在于将SA替换为patch内的注意力计算和单通道特征图的注意力计算，在不增加计算成本的前提下更好的捕获局部和全局信息。


本文希望可以作为有效聚合CNN和Transformer的一种有效的尝试。