[论文解读] Efficient Attention: Attention with Linear Complexities

Zhuoran Shen, Mingyuan Zhang+3 more|arXiv (Cornell University)|Dec 4, 2018

Advanced Neural Network Applications参考文献 22被引用 88

一句话总结

提出一种高效的注意力机制，与线性内存与计算下的点积注意力在数学上等价，使高分辨率和资源密集型任务成为可能，并在 MS-COCO 和 Scene Flow 数据集上带来改进。

ABSTRACT

Dot-product attention has wide applications in computer vision and natural language processing. However, its memory and computational costs grow quadratically with the input size. Such growth prohibits its application on high-resolution inputs. To remedy this drawback, this paper proposes a novel efficient attention mechanism equivalent to dot-product attention but with substantially less memory and computational costs. Its resource efficiency allows more widespread and flexible integration of attention modules into a network, which leads to better accuracies. Empirical evaluations demonstrated the effectiveness of its advantages. Efficient attention modules brought significant performance boosts to object detectors and instance segmenters on MS-COCO 2017. Further, the resource efficiency democratizes attention to complex models, where high costs prohibit the use of dot-product attention. As an exemplar, a model with efficient attention achieved state-of-the-art accuracies for stereo depth estimation on the Scene Flow dataset. Code is available at https://github.com/cmsflash/efficient-attention.

研究动机与目标

• 由于标准点积注意力在大输入尺寸上的二次成本，需要进行全局依赖建模的动机。
• 引入一种在缩放归一化下等价于点积注意力、具有线性内存和计算的替代注意力公式。
• 展示高效注意力使得在网络中集成更多注意力模块和更高分辨率特征成为可能。
• 在资源受限的情况下，在目标检测、实例分割和立体深度估计上展示显著的性能提升。

提出的方法

• 通过重新排列矩阵乘法来计算全局上下文 Q 和 K'V，从而得到线性内存和计算的点积注意力重构。
• 定义高效注意力 E(Q,K,V) = rho_q(Q) (rho_k(K)^T V)，其中 rho_q 和 rho_k 作为归一化函数（缩放或 softmax）。
• 在缩放归一化下证明等价性：使用缩放归一化时 E(Q,K,V) = D(Q,K,V)。
• 提供针对视觉数据的实现细节，通过展平输入并在残差连接中应用该机制。
• 给出一种解释，将 K 视为 d_k 维的模板注意力图，形成全局上下文向量，由每个位置的 Q 系数组合。

实验结果

研究问题

• RQ1 在降低资源消耗的同时，高效注意力是否能够达到与标准点积注意力相同的表示能力？
• RQ2 归一化选择（缩放 vs softmax）如何影响等价性和准确性？
• RQ3 在高分辨率视觉模型中集成高效注意力时，在内存、计算和准确性方面的实际收益有哪些？

主要发现

• 在缩放归一化下，高效注意力与点积注意力相匹配，在 softmax 下也能近似，与经验准确性保持一致。
• 内存和计算量随输入规模线性增长，而传统点积注意力为二次增长，使高分辨率成为可能。
• 在 MS-COCO 上，高效注意力实现了更优的性能-成本权衡，并在内存受限的场景下比非局部模块更易部署。
• 在立体深度估计（Scene Flow）中，高效注意力带来显著提升，达到最先进水平，同时使用的内存远低于非局部变体。
• 消融研究显示归一化方法和 key 维数对性能影响有限，便于实现节省资源的配置。