[1812.01243] Efficient Attention: Attention with Linear Complexities

“Efficient Attention: Attention with Linear Complexities” 提出了一种高效注意力机制，有效解决了点积注意力机制在处理大规模输入时内存和计算成本过高的问题，在多个任务中展现出卓越性能，推动了注意力机制在深度学习领域更广泛、高效的应用。

  1. 研究背景与动机

    • 点积注意力在神经网络中广泛用于长距离依赖建模，但内存和计算成本随输入规模呈二次增长，限制了其在高分辨率输入或资源敏感任务中的应用。例如，在处理 64 通道的 128×128 特征图或 64 通道的 64×64×32 3D 特征体时，非局部模块（基于点积注意力）会消耗大量 GPU 内存和计算资源。
    • 为满足大规模输入下全局依赖建模需求，需探索资源高效的注意力机制。

  2. 相关工作

    • 点积注意力：在机器翻译、自然语言处理和计算机视觉等领域取得成功，但资源消耗问题突出。本文主要基于非局部模块中的点积注意力进行改进，并在相同性能和资源约束下对比高效注意力模块与非局部模块。
    • 缩放注意力：如挤压与激励（SE）模块、CBAM 和 GCNet 等，专注于突出重要特征和抑制无信息特征，与点积注意力目标不同，本文未对两者进行比较。
    • 高效非局部操作：LatentGNN 通过低秩矩阵近似非局部模块中的亲和矩阵；CGNL 虽也处理全局依赖，但将输入展平为向量后信息交互受限。实验表明高效注意力机制性能更优。

  3. 方法介绍

    • 点积注意力回顾：通过三个线性层将输入特征向量转换为查询、键和值，计算所有位置对的相似度并加权求和得到输出。其内存复杂度为，计算复杂度为，在大输入时资源需求过高。
    • 高效注意力机制：同样经三个线性层转换，但将键视为个单通道特征图，聚合值特征形成全局上下文向量，再结合查询得到输出。其内存复杂度为，计算复杂度为，显著降低了资源消耗。
    • 等价性证明：在缩放归一化下，通过公式推导证明高效注意力与点积注意力等价；在 softmax 归一化下，两者近似等价且实验验证了这种近似不影响精度。
    • 机制解释：点积注意力为每个位置生成注意力图来聚合值；高效注意力将键解释为全局注意力图，对应输入的语义方面，先聚合值形成全局上下文向量，再由各位置结合自身查询进行进一步处理。
    • 效率优势分析：高效注意力避免计算所有位置对的相似度，消除了复杂度中的项，在不同输入规模下，相比非局部模块（基于点积注意力）可大幅节省内存和计算资源，如在处理 256×256 特征图时，内存和计算节省倍数显著。

  4. 实验部分

    • MS - COCO 任务集实验

    • 与非局部模块比较：在相同位置插入模块时性能相近，但高效注意力资源消耗少几个数量级；在相同资源上限下，高效注意力性能更好，如 res3 - 4 + fpn3 - 5 配置中表现突出。

    • 与竞争方法比较：在所有设置下，高效注意力模型（EA 模型）性能和性能提升最高且资源消耗最少。

    • 消融研究： 实验表明缩放和 softmax 归一化方法对高效注意力效果影响不大，后续实验采用 softmax 归一化；降低键的维度对精度影响极小，验证了可通过降低键和查询维度节省资源。

    • 其他任务实验

    • 立体深度估计：在 Scene Flow 数据集上，基于 PSMNet 基线模型插入高效注意力模块后性能提升，达到新的 state - of - the - art，而插入非局部模块会因内存需求过大无法验证有效性。

    • 时间动作定位：在 THUMOS14 数据集上，在 ResNet - 50 骨干网络的 res3 和 res4 后添加高效注意力模块，显著提高了任务性能。

  5. 可视化结果：对高效注意力模块的全局注意力图可视化显示，不同组的全局注意力图具有不同语义焦点，如捕捉前景、物体核心和周边部分，支持了注意力图是模板注意力图线性组合的分析。

  6. 研究结论：高效注意力机制在内存和计算效率上比点积注意力有大幅提升，在对象检测、实例分割和立体深度估计等任务中验证了有效性并取得显著改进，在其他注意力应用领域也有潜在应用价值，未来计划推广到更多领域。