Vision Attention之X-volution详细解读，一个不错的想法

X-volution: On the Unification of Convolution and Self-attention

本文原始文档：https://www.yuque.com/lart/papers/zpdmek

今天，上交&华为海思提出了新的“卷王”Xvolution：它对卷积与自注意力进行了统一集成，同时利用了卷积的局部特征提取能力与自注意力的全局建模能力。更重要的是，它通过结构重参数化思想将训练与推理进行了解耦：在训练阶段采用多分支结构进行训练，在推理阶段等价转换为单一动态卷积形式。
​
——“新卷王”X-volution | 将卷积与自注意力进行高效集成，上交与华为海思提出了 Xvolution：https://mp.weixin.qq.com/s/yE9nwWjAUjAkmDSpnVuzKA

从摘要读论文

Convolution and self-attention are acting as two fundamental building blocks in deep neural networks, where the former extracts local image features in a linear way while the latter non-locally encodes high-order contextual relationships.

这里很精炼的总结了卷积和自注意力计算的计算形式的差异。一个是线性提取局部特征，另一个是计算非局部的上下文关系。

Though essentially complementary to each other (i.e., first-/high-order), stat-of-the-art ar-chitectures (i.e., CNNs or transformers) lack a principled way to simultaneously apply both operations in a single computational module, due to their heterogeneous computing pattern and excessive burden of global dot-product for visual tasks.

由于两种操作的不同计算模式和自注意力在视觉任务中全局点积计算的巨大负担，所以现有的架构缺少一种在单一计算单元同时应用两种具有互补价值的操作有原则的方式。
难道作者要降低自注意力运算的负担么？

In this work, we theoretically derive a global self-attention approximation scheme, which approximates self-attention via the convolution operation on transformed features.

果然是简化了自注意力操作。

Based on the approximate scheme, we establish a multi-branch elementary module composed of both convolution and self-attention operation, capable of unifying both local and non-local feature interaction.
Importantly, once trained, this multi-branch module could be conditionally converted into a single standard convolution operation via structural re-parameterization, rendering a pure convolution styled operator named X-volution, ready to be plugged into any modern networks as an atomic operation.

自注意力操作与卷积相结合，通过近似策略以及重参数化策略进行转化后，整体形成了一个多分支的卷积结构——X-volution。

Extensive experiments demonstrate that the pro-posed X-volution, achieves highly competitive visual understanding improvements (+1.2% top-1 accuracy on ImageNet classification, +1.7 box AP and +1.5 mask AP on COCO detection and segmentation).

从摘要整体来看，这篇文章主要做的就是将自注意力操作通过特定的近似策略实现了与卷积运算的统一。
也就是说，如何保证这种近似足够有效果和有效率才是关键。

主要方法

右侧结构为推理阶段重参数化后的结构。而模型细节应该重点看图中左侧的结构，也就是训练阶段的模型。

左侧结构的核心是一个双分支结构：

• 右侧支路包含级联的卷积层和 BN 层（这样的结构可以在推理阶段被融合为一个单一的卷积层），这一结构可以保证模型利用卷积来获得局部信息的增强。
• 左侧支路构建在自注意力操作之上。是一个基于（空间）像素偏移操作（）与卷积构成的单元和一个 BN 层，这用于近似全局自注意力操作。

关于 Conv 和 BN 层的推理时融合是一个常用的加速推理的策略，因为 BN 层的参数在推理时是固定的，即失去了训练时需要跨 batch 计算均值标准差的需求，所以可以被简单等价于 batch 共享的分组数输入通道数（也等于输出通道数）的$1\times 1$卷积，从而便于与前一层合并。当然，实际代码可以直接通过将两个结构的计算过程合并后的公式进行构造。

为了更好的表述这个将注意力和卷积结合起来的过程，接下来将按照论文中的构思过程来描述。

卷积运算

对于输入 X 使用空间共享的权重矩阵 W（参数与输入独立）来聚合大小为$K\times K$的局部窗口${\Delta }_K$中的特征（这里公式中的$\delta$是相对计算的局部区域的中心的偏移量，有正有负，而在 W 中索引的形式里，通过加上了一个矩阵中心的最大距离，从而将其调整为了 W 索引形式），同时使用偏置 B 来进行调整。

这可以看作是一种一阶的线性加权操作。

自注意力运算

注意，这里讨论的是原始的全局自注意力机制，而非后续提出的各种简化版本。

自注意力通过对线性变换过的 X 计算全局的相似性关系，从而获得一个对于输入 X 的全局不共享的动态权重（参数与输入有关）。在文中所述的形式中，下式最后的等价权重可以看作是a dynamic (element-wise content-dependent) and spatially varying convolutional filter。

这可以看作是一种高阶的全局操作。

这里要注意的是，这个公式实际上有问题。

假定按照文中的设定，输入序列设为$X\in {\mathbb{R}}^{C\times L}$，序列长度$L=H\times W$，操作中包含有独立的空间共享的线性嵌入变换权重$W^Q,W^K,W^V\in {\mathbb{R}}^{D\times C}$（这里按照常规设定，三者形状一致）。那么这里最后的$W^{V}X$应该有形状$L\times D$才对，但是按照其实际形状为$D\times L$。所以合理的形式应该是：

$$Y\in {\mathbb{R}}^{L\times D}\triangleq \text{softmax}(\underset{L\times D}{(W^{Q}X{)}^{\top }}\underset{D\times L}{W^{K}X})\underset{L\times D}{(W^{V}X{)}^{\top }}$$

这里是无法等价于论文中公式最终的形式的，（Y 的形状与 X 并不一致），除非是一个概念意义上的等价，而非公式推导。

更合理的表示应该将该过程表示为$Y\in {\mathbb{R}}^{D\times L}\triangleq \underset{D\times L}{W^{V}X}\text{softmax}(\underset{L\times D}{(W^{K}X{)}^{\top }}\underset{D\times L}{W^{Q}X})$。

但是该形式虽然在形状上符合，但是却并不能得到对应于论文公式最终的权重$\bar{W}(X)$，除非去掉 softmax。

如果去掉 softmax，则式子可以转化为：

$$\begin{gathered} Y\in {\mathbb{R}}^{D\times L}\triangleq \underset{D\times L}{W^{V}X}\underset{L\times D}{(W^{K}X{)}^{\top }}\underset{D\times L}{W^{Q}X}=\underset{D\times C}{(W^{V}X{X}^{\top }{W}^{K\top }{W}^Q)}\underset{C\times L}{X} \\ \bar{W}(X)=W^{V}X{X}^{\top }{W}^{K\top }{W}^Q \end{gathered}$$

看起来挺合适的 :>。
不过需要注意的是，作者并没有按照 self-attention 本身的形式来构建结构，并且后续的推导也不是基于 softmax 形式的。所以 logits 一直都是 logits。

全局自注意力机制的近似

为了降低自注意力机制在计算动态权重过程中的$O(n^2)$的复杂度，作者们将目光放到了自注意力操作中的动态权重的计算上。即通过现成的（off-the-shelf）操作符例如卷积或者元素乘积等将其计算过程进行简化。

在最终的形式上，作者实际上使用空间元素偏移和元素点乘的策略将自注意力操作表示成了卷积的形式。

需要注意的是，这里的运算和注意力机制中的运算并不在数学推导上等价，而仅仅是概念上的等价。

不过，从另一个角度来说，为什么我们必须延续 self-attention 的形式来构造运算呢？self-attention 的有效点应该是其运算的逻辑，而非形式吧？我又为什么始终会期待其从 self-attention 上开始推导呢？

但是我似乎意识到了我为什么会产生这样的纠结：作者旨在近似自注意力操作，却将推导过程建立在一个仅仅是概念等价的构造之上，这就有些立不住脚了。

这不同于之前的类似于 lambda network https://www.yuque.com/lart/papers/lhnua2那样，人家直接就说了是一种新的形式，本就不再遵循 self-attention 的固有运算，所以后续的改造也无可非议，是顺理成章地。但是本文在前文的描述中却缺少了这样的铺垫与过渡。让阅读增加了不必要的期待，换句话说，如果延续 self-attention 自身的运算结构，那还是很难被简化的到本文的框架上的。

所以为了理解本文的推导，我们需要抛开 self-attention 本身的形式，而要以其运算的核心概念来作为基础，即特定位置的结果是通过利用其与全局上下文的相似关系来整合全局信息得到的。

self-attention 是一种基于相似性的聚合运算。

convolution 是一种基于预设权重的聚合运算。这里只考虑普通卷积，不考虑动态权重。

所以在本文新构造的形式中，包含了这两点——相似性计算（特征矢量的点积）、聚合操作（基于点积得到的相似性加权聚合全局特征）。

接下来介绍下作者的推导思路。

首先按照文中的表述，输入位置$x_0$对应的 attention logits 表示如下：

从这个式子可以知道，这里的 attention logits 并不是指代原始 self-attention 中的 attention logits map（即Q^TK），而是指代Q^TKV^T。

由点乘运算可知，这里的${\alpha }_t\in {\mathbb{R}}^{C\times 1}$是一个向量，其中用到输入 X 中的向量元素$x_t\in {\mathbb{R}}^{C\times 1}$以及对应于输入的三个权重矩阵中的元素，这里的运算都是 element-wise 乘法。

而这里使用相似性计算基本可以对标 QK 的计算。而在 QK 的计算中，每个输入特征对应的权重 w 都是全局共享通道不共享的，而这里$w^q$和$w^k$被取了出来，独立于相似性计算的过程，并整合到了对于“V”的变换中。相当于直接基于了原始的输入计算了相似性。

原本对于全局的聚合，作者将其拆分成了两部分，一部分是如同卷积那样，将其局限于特定窗口内的局部区域，另一部分则是窗口之外的其他非局部区域。

考虑到图像自身具有较强的 markovian property[Efficient Inference in Fully Connected CRFs with Gaussian Edge Potentials]，所以特定位置可以近似地由其局部邻域线性表示：$x_0\approx {\sum }_{x_k\in \mathring{A}}{\beta }_k{x}_k$，这里${\beta }_k$是线性权重。于是前面式子的非局部项可以被改造如下：

式子中的当前位置的特征$x_0$被局部近似并替换成了特定窗口内的特征组合。

不失一般性，这里将非局部区域运算中被排除在外的“局部区域”的系数$\alpha$设为 0，从而将空缺填补上去：

考虑到图像的 markovian property，作者假设对于参考点$x_0$（以及附近的位置$x_k$）和远离参考点的位置$x_i$的交互是非常弱的。所以进一步简化得到：

即将$x_i$进一步约束到$x_0$的临近点$x_k$的局部区域内。

这里有一种类似于二阶近似的感觉，通过特定位置局部邻域和其自身邻域的相似性（即邻域像素的注意力结果），从而近似表征该位置与更大范围邻域的相似性。即保留了最相关的两层关系。

通过整体替换原始式子，可以得到最终的 attention logits 形式：

这里的结合过程似乎有点让人迷惑。$A$头顶的圈是怎么没的？

$\mathring{A}$表示的区域中排除了中心点$x_0$。

$$\sum _{x_i\in A}{\alpha }_i<x_0,x_i>=\sum _{x_i\in U(x_k)}{\alpha }_i{\beta }_k<x_k,x_i>,k=0,{\beta }_0=1.$$

这局部项正好填补了被排除的中心点。所以$A$头顶的圈被去掉了。

于是这里公式对应的具体含义为：局部项中“局部区域中心点与其邻居的关系的加权和”加上非局部项中“局部区域中心点之外的其他点和其自身邻居的关系的加权和结果的加权和”，得到了“以$x_0$为中心的局部区域内所有点与各自邻居的关系加权和结果的加权和”。

两项的结合好像显得前面的“局部”与“非局部”的拆分多此一举？直接按照全局计算前面的推导看起来也是满足的？

实际上还是有必要的，因为这里的近似仅仅是在非局部项中。前式可以直接合并仅仅是因为局部近似的窗口大小和前一项局部窗口大小一样。如果不一样了，那就不能这样融合了。

最后通过这样一个局部邻域的聚合过程，实现了对于全局关系的近似，即Pixel Shift Self-Attention (PSSA)。

接下来就是本文的重头戏，即如何将前述的推导得到的形式使用空间偏移和卷积操作实现出来。

Pixel Shift Self-Attention (PSSA)的实现方式

上图展示了附录中 PSSA 的具体过程。文章正文中的描述对此非常容易让人产生误解（这个糟糕的写作风格:<）。

首先对特征图进行空间偏移，沿着给定方向的八个方向各偏移 L 个像素位置，之后使用 Key 变换整合特征。

然后对经过 Query 变换后的原始特征和 Key 变换之后的特征使用 element-wise 乘法获得转换后的特征。

实际上，shift-product 操作构建邻域点之间的上下文关系，并且通过分层堆叠的形式，上下文信息被逐渐传播到全局范围。

最后通过使用卷积加权求和聚集这些邻域信息来获得一个近似的 self-attention map。

由于空间偏移、element-wise product 和加权求和的复杂度都是$O(n)$，所以整体构成的 PSSA 也是$O(n)$复杂度的运算单元。

总而言之，通过这样的策略，self-attention 被转换为一个作用域变换特征上的标准卷积操作。通过多层堆叠，该结构可以通过上下文关系传播来实现对于全局自注意力的有效近似。

等等~说到这里，我们需要考虑，为什么前面的式子就可以通过这样的偏移加变换的形式来等效呢？

在最终的公式中，attention 过程可以拆分成级联的相似性计算和局部集成拆两步操作，即先在每个小的局部区域计算相似性，之后再集成所有的这些小区域的结果来得到最终的结果。

• 每个小的局部区域的相似性计算可以直接通过空间偏移和乘法来实现。特定方向的偏移和乘法实现了特定方向上距离等于偏移量的两点之间的相似性计算。通过使用（共享？）权重$w^k\&w^q$（key 变换和 query 变换），不同方向的特征图会被单独处理。从而可以得到每个位置和相邻局部区域内各点之间的相似性。
  • 这里的计算过程的加权明显是在计算相关性（乘法）之前的，和前面公式中的设计并不一致（真不知道作者怎么回事）。
• 通过 query 变换后的原始特征与 key 变换后的偏移特征计算得到相似性之后，接下来就该第一次聚合了（即后续推导中和忽略的系数${\beta }_k$），正常来说，这该将不同的偏移的结果进行合并，从而得到局部聚合后的结果。之后该第二次加权聚合了，在特征图上应用卷积来整合局部特征（我觉得这里的参数可以看作是$w^v$）。
• 两次加权聚合进而实现了对于全局信息的近似聚合。获得的值可以看作是对于局部区域内的位置$x_i$**的一个动态加权后的结果。**作者提到“实际上，我们通常通过分层堆叠使用 PSSA，并且可以省略堆叠结构中的加权操作，因为分层堆叠意味着加权邻居像素的操作。”所以应该是第一次聚合被省略了，直接一次聚合搞定。

对卷积和自注意力的统一

关于卷积和自注意力是互补的，作者总结的内容非常精辟，这里摘录下来：

The convolution employs the inductive bias of locality and isotropy endowing it the capability of translation equivariance [1]. However, the local inherent instinct makes convolution failed to establish the long-term relationship which is necessary to formulate a Turing complete atomic operator [4, 23].

In contrast to convolution, self-attention discards mentioned inductive bias, so-called low-bias, and strives to discover natural patterns from a data-set without explicit model assumption. The low-bias principle gives self-attention the freedom to explore complex relationships (e.g., long-term dependency, anisotropic semantics, strong local correlations in CNNs [7], e.t.c.), resulting in the scheme usually requires pre-training on extra oversize data-sets (e.g., JFT-300M, ImageNet21K). In addition, self-attention is difficult to optimize, requiring a longer training cycle and complex tricks [31, 20, 38, 2].

Witnessing this, several works [1, 36] propose that convolution should be introduced into the self-attention mechanism to improve its robustness and performance.

In short, different model assumptions are adopted to make convolution and self-attention complement each other in terms of optimization characteristics (i.e., well-condition/ill-condition), attention scope (i.e., local/long-term), and content dependence (content-dependent/independent) e.t.c…

这里总结道：不同的模型假设被用来使卷积和自注意互相补充，例如在优化特性方面（即 well-condition/ill-condition），注意范围方面（即局部还是长距离）以及内容依赖性方面 (内容依赖的还是内容独立的）等。

为了实现对于二者的统一，现有方法已经有一些尝试，例如 Cvt: Introducing convolutions to vision transformers 和 AANet: Attention augmented convolutional networks。然而他们的策略过于粗糙，例如直接分层堆叠或者是拼接，这并不便于获得一个单独的原子操作（即在相同模块中应用卷积和自注意操作），并且使得结构不够标准化。虽然两种操作之间差异明显的计算模式带来了许多阻碍，但是通过前面的近似处理，二者可以实现基于统一的模式来实现，即卷积操作。从另一个视角来看，卷积反过来也可以看作是自注意力操作的空间归纳偏置的形式。由此，作者们将利用二者构建了图 1 所示的多分支结构，这可以同时受益于二者。

于是训练阶段，该结构的公式表达如下：

这里的运算中的相关计算，由于使用空间偏移来模拟领域与领域的关系建模，所以上式展现的局部运算形式中只体现了偏移后的运算。

通过已经计算得到的相关性和卷积操作，自注意力运算被转化成了一个局部的动态卷积的过程。

这样的结构通过重参数化在推理时可以合并为一个独立的动态卷积算子，命名为X-volution。

实验结果

感想

论文的方法包装的有点过度，想法确实挺有趣，将全局的自注意力计算近似为一种局部范围内各点与其自身局部范围的关系的整体聚合。

但是文章的方法介绍却并不清晰。

最为让我不舒服的一点就是，公式与方法的对应并不紧密，看了半天也对不上方法的介绍。不知道这些公式在写给谁看。

极其不痛快的一次阅读经历。

链接

• 论文：https://arxiv.org/pdf/2106.02253.pdf
• Vision Transformer: https://www.yuque.com/lart/architecture/nmxfgf
• Pytorch中Spatial-Shift-Operation的5种实现策略：https://www.yuque.com/lart/ugkv9f/nnor5p