作者： Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
提交/修订日期： 2022年5月27日提交，2022年6月23日修订
主题/分类： Machine Learning (cs.LG)
摘要：
Transformer 模型在处理长序列时速度慢且内存消耗大，因为自注意力机制的时间和内存复杂度与序列长度呈平方关系。近似注意力方法试图通过牺牲模型质量来降低计算复杂度以解决此问题，但通常无法实现实际运行速度的提升。本文认为，一个缺失的原则是使注意力算法具备 IO 感知能力——即考虑 GPU 内存层级之间的读写操作。
我们提出了 FlashAttention，一种 IO 感知的精确注意力算法，它使用分块技术来减少 GPU 高带宽内存（HBM）与 GPU 片上 SRAM 之间的内存读写次数。我们分析了 FlashAttention 的 IO 复杂度，表明它比标准注意力需要更少的 HBM 访问次数，并且对于一定范围的 SRAM 大小是最优的。我们还将 FlashAttention 扩展到块稀疏注意力，从而得到一个比任何现有近似注意力方法都更快的近似注意力算法。
FlashAttention 训练 Transformer 的速度比现有基线更快：与 MLPerf 1.1 训练速度记录相比，在 BERT-large（序列长度 512）上实现了 15% 的端到端实际运行速度提升，在 GPT-2（序列长度 1K）上实现了 3 倍加速，在长距离竞技场（序列长度 1K-4K）上实现了 2.4 倍加速。FlashAttention 和块稀疏 FlashAttention 使得 Transformer 能够处理更长的上下文，从而产生更高质量的模型（GPT-2 困惑度提升 0.7，长文档分类提升 6.4 个百分点）以及全新的能力：首次使 Transformer 在 Path-X 挑战（序列长度 16K，准确率 61.4%）和 Path-256（序列长度 64K，准确率 63.1%）上实现了优于随机猜测的性能。