[论文解读] X-volution: On the unification of convolution and self-attention

Xuanhong Chen, Hang Wang+1 more|arXiv (Cornell University)|Jun 4, 2021

Computability, Logic, AI Algorithms参考文献 34被引用 23

一句话总结

本文提出X-volution，一种新颖的神经网络模块，通过可学习的、卷积化的位移特征图变换，近似全局自注意力机制，统一了卷积与自注意力。该方法通过多分支架构实现局部与非局部特征学习的联合建模，可在推理阶段重参数化为单一标准卷积，实现SOTA性能，在ImageNet上实现+1.2%的top-1准确率提升，在COCO目标检测任务上实现+1.7 AP提升。

ABSTRACT

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 architectures, 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 a self-attention via the convolution operation on transformed features. Based on the approximated 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. Extensive experiments demonstrate that the proposed 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).

研究动机与目标

• 为解决由于卷积与自注意力机制计算模式异质，缺乏一种系统性方法将二者统一集成于单一神经网络模块中的问题。
• 开发一种低复杂度、可微分的全局自注意力近似方法，以实现在视觉任务中的高效计算。
• 设计一种多分支模块，在保持卷积优化稳定性的同时，通过自注意力机制捕捉长距离依赖关系。
• 实现结构化重参数化，使训练后的多分支模块可转换为单个高效卷积层，以支持高效推理。

提出的方法

• 提出像素位移自注意力（PSSA），一种O(n)的全局自注意力近似方法，通过在八个方向上位移特征图后进行逐元素运算，替代点积注意力机制。
• 构建包含标准3×3卷积分支、空洞5×5卷积分支以及PSSA分支的多分支X-volution模块，用于非局部特征交互。
• 通过可学习的位移操作与原始特征图和位移后特征图之间的逐元素乘法，生成动态注意力图。
• 在最终输出前应用批量归一化，以稳定训练过程并提升泛化能力。
• 利用结构化重参数化技术，在训练完成后将多分支结构合并为单个标准卷积层，从而实现高效推理。

实验结果

研究问题

• RQ1 能否仅通过在变换后的特征图上使用卷积操作，有效且高效地近似全局自注意力？
• RQ2 能否设计一种统一模块，将卷积与自注意力结合，实现稳定训练，并在性能上超越独立的CNN或Transformer？
• RQ3 多分支结构结合结构化重参数化，是否能同时保留卷积的优化优势与自注意力的表征能力？
• RQ4 所提出的X-volution模块能否作为即插即用的替代模块，无缝部署于现有网络而无需修改网络架构？

主要发现

• 与ResNet-50基线相比，X-volution在ImageNet-1K图像分类任务上实现了+1.2%的top-1准确率提升。
• 在COCO目标检测任务中，X-volution相较基线模型，将框AP提升+1.7点，掩码AP提升+1.5点。
• 在更长的训练周期（如3倍和6倍设置）下，X-volution的性能增益进一步提升，在3倍设置下COCO的AP提升达到+2.0点。
• PSSA近似在短周期训练下性能与标准卷积相当，但随着训练进程推进，其性能超越标准卷积，表明其具备更优的长距离建模能力。
• 重参数化后的X-volution模块在保持高性能的同时，可作为标准卷积部署，实现高效推理。