优化Swin Transformer的窗口注意力，速度提升300%
Flash Window Attention: speedup the attention computation for Swin Transformer
- 针对窗口注意力设计专用加速算法
- 计算效率最高提升300%，端到端速度提升30%
- 适合图像模型优化与高效部署场景
为应对高分辨率图像像素带来的计算压力，Swin Transformer引入了窗口注意力机制，将图像划分为不重叠的窗口，并限制注意力计算仅在窗口内进行，显著提升了计算效率。然而，直接将标准注意力替换为已在语言模型中证明高效的Flash Attention效果不佳。原因在于Flash Attention针对长序列设计，而窗口注意力处理的是短序列但需并行处理大量窗口。本文提出专为窗口注意力优化的Flash Window Attention方案，可将注意力计算效率提升高达300%，端到端运行效率提升最多达30%。代码已公开。
原文摘要 · Abstract (English)
To address the high resolution of image pixels, the Swin Transformer introduces window attention. This mechanism divides an image into non-overlapping windows and restricts attention computation to within each window, significantly enhancing computational efficiency. To further optimize this process, one might consider replacing standard attention with flash attention, which has proven to be more efficient in language models. However, a direct substitution is ineffective. Flash attention is designed for long sequences, whereas window attention deals with shorter sequences but must handle numerous of them in parallel. In this report, we present an optimized solution called Flash Window Attention, tailored specifically for window attention. Flash Window Attention improves attention computation efficiency by up to 300% and enhances end-to-end runtime efficiency by up to 30%. Our code is available online.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。