Flash Window Attention：加速Swin Transformer的注意力计算
计算机视觉与模式识别
2025-01-15 v2
摘要
为了解决图像像素的高分辨率问题，Swin Transformer引入了窗口注意力机制。该机制将图像划分为不重叠的窗口，并将注意力计算限制在每个窗口内，显著提高了计算效率。为了进一步优化这一过程，可以考虑用flash attention替代标准注意力，后者在语言模型中已被证明更高效。然而，直接替换是无效的。Flash attention专为长序列设计，而窗口注意力处理的是短序列，但必须并行处理大量窗口。在本报告中，我们提出了一种针对窗口注意力优化的解决方案，称为Flash Window Attention。Flash Window Attention将注意力计算效率提高了高达300%，并将端到端运行时效率提高了高达30%。我们的代码已在线公开。
引用
@article{arxiv.2501.06480,
title = {Flash Window Attention: speedup the attention computation for Swin Transformer},
author = {Zhendong Zhang},
journal= {arXiv preprint arXiv:2501.06480},
year = {2025}
}