Fast attention mechanisms: a tale of parallelism

📄 arXiv: 2509.09001v1 📥 PDF

作者: Jingwen Liu, Hantao Yu, Clayton Sanford, Alexandr Andoni, Daniel Hsu

分类: cs.LG

发布日期: 2025-09-10

💡 一句话要点

提出ANNA注意力机制，加速Transformer并保持大规模并行计算能力

🎯 匹配领域: 支柱一：机器人控制 (Robot Control)

关键词: 注意力机制 Transformer 近似最近邻搜索 大规模并行计算 模型加速

📋 核心要点

  1. Transformer的二次复杂度限制了其在大规模任务中的应用，成为一个亟待解决的问题。
  2. 论文提出近似最近邻注意力（ANNA）机制，旨在降低计算复杂度，同时保持Transformer的表达能力。
  3. 理论证明ANNA-Transformer在MPC算法模拟和关键推理任务上具有优越性能，并能模拟低秩Transformer。

📝 摘要（中文）

Transformer模型具有模拟大规模并行计算（MPC）算法的表征能力，但其二次时间复杂度严重限制了其可扩展性。本文提出了一种高效的注意力机制，称为近似最近邻注意力（ANNA），它具有亚二次时间复杂度。我们证明了ANNA-Transformer：（1）在匹配MPC算法能力方面，保留了先前为标准注意力机制建立的表达能力；（2）能够以接近最优的深度解决诸如Match2和k-hop等关键推理任务。利用MPC框架，我们进一步证明了常数深度的ANNA-Transformer可以模拟常数深度的低秩Transformer，从而为推理一类广泛的高效注意力近似方法提供了一种统一的方式。

🔬 方法详解

问题定义：Transformer中的标准注意力机制具有二次时间复杂度，这限制了其在处理长序列时的可扩展性。现有的高效注意力机制往往难以保证表达能力，或者缺乏统一的理论框架来分析其性能。因此，如何设计一种既高效又具有足够表达能力的注意力机制是一个关键问题。

核心思路：论文的核心思路是利用近似最近邻搜索（Approximate Nearest Neighbor Search, ANNS）来近似计算注意力权重。通过只关注与查询向量最相关的键向量，可以显著减少计算量，从而降低时间复杂度。这种方法旨在在计算效率和模型表达能力之间取得平衡。

技术框架：ANNA-Transformer的整体架构与标准Transformer类似，主要区别在于注意力机制的实现。在ANNA中，对于每个查询向量，首先使用ANNS算法找到其近似最近邻的键向量集合。然后，只在这个集合上计算注意力权重，并进行加权求和。整个过程可以并行化，从而进一步提高计算效率。

关键创新：ANNA的关键创新在于将近似最近邻搜索引入到注意力机制中。与标准注意力机制需要计算所有查询-键对之间的相似度不同，ANNA只计算与查询向量最相关的键向量的相似度，从而显著降低了计算复杂度。此外，论文还从理论上证明了ANNA-Transformer的表达能力，并证明其可以模拟低秩Transformer。

关键设计：ANNA的关键设计包括ANNS算法的选择、近邻数量的设置以及注意力权重的计算方式。ANNS算法的选择会影响计算效率和近似精度。近邻数量的设置需要在计算复杂度和模型表达能力之间进行权衡。注意力权重的计算方式可以采用标准的softmax函数，也可以采用其他更高效的近似方法。论文中具体使用的ANNS算法和参数设置未知。

📊 实验亮点

论文通过理论分析证明了ANNA-Transformer在MPC算法模拟和关键推理任务（如Match2和k-hop）上具有接近最优的性能。此外，论文还证明了常数深度的ANNA-Transformer可以模拟常数深度的低秩Transformer，从而为一类广泛的高效注意力近似方法提供了一种统一的理论框架。具体的实验数据未知。

🎯 应用场景

ANNA注意力机制可以应用于各种需要处理长序列的任务，例如自然语言处理中的机器翻译、文本摘要和问答系统，以及计算机视觉中的视频理解和图像生成。该方法降低了计算复杂度，使得Transformer模型能够处理更长的序列，从而提升模型性能和应用范围。此外，该研究为设计高效注意力机制提供了一种新的思路。

📄 摘要（原文）

Transformers have the representational capacity to simulate Massively Parallel Computation (MPC) algorithms, but they suffer from quadratic time complexity, which severely limits their scalability. We introduce an efficient attention mechanism called Approximate Nearest Neighbor Attention (ANNA) with sub-quadratic time complexity. We prove that ANNA-transformers (1) retain the expressive power previously established for standard attention in terms of matching the capabilities of MPC algorithms, and (2) can solve key reasoning tasks such as Match2 and $k$-hop with near-optimal depth. Using the MPC framework, we further prove that constant-depth ANNA-transformers can simulate constant-depth low-rank transformers, thereby providing a unified way to reason about a broad class of efficient attention approximations.