On the Integration of Self-Attention and Convolution

On the Integration of Self-Attention and Convolution

本文是2022 CVPR的一篇关于将卷积和self-attention进行组合，从而提高下游任务性能的文章。卷积和self-attention是图像特征提取时常用的两大范式，然而二者的优缺点也都十分明显，为了融合二者的优点，本文从一种新的角度对二者进行了比较，并提出了一种新的图像特征提取范式，这种范式同时有着卷积和self-attention的优势，又将计算量控制在可接受范围之内。

首先我们先来看一下作者是如何来对比卷积和self-attention的：
对于卷积，作者将其公式化为了：

假设输入和输出特征的通道数都是$C$，卷积核的大小是$k\times k$的，其中$g_{i,j}$是卷积后特征图$(i,j)$处的特征，$f_{i,j}$是输入特征图$(i,j)$处的特征，$K_{p,q}$是一个$C\times C$的矩阵。那么上式可以解释为对于位置为$(i,j)$处的输出特征，其获得方式是卷积核的每个位置（是一个矩阵）与输入特征的对应位置（是一个向量）相乘，得到的$k\times k$个向量再相加。作者将这个过程拆分为了两个阶段，第一阶段：特征提取，可以看成是用$k\times k$个$1\times 1$的卷积对特定位置进行计算，公式如下
第二阶段，上述卷积对全局进行扫描（shift），并且每次求得的这$k\times k$个特征进行相加：
在上述两个过程中，第一个过程涉及到的计算量是：$k^2$个$[C,C]\times [C,1]=k^2{C}^2$，参数量是$k^2{C}^2$；第二个过程涉及到的计算量是： $k^2$个$[C,1]$相加$=k^{2}C$，没有额外参数。
为了更明显的看出两个阶段的差异，作者ResNet50为例，第一阶段和第二阶段的浮点数运算和参数量分别如下表所示
之后，作者将self-attention也进行了类似的公式化：

这里是一个经典的mult-head self-attention，其中head数量为$N$，对于每一个head，先是通过两个参数矩阵$W_q,W_k$与特征图$f$相乘获得两个特征特异性矩阵，之后特异性矩阵相乘计算相似性，获得attention权重矩阵$A(W_q{f}_{i,j},W_k{f}_{ab})$，最后将$W_v{f}_{ab}$和attention权重矩阵相乘并将每处的结果相加，就得到了在$(i,j)$处的self-attention值。之后将这$N$个head得到的结果concat在一起，就是最终的输出。同样，作者也将这个过程拆分成了两部分，第一部分是三个参数矩阵$W_q,W_k,W_v$对特征图的卷积，可以表示为：

第二部分是三个得到的卷积结果做self-attention，可以表示为：

为了公平比较，假设attention范围是$k\times k$的区域。那么可以看出，第一步的计算量就是$3$个$[C,C]\times [C,1]=3C^2$，参数量也是$3C^2$，这一步可以看成是用三组$k^2$的卷积核对每个位置进行$k^2$次$1\times 1$的卷积；第二步的计算量是（1）获得attention权重$k^2个[1,C]与一个[C,1]$求相似性；（2）依照相似性对$k^2个[C,1]$特征进行相加，其计算量是$2k^{2}C$，但是没有新的参数。
同样以ResNet50为例，给出两个阶段的浮点运算量和参数：

通过对以上过程两相比较，可以看出二者在第一阶段都在进行多个$1\times 1$的卷积运算，可以将这一阶段理解成特征提取阶段。并且通过分析可以看出，二者的计算量大头和参数量也都在第一阶段，那么如果将第一阶段提取的特征进行重用，就可以在几乎只引入一半计算量的条件下同时实现卷积和self-attention。为此，作者提出了一种新的计算范式ACmix(Attention-Conv-mix)，其流程如下图所示：

ACmix同样是分为两个阶段，第一阶段作者通过三组$1\times 1$的卷积对输入特征图（$H\times W\times C$）进行计算，得到三组$H\times W\times C$的特征图，之后对于每一组特征图将其切片为$H\times W\times C/N\times N$的特征图。在第二阶段，这些特征图分别进入到两个流程中。流程1：attention流程。三组特征图中前两组看成是$Q和K$，计算相似性矩阵，之后将得到的相似性矩阵转换成attention权重矩阵，与第三组进行相乘，完成一组self-attention操作。注意，因为每一组特征图的$N$组切片都是单独计算attention的，因此这一阶段可以看成是N头的self-attention。之后，将得到的N个头进行concat，获得$H\times W\times C$的特征就是流程1的输出；流程2：卷积流程。首先通过一个$[3C,k^2]$的全连接，将输入转换为$H\times W\times 3C--->H\times W\times C/N\times N{k}^2$，之后通过对每一组$H\times W\times C/N\times k^2$个特征向量相加，得到$H\times W\times C/N$的特征，之后将N组特征concat在一起，就得到了一个$H\times W\times C$的特征，这就是流程2的输出。最后将两个流程的输出通过权重相加，就得到了第二阶段的最终输出。
此外还需要提一点的是，在第二阶段的流程2中，在进行特征相加时需要对特征进行shift，如果直接简单的对每个区域进行遍历相加，则会导致结果失去local性以及难以被向量化实现（意味着难以并行运算）。为了解决这个问题，作者在这里提出了一个小技巧，即通过固定核的方式将shift转换为卷积操作。作者给了一个例子，假如我们需要对$(i-1,j-1)$处的一个向量进行相加，那么我们只需要对这个区域乘一个

卷积核即可。进一步，作者又提出可以将固定核衍生为可训练核以及多组可训练核
最后，我们来分析一下ACmix的计算量和参数量：
首先在第一阶段，其计算量和self-attention一样，是3个$[1,C]\times [C,C]=3C^2$，参数量是$3C^2$。在第二阶段，假设attention范围是$k\times k$，则流程1的计算量是$2k^{2}C$，参数量为0。流程2的计算量为 ①、全连接 $[H,W,C/N,3N]\times [3N,N{k}^2]=3C{k}^2（将3C转换成$k^2C$）$， 参数量为$3k^{2}N$；②、特征相加 $k^2个[C/N]=k^{2}C/N$， 这个过程共进行N次，共Ck^2；③、可训练核shift，$k^2$个核进行shift，之后$k^2$个$[1,C/N]$相加，共$C{k}^4$，参数量为$k^{4}C$。
因此其两个阶段的参数量和计算量分别如上表所示。