无需训练即可加速各类模型的注意力计算，保持性能同时大幅提升速度。
SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference
- 两阶段在线过滤：快速预测注意力图并跳过部分矩阵乘法。
- 无额外开销的softmax感知过滤，进一步减少计算量。
- 适用于文本、图像、视频生成等多类模型，通用性强。
由于注意力机制具有二次时间复杂度，其高效实现对大模型至关重要。值得注意的是，注意力矩阵通常呈现稀疏性，即许多值接近零，因此可跳过对应计算。现有方法多针对特定模型，利用其注意力图的特定稀疏模式进行优化。然而，能兼顾速度提升与端到端性能的通用稀疏注意力方案仍缺失。本文提出SpargeAttn，一种适用于任意模型的通用稀疏量化注意力机制。方法采用两阶段在线过滤：第一阶段快速准确预测注意力图，跳过部分矩阵乘法；第二阶段设计无额外开销的softmax感知过滤，进一步跳过计算。实验表明，该方法在多种模型（包括语言、图像、视频生成）上显著加速推理，且不损失端到端指标。代码已公开于https://github.com/thu-ml/SpargeAttn。
原文摘要 · Abstract (English)
An efficient attention implementation is essential for large models due to its quadratic time complexity. Fortunately, attention commonly exhibits sparsity, i.e., many values in the attention map are near zero, allowing for the omission of corresponding computations. Many studies have utilized the sparse pattern to accelerate attention. However, most existing works focus on optimizing attention within specific models by exploiting certain sparse patterns of the attention map. A universal sparse attention that guarantees both the speedup and end-to-end performance of diverse models remains elusive. In this paper, we propose SpargeAttn, a universal sparse and quantized attention for any model. Our method uses a two-stage online filter: in the first stage, we rapidly and accurately predict the attention map, enabling the skip of some matrix multiplications in attention. In the second stage, we design an online softmax-aware filter that incurs no extra overhead and further skips some matrix multiplications. Experiments show that our method significantly accelerates diverse models, including language, image, and video generation, without sacrificing end-to-end metrics. The code is available at https://github.com/thu-ml/SpargeAttn.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。