注意力的 I/O 复杂性：闪快注意力有多优化？
内容提要
本文综述多种高效注意力机制：FlashAttention通过分块减少HBM与SRAM读写，提升Transformer速度并支持长上下文，后续扩展至稀疏注意力，训练提速2至3.3倍；另有研究将复杂度降至O(N)，或提出HyperAttention、Fast Multipole Attention、Routing Transformer、Linformer等，在保持性能的同时降低时间与内存开销，加速长序列推理。
延伸解读
IO 感知：FlashAttention 的优化核心
FlashAttention 的关键在于从内存读写角度优化注意力计算。它通过分块技术减少 GPU 高带宽内存（HBM）与片上 SRAM 之间的数据传输，而非改变计算复杂度。这种 IO 感知设计使 Transformer 在长上下文任务中速度更快，并首次在 Path-X 挑战中成功应用。后续的 FlashAttention-2 进一步提升了效率，训练速度可提高 2 至 3.3 倍。
稀疏化与线性化：降低复杂度的不同路径
为应对标准注意力的 O(N^2) 复杂度，研究者探索了多种方向。FlashAttention 可扩展至块状稀疏注意力，在低稀疏度下随序列增长而提升性能。另有工作将复杂度降至 O(N)，并保持注意力矩阵的完整表示。Routing Transformer 通过在线 K 均值稀疏路由，将复杂度从 O(n^2 d) 降至 O(n^{1.5} d)。Linformer 则通过近似注意力矩阵，将时间和空间复杂度降为 O(n)。
近似注意力的性能权衡：以 HyperAttention 为例
HyperAttention 是一种近似注意力机制，旨在解决长上下文带来的计算挑战。它利用局部敏感哈希识别大条目，实现线性时间采样，并可模块化集成 FlashAttention 等快速实现。实证显示，在 ChatGLM2 上推理时间加快 50%，但困惑度从 5.6 增至 6.3；在 131k 上下文且因果屏蔽下，单注意力层获得 5 倍加速。这表明速度提升可能伴随质量损失，需根据场景权衡。
Q&A
FlashAttention 是如何提升 Transformer 速度的？
FlashAttention 使用平铺（tiling）技术减少 GPU 高带宽内存（HBM）和片上 SRAM 之间的内存读写次数，从而加速 Transformer，使其能处理更长上下文并提升模型质量。
FlashAttention 能处理多长的序列？训练速度提升多少？
FlashAttention 可处理长达 8k 和 16k 的序列，在不牺牲困惑度的情况下，训练速度分别提升 2.0 倍和 3.3 倍。
有哪些方法可以将注意力复杂度从 O(N^2) 降到 O(N)？
Linformer 通过近似自注意力矩阵将时间和空间复杂度从 O(n^2) 降至 O(n)；此外，有研究引入可分解注意力机制，在保持完整注意力矩阵表示的同时将计算和内存复杂度降至 O(N)。
HyperAttention 是什么？它有什么优势？
HyperAttention 是一种近似注意力机制，用于解决大语言模型长上下文带来的计算挑战。它通过局部敏感哈希（LSH）识别大条目，实现线性时间采样，可模块化集成 FlashAttention 等快速实现。在 ChatGLM2 上推理时间加快 50%，困惑度从 5.6 增至 6.3；在 131k 上下文长度下，单个注意力层实现 5 倍加速。
Routing Transformer 如何降低注意力复杂度？
Routing Transformer 结合了稀疏注意力和局部、时间稀疏注意力，并采用基于在线 K 均值的稀疏路由模块，将注意力总体复杂度从 O(n^2 d) 降至 O(n^{1.5} d)，在 Wikitext-103 和 ImageNet-64 上表现良好。
Fast Multipole Attention 解决了什么问题？
Fast Multipole Attention 针对自注意力二次复杂度限制长序列应用的问题，通过分层方法降低时间和内存复杂度，同时保持全局感受野。
