摘要
arXiv:2502.18137v2 宣告类型: replace-cross
摘要:由于注意力机制的时间复杂性呈二次方增长,高效的注意力实现对于大型模型至关重要。幸运的是,注意力通常表现出稀疏性,即注意力图中的许多值接近零,允许省略相应的计算。许多研究利用稀疏模式来加速注意力。然而,大多数现有工作集中在通过利用注意力图中特定的稀疏模式来优化特定模型的注意力。一种既能保证各种模型提速又能保持端到端性能的通用稀疏注意力依然难以实现。在本文中,我们提出了SpargeAttn,这是一种适用于任何模型的通用稀疏和量化注意力方法。我们的方法使用两阶段在线过滤器:第一阶段,我们迅速且准确地预测注意力图,从而省略部分矩阵乘法。第二阶段,我们设计了一种在线 softmax 意识过滤器,不会增加额外开销,并进一步省略部分矩阵乘法。实验表明,我们的方法在不牺牲端到端指标的情况下,显著加速了包括语言、图像和视频生成等各种类型的模型。代码可在 https://github.com/thu-ml/SpargeAttn 获取。