[论文解读] On the Integration of Self-Attention and Convolution

Xuran Pan, Chunjiang Ge+5 more|arXiv (Cornell University)|Nov 29, 2021

Domain Adaptation and Few-Shot Learning参考文献 35被引用 9

一句话总结

该论文提出ACmix，一种新颖的神经网络模块，通过共享相同的1×1卷积操作作为核心计算，实现自注意力与卷积的无缝融合，从而实现高效的联合特征学习。与纯自注意力或卷积模型相比，该方法在图像识别和目标检测基准上实现了最先进性能，且计算开销极低。

ABSTRACT

Convolution and self-attention are two powerful techniques for representation learning, and they are usually considered as two peer approaches that are distinct from each other. In this paper, we show that there exists a strong underlying relation between them, in the sense that the bulk of computations of these two paradigms are in fact done with the same operation. Specifically, we first show that a traditional convolution with kernel size k x k can be decomposed into k^2 individual 1x1 convolutions, followed by shift and summation operations. Then, we interpret the projections of queries, keys, and values in self-attention module as multiple 1x1 convolutions, followed by the computation of attention weights and aggregation of the values. Therefore, the first stage of both two modules comprises the similar operation. More importantly, the first stage contributes a dominant computation complexity (square of the channel size) comparing to the second stage. This observation naturally leads to an elegant integration of these two seemingly distinct paradigms, i.e., a mixed model that enjoys the benefit of both self-Attention and Convolution (ACmix), while having minimum computational overhead compared to the pure convolution or self-attention counterpart. Extensive experiments show that our model achieves consistently improved results over competitive baselines on image recognition and downstream tasks. Code and pre-trained models will be released at https://github.com/LeapLabTHU/ACmix and https://gitee.com/mindspore/models.

研究动机与目标

• 揭示先前研究中被忽视的自注意力与卷积之间在基本计算层面的相似性。
• 设计一种混合神经模块，充分利用自注意力与卷积的优势，同时避免计算成本翻倍。
• 降低通常与结合注意力和卷积模块相关的推理与训练开销。
• 通过实证验证，所提出的ACmix模块在多个视觉基准上优于纯卷积或自注意力模型。

提出的方法

• 将标准的k×k卷积分解为k²个独立的1×1卷积，随后进行移位与求和操作。
• 将自注意力中的查询、键和值投影重新解释为多个1×1卷积，随后进行动态注意力权重计算与值聚合。
• 在ACmix中，将相同的1×1卷积投影共享于自注意力与卷积路径之间，消除冗余计算。
• 使用可学习的缩放参数α和β，动态组合自注意力与卷积路径的输出，实现特征的自适应路由。
• 通过分组卷积或可学习核实现移位操作，以在聚合阶段保持高效与灵活。
• 采用端到端优化训练模型，其中共享的1×1卷积与可学习的α、β参数联合更新。

实验结果

研究问题

• RQ1 当自注意力与卷积被分解为其原始操作时，是否存在基本的计算等价性？
• RQ2 共享的1×1卷积操作能否作为注意力与卷积模块的统一基础？
• RQ3 一种重用相同特征投影的混合模块，是否在性能与效率上优于独立的注意力与卷积模块？
• RQ4 自注意力与卷积在不同网络深度下的相对贡献如何演变？

主要发现

• ACmix在ImageNet-1K上的top-1准确率达到81.9%，优于相同设置下的Swin-T（81.5%）与Conv-Swin-T（81.4%）。
• 消融实验表明，使用可学习参数α与β组合路径的性能优于固定比例或单路径模型。
• 采用基于分组卷积的移位模块的模型在单张RTX 2080Ti GPU上达到419 FPS的推理速度，展现出极高的推理效率。
• 对学习到的α与β参数的可视化显示：在浅层中卷积占主导地位，而在深层中自注意力的影响逐渐增强。
• 两条路径的组合在所有评估基准与设置下，始终优于仅使用单一模块的模型。
• 与分别训练注意力与卷积模块相比，所提方法显著降低了计算开销，同时保持或提升了性能。