当前位置：首页 > ai >正文

Efficient Attention 理解

ai 2025/6/21 20:28:44

论文链接：arxiv.org/pdf/1812.01243

引言

点积注意力 [1, 22, 23] 是神经网络中一种普遍的机制，用于建模长距离依赖性，这是深度学习面临的一项关键挑战，而卷积和递归网络在这方面的表现有限。该机制通过对前一层中所有位置的特征进行加权求和，计算每个位置的响应。与卷积或递归层的有限感受野不同，点积注意力在一次传递中将感受野扩展到整个输入。使用点积注意力高效建模长距离依赖性使得卷积和递归能够专注于局部依赖建模，这是它们的专长。

点积注意力 Dot-Product Attention

非局部模块 [23] 是点积注意力在计算机视觉中的一种改编，已经在视频分类 [23] 和生成对抗图像建模 [28, 2] 中取得了最先进的性能，并在目标检测 [23]、实例分割 [23]、行人再识别 [14]、图像去雨 [13] 等任务中展示了显著的改进。然而，在大输入（例如长序列、高分辨率图像、大型视频）上进行全局依赖建模仍然是一个未解决的问题。点积注意力的内存和计算复杂度与输入大小呈平方关系，限制了它在大输入上的应用。例如，非局部模块在处理 64 通道 128 × 128 的特征图时，使用超过 1 GB 的 GPU 内存和超过 25 GMACC 的计算，或者在处理 64 通道 64 × 64 × 32 的三维特征体积时，使用超过 68 GB 的内存和超过 1.6 TMACC 的计算（例如用于深度估计或视频任务）。高内存和计算成本限制了点积注意力在模型低分辨率部分的应用 [23, 28, 2]，并禁止其在对分辨率敏感或资源消耗大的任务中的使用。因此，对大输入进行全局依赖建模的需求激励了资源高效的注意力机制的探索。

对非局部模块的调查揭示了一个有趣的发现。如图 1 所示，撇开归一化不谈，点积注意力涉及两次连续的矩阵乘法。

第一次乘法 $S=QK^T$ 计算像素之间的成对相似性，并形成逐像素的注意力图。

第二次乘法 D=SV 则通过逐像素的注意力图对值 V 进行聚合，从而产生输出。

由于矩阵乘法是结合律的，将顺序从 $(QK^T)V$ 改为 $Q(K^TV)$ 并不会影响效果，但会将复杂度从 $O(n^2)$ 改变为 $O(d_kd_v)$ ，其中 n 是输入大小，d_k和 d_v 分别是键和值的维度。这一改变消除了模块复杂度中的 $O(n^2)$ 项，使其在复杂度上变为线性。

此外，在实际情况下， $d_kd_v$ 显著小于 $n^2$ ，因此这个新项不会成为新的瓶颈。

由此，将乘法顺序改为 $Q(K^TV)$ 产生了一种显著更高效的机制，本文将其称为高效注意力。这种新机制在数学上等价于带有缩放归一化的点积注意力，并且与 softmax 归一化近似等价。实验实证验证了当等价性是近似时，并不会影响准确性。此外，实验显示其效率允许在网络中集成更多的注意力模块，并集成到高分辨率的网络部分，从而显著提高了准确性。此外，实验表明，高效注意力使得在由于资源限制而无法使用点积注意力的任务中广泛应用注意力成为可能。

另一个发现是，高效注意力为注意力机制提供了新的解释。假设键的维度为 $d_k$ ，输入大小为 n，可以将 $d_k\times n$ 的键矩阵解释为 $d_k$ 个模板注意力图，每个图对应输入的一个语义方面。然后，每个像素的查询是对每个 $d_k$ 模板注意力图的 $d_k$ 个系数。