X-volution On the Unification of Convolution and Self-attention 论文学习

Abstract

卷积和自注意力是DNN中的两个基本模块，前者以线性的方式提取局部图像特征，后者从全局的角度对高阶语义关系做编码。尽管二者在本质上是可以互补的（一阶/高阶），但 CNN 或 Transformers 都无法同时将二者应用在单个计算模块中，因为二者的计算模式不同，图像运算中点积过多。本文，作者系统地推导了一个全局自注意力近似机制，对变换特征，通过卷积操作来近似实现自注意力。基于该机制，作者构建了一个多分支基础模块，由卷积和自注意力组成，能够统一局部和非局部特征的传递。一旦训练好了，这个多分支模块能够通过结构重参数化，转换为一个标准的卷积操作，该算子称作 X-volution，可以作为原子操作而用在任意的网络中。大量实验证明 X-volution 取得了极高的视觉理解性能提升（在 ImageNet 分类任务上为 $+1.2\%$ top-1 准确率，在 COCO 检测和分割任务上为 $+1.7$ box AP 和 $+1.5$ mask AP）。

1. Introduction

对于DNN来说，卷积和自注意力是信息提取和编码最重要的两种计算组成。一方面，卷积是一个线性操作，通过离散的卷积操作（点积）提取局部特征，该操作是对中心点是某像素点/图像特征的图像块$x$和一个窗口形状的滤波器$\mathcal{W}$做运算，即 $\mathcal{W}x+\mathcal{B}$。另一方面，自注意力操作是一个高阶运算，通过对图像中的给定位置和其它所有位置做缩放点积计算，编码出非局部语义信息，即 $softmax(({\mathcal{W}}^{Q}x{)}^T{\mathcal{W}}^{K}x){\mathcal{W}}^{V}x$。

自然就会想到这二者是互补的，即局部和非局部信息，将二者整合到一个算子中是合理的，模型性能可以从中得到提升。卷积采用了局部处理的归纳偏置（inductive bias，空间相近的grid elements有联系而远的没有）和各向同性的特点，使之具备容易优化和平移不变的特性。与卷积相反，自注意力则具备较弱的归纳偏置，使自注意力能在较大的空间内自由地探索数据集的内在特征，从而取得更优的表现和泛化性。此外，将二者整合起来有助于抵消各自的弱点。比如，卷积的局部本质使它无法构建长期联系，但它通常对视觉理解是有用处的。此外，自注意力因为缺乏局部先验，很难优化，故而需要在超大数据集上预训练（如 JFT300M 和 ImageNet-21K）。但这两个算子属于不同的计算模式，很难整合到一起。卷积算子会对局部窗口进行一个预先定义好的加权平均计算，而自注意力则利用带有动态系数的全局加权操作来计算。此外，在视觉领域，计算图像上所有位置的点积在算力上是不允许的，这就给在 CNN 框架中引入非局部操作带来挑战。因此，SOTA 网络要么只进行卷积操作，要么只进行自注意力操作。人们还没有找到一个灵活、有效的方式来融合这两个算子于一个原子操作中，但如果有，那么它很大可能会成为视觉理解任务上重要的计算构成。

本文作者解决了这些障碍，提出了一个新的原子操作，叫做 X-volution，整合了卷积和自注意力。首先，作者在理论上证明了近似实现全局自注意力的可能性，通过从局部区域到全局传播上下文关系。根据这个思路，作者设计了一个新颖的自注意力近似方案，其复杂度是$O(n)$，叫做 Pixel Shift Self-Attention(PSSA)。与现有的自注意力方法相比，PSSA 将自注意力转化为一个卷积操作，通过顺序的 element-wise shift 和点积计算得到变换特征。其次，基于该近似方案，作者构建了一个多分支网络模块，同时整合卷积和自注意力。多分支拓扑结构不仅具有容易优化的特性，也能获取长期关系，因此展现了更强的学习能力和更优的表现。更重要的是，通过网络结构的重参数化，该多分支结构能有条件地转化为一个标准卷积操作，得到一个纯卷积风格的算子 X-volution，可当作原子操作而加入任意一个网络中。实际上，卷积和自注意力可分别看作为静态卷积（content-independent kernel）和动态卷积（content-dependent kernel）。

在多个视觉任务上，作者对 X-volution 进行了大量的定量和定性实验。结果证明 X-volution 算子取得了非常好的提升（即 ImageNet 图像分类：$+1.2\%$ top-1 准确率，COCO目标检测：$+1.7\%$ box AP，COCO实例分割：$+1.5\%$ mask AP）。

2. Related Work

卷积网络在目标检测和分类任务上非常强大。但是，最近的工作表明卷积算子的有效性会受到局部感受野的限制，带有全局信息融合的自注意力机制可以取得更好的性能。受此启发，一些工作加入了复杂的机制来增强网络的表征能力。最近，研究人员意识到自注意力作为一个计算组成，可以很好地处理计算机视觉问题，出现了许多 transformers。Cordonnier 等人[4]发现 multi-head 自注意力能够学习卷积的特征，二者在特殊条件下可以互相转化。Bello 等人没有单独地使用卷积操作，而通过直接 concat 将卷积和自注意力结合起来，取得了不少提升。这说明将二者结合能有效提升表现。Wu等人从另一个角度肯定了这一点。他们在视觉 transformers 中引入了卷积操作，进一步提升了性能。

3. Methodology

作者提出了一个新颖的原子操作，X-volution，将基础的卷积算子和自注意力算子整合到一个统一的计算模块中，期待这样能取得显著的性能提升。这一部分，作者首先回顾了卷积和自注意力的基本数学公式，然后提出了一个全局自注意力的近似形式，直接转化为可兼容卷积的形式。最后作者介绍了，在推理阶段如何将卷积分支和本文提出的自注意力近似操作融合到单个卷积风格的原子操作中。

3.1 Convolution and Self-Attention Operator Revisit

3.1.1 The Convolution Module

卷积算子是构建 CNN 中的重要计算组成部分，通过在一个有限的局部区域内线性加权来预测输出。给定一个特征张量 $\mathbf{X}\in {\mathbb{R}}^{C_i\times H\times W}$，$C_i$表示输入通道个数，$H,W$是高度和宽度。卷积的预测结果$\mathbf{Y}\in {\mathbb{R}}^{C_o\times H\times W}$定义如下：

$${\mathbf{Y}}_{c_o,i,j}\stackrel{def}{=}\sum _{c_i=0}^{C_i}\sum _{({\delta }_i,{\delta }_j)\in {\Delta }_K}{\mathcal{W}}_{c_o,c_i,{\delta }_i+\lfloor K/2\rfloor ,{\delta }_j+\lfloor K/2\rfloor }{\mathbf{X}}_{c_i,i+{\delta }_i,j+{\delta }_j}+{\mathcal{B}}_{c_o}\text{\hspace{1em}(1)}$$

其中$C_o$表示输出通道个数。$\mathcal{W}\in {\mathbb{R}}^{C_o\times C_i\times K\times K}$表示卷积核，其中${\mathcal{W}}_{c_o,c_i,{\delta }_i+\lfloor K/2\rfloor ,{\delta }_j+\lfloor K/2\rfloor }$表示每个位置的卷积核的值。$K$是卷积窗口大小，$\mathcal{B}\in {\mathbb{R}}^{C_o}$表示偏置向量，${\Delta }_K\in {\mathbb{Z}}^2$ 表示$K\times K$卷积窗口所有可能的偏移量的集合。从等式1可以看出，卷积实际上就是一阶线性加权操作。

3.2 The Self-Attention Module

自注意力是最近兴起的视觉任务计算模块，核心思想就是通过向量内积操作构建一个长期语义交互关系，与卷积不同，自注意力无法直接处理图像张量，输入特征张量首先 reshape 为一个向量$\mathbf{X}\in {\mathbb{R}}^{C\times L}$。$L$代表向量长度，$L=H\times W$。${\mathcal{W}}^Q,{\mathcal{W}}^K,{\mathcal{W}}^V$ 分别表示 Query, Key 和 Value 的 embedding 变换，它们是空间共享的线性变换。自注意力计算的方式如下：

$$\mathbf{Y}\stackrel{def}{=}softmax(({\mathcal{W}}^Q\mathbf{X}{)}^T{\mathcal{W}}^K\mathbf{X}){\mathcal{W}}^V\mathbf{X}=\overline{\mathcal{W}}(\mathbf{X})\mathbf{X}\text{\hspace{1em}(2)}$$

其中$\overline{\mathcal{W}}(\mathbf{X})$代表自注意力最终的等价系数矩阵，可看作为一个动态（element-wise content-dependent）的、空间变化的卷积核。等式2说明自注意力具有高阶的全局运算。

3.3 The Approximation of Global Self-Attention

全局自注意力是最重要的注意力机制，得益于全局性，它具备优异的性能优势。但是其复杂度过高，$O(n^2)$（$n$表示所有像素点的个数），使其在视觉任务上应用很有限。问题就变成了，我们是否能找到等式2中$\overline{\mathcal{W}}(\mathbf{X})$的一个合适近似方法，即我们是否能从现有的算子中，比如卷积或逐元素相乘，找到一个与$\overline{\mathcal{W}}(\mathbf{X})$兼容的算子？这一章节，作者证明，通过 element-wise shift 和点积，我们通过卷积的形式可以近似做到全局自注意力算子。给定特征张量$\mathbf{X}$的位置，我们将它的特征向量表示为$x_0$，注意力对数$s_0$可以写作：

$$s_0=\sum _{x_t\in \Omega }{\alpha }_t⟨x_0,x_t⟩=\underset{\text{Local Region}}{\underbrace{\sum _{x_j\in A}{\alpha }_j⟨x_0,x_j⟩}}+\underset{\text{Non-Local Region}}{\underbrace{\sum _{x_i\in (\Omega \backslash A)}{\alpha }_i⟨x_0,x_i⟩}}\text{\hspace{1em}(3)}$$

图1. X-volution 算子的结构。X-volution 训练和推理时的结构是解耦的。它的训练结构展示在中间位置，有两个分支。右边分支有级联卷积和BN组成，可以整合卷积的能力。左边的分支由一个 pixel shift self-attention 组成，提供近似的全局自注意力特征。训练完成后，我们将其重新参数化，形成一个原子操作，如右边所示。在推理时，X-volution 实际上是一个动态卷积算子，权重由卷积权重和注意力图组成（需要被动态计算得到）。

其中${\alpha }_t=w^q{w}^k{w}^v{x}_t$，$\Omega$表示全局区域，$A$表示中心点位于$x_0$的局部区域。作者在图1左边展示了局部区域和非局部区域。灰色框表示输入特征$\mathbf{X}$ 的全局区域，绿色框表示中心点位于$x_0$的局部区域。非局部区域是局部区域以外的区域。因为图像具有很强的 Markovian 性质，$x_0$可以用局部区域的像素点来近似地线性表示：$x_0\approx {\sum }_{x_k\in \mathring{A}}{\beta }_k{x}_k$，其中${\beta }_k$是线性权重。将其代入上面等式的第二项，可以得到下面的等式：

$$\sum _{x_i\in (\Omega \backslash A)}{\alpha }_i⟨x_0,x_i⟩\approx \sum _{x_i\in (\Omega \backslash A)}{\alpha }_i⟨\sum _{x_k\in \mathring{A}}{\beta }_k{x}_k,x_i⟩=\sum _{x_i\in (\Omega \backslash A)}\sum _{x_k\in \mathring{A}}{\alpha }_i{\beta }_k⟨x_k,x_i⟩\text{\hspace{1em}(4)}$$

不失一般性，我们可对系数不为0的区域$A$增加几项。通过设计，非局部区域也出现在了局部区域边界像素点的感受野之内了。因此，作者可将上面等式变换为下面形式：

$$\sum _{x_i\in (\Omega \backslash A)}\sum _{x_k\in \mathring{A}}{\alpha }_i{\beta }_k⟨x_k,x_i⟩=\sum _{x_i\in \Omega }\sum _{x_k\in \mathring{A}}{\alpha }_i{\beta }_k⟨x_k,x_i⟩\text{\hspace{1em}(5)}$$

根据图像中 Markovian 性质，我们可以假设对于$x_k\in A$，$x_i$（与$x_k$相距很远）与$x_k$的关系是很弱的。因此，等式5可以简化为：

$$\sum _{x_i\in \Omega }\sum _{x_k\in \mathring{A}}{\alpha }_i{\beta }_k⟨x_k,x_i⟩=\sum _{x_k\in \mathring{A}}\sum _{x_i\in U(x_k)}{\alpha }_i{\beta }_k⟨x_k,x_i⟩\text{\hspace{1em}(6)}$$

其中$U(x_k)$表示$x_k$的局部区域。将等式6代入等式3的第二项，我们可以将其重写为下面样子：

$$\sum _{x_i\in \Omega }{\alpha }_t⟨x_0,x_t⟩\approx \sum _{x_i\in A}{\alpha }_i\cdot 1⟨x_0,x_i⟩+\sum _{x_k\in \mathring{A}}\sum _{x_i\in U(x_k)}{\alpha }_i{\beta }_k⟨x_k,x_i⟩$$

$$=\sum _{x_k\in A}\sum _{x_i\in U(x_k)}{\alpha }_i{\beta }_k⟨x_k,x_i⟩=\sum _{x_k\in A}{\beta }_k\sum _{x_i\in U(x_k)}{w}^q{w}^k{w}^v{x}_i⟨x_k,x_i⟩\text{\hspace{1em}(7)}$$

注意$⟨x_k,x_i⟩$是$x_k$和$x_i$的内积，计算二者的相似度。${\sum }_{x_i\in U(x_k)}{\alpha }_i{\beta }_k⟨x_k,x_i⟩$是$x_k$在该相邻区域内$x_k$的注意力结果。因此，$x_0$处的全局注意力对数值可以通过对相邻区域内的像素点的注意力结果求和，近似得到，如等式7所示。我们可以设计一个近似的算子，通过逐点的上下文关系传递，预测全局注意力。所以作者提出了一个全局注意力近似方法，Pixel Shift Self-Attention (PSSA)，基于像素点 shift 和卷积来近似实现全局注意力。我们首先沿着给定的方向（即左、右、上）移动特征图$L$个像素，然后利用原始特征图和移动后的特征图逐点相乘，得到变换特征。实际上，shift-product 操作在相邻点之间构建起了上下文关系，通过层级堆叠，我们可以将上下文关系传递到全局区域。最后，我们对这些变换特征进行加权求和（可通过卷积算子实现），得到近似自注意力图。Shift、element-wise product 和加权求和的复杂度都是 $O(n)$，所以 PSSA 的复杂度是$O(n)$。值得注意，PSSA 将自注意力转化为了一个对变换特征做的标准卷积操作。通过层级堆叠，这个结构通过上下文关系传递，能够预测出全局自注意力对数。

3.4 The Unification of Convolution and Self-Attention: X-volution

3.4.1 Convolution and Self-Attention are Complementary

卷积使用了位置的归纳偏置和各向同性的特性，使之具备平移同变性。但是它天然的局部特性无法构建起长期关系，该长期关系对图灵完备的原子操作来说是必要的。和卷积相反，自注意力丢弃了归纳偏置的特性（低偏置），无需模型假设就可探索数据集的自然模式。低偏置特点使自注意力可以自由地探索复杂的关系（比如长期依赖关系、各向异性的语义、CNN 的强局部关系等），但该机制需要在超大规模数据集上预训练。此外，自注意力很难优化，需要更长的训练周期和复杂的tricks。基于此发现，一些学者认为卷积应该被引入到自注意力机制中去，提升鲁棒性和表现。根据优化特点（well-condition/ill-condition）、注意力范围（局部/长期）及内容依赖性（content-dependent/independent）等，他们采用了不同的模型假设让卷积和自注意力可以互补。

3.4.2 The Multi-branch Topology for the Unification

有多篇工作尝试将卷积和自注意力结合起来，但是它们的拓扑组合都比较粗糙（层级堆叠、concat），无法得到一个原子操作（在一个模块中使用卷积和自注意力），结构都很不规则。例如 AANet 直接将卷积层和自注意力层的输出 concat 起来，组合它们的结果。实验表明，单个卷积或自注意力会造成性能退化，而当它们一起出现时性能则会显著提升。

尽管卷积和自注意力的计算模式很不一样，本文作者研究了卷积和自注意力的数学形式（等式1和2），在等式7中给出了近似形式，在特定的点积 map 上等价于卷积操作，发现全局 element-wise 交互（即点积）可以通过局部 element-wise 交互的传递来近似。因此，这两个算子可以在一个统一的计算模式（即卷积操作）内完成。从另一个角度看，卷积操作可以看作为自注意力的空间不变偏置。于是作者将这些算子结合得到一个多分支拓扑结构，如图1所示，兼具卷积和自注意力的优势。该多分支结构有两个主要分支。左边的分支由级联 Shift Pixel Self-Attention 和BN组成，近似地实现全局自注意力操作。右边分支是卷积分支，由级联卷积和BN组成。

3.4.3 Conditionality Converting Multi-branch Scheme to the Atomic X-volution

该多分支模块将卷积和自注意力组合起来。但是，它只是一个粗粒度的算子组合，让网络非常负责且不规则。从硬件实现的角度看，该多分支结构需要更多的缓存来处理多个分支。然而，单个原子操作要更加高效，内存占用也更低。为了简洁点，作者省略了BN。实际上，BN可以看作为一个 $1\times 1$分组卷积（组数等于通道数），可以融合到卷积/自注意力层中。事实上，作者通过层级堆叠的方式来用 PSSA，堆叠结构的加权操作可以省略，因为层级堆叠隐含着相邻像素点加权的操作。多分支模块的训练表达式可以写作：

$$y_0=\underset{\text{PSSA Branch}}{\underbrace{\sum _{x_i\in A}{\alpha }_i⟨x_0,x_i⟩}}+\underset{\text{Conv Branch}}{\underbrace{\sum _{x_i\in A}{w}^c{x}_i+b^c}}=\sum _{x_i\in A}{w}^q{w}^k{w}^v{x}_i⟨x_0,x_i⟩+\sum _{x_i\in A}{w}^c{x}_i+b^c\text{\hspace{1em}(8)}$$

其中$w^c$表示卷积权重，$b^c$是相应的偏置。

$$y_0=\sum _{x_i\in A}(w^q{w}^k{w}^v⟨x_0,x_i⟩+w^c)x_i+b^c=\sum _{x_i\in A}(w^{\mathcal{A}}(x_0,x_i)+w^c)x_i+b^c\text{\hspace{1em}(9)}$$

其中$w^{\mathcal{A}}(x_0,x_i)\stackrel{def}{=}{w}^q{w}^k{w}^v⟨x_0,x_i⟩$代表 PSSA 分支的 content-dependent/动态系数。$w^c$表示卷积分支的content-independent/静态系数，一旦训练完成它们就会固定住。观察等式9，我们可以发现，只做一个简单的变换，多分支结构就可变为一个卷积形式。这一步在 CNN 中很常用，叫做结构重参数化。作者首先在卷积和自注意力融合中扩展。根据等式9，将由卷积和自注意力组成的多分支结构转化为一个动态的卷积算子，叫做 X-volution。它可以插入到任何主流模型中去，作为一个原子算子使用。

4. Experiments

6. Appendix

6.1 More Implementation Details

6.1.1 Network Architecture

表5展示了 ImageNet 分类任务上 X-volution模型（X-volution, stage5）的结构。与 ResNet-50 唯一的区别就是在 conv5_x 中用 X-volution 代替 $3\times 3$卷积。

6.1.2 Implementation of X-volution

图4：PSSA 算子的详细结构

X-volution 由两个分支组成，一个是卷积，另一个是自注意力。对于卷积分支而言，作者使用$3\times 3$卷积，以及一个额外的$5\times 5$膨胀卷积（stride=1, dilation=2）与之平行。对于自注意力分支而言，作者使用了 PSSA，它是全局自注意力的近似实现，图4中有展示。PSSA中，输入特征图首先朝八个方向移动$L$个像素，$L\in [1,3,5]$。在 Query 和 Key 变换后，我们计算原始特征图和移动后特征图的 element-wise 乘积，得到变换特征。然后对变换特征进行加权相加，得到近似的自注意力图。在最终输出前是一个BN层。