FLASHSWIN：利用记忆高效注意力解锁 Swin Vision Transformers 中的大窗口和密集token

高分辨率视觉骨干模型长期以来被迫牺牲局部token密度来提供更大的感受野。分层 Swin 变换器强加了这种妥协，因为标准窗口注意力实现了每个窗口的 $M^2\times M^2$ 分数矩阵，随着窗口或token网格的增长，导致 $O(M^4)$ 记忆 。此外，Swin 在注意力分数中添加了学习到的相对位置偏差元素，需要分数矩阵及其梯度的完全具体化。这使得 Swin 和 SwinV2 陷入小窗口 ($M=8,16$)、补丁大小为 $4\times4$ ($p=4$) 的粗token状态，从而限制了细粒度任务的性能。我们引入了 FLASHSWIN，它用 FlashAttention 实现替代了标准窗口注意力，该实现可以计算精确的 softmax 注意力，而无需具体化分数矩阵，从而将每个窗口的记忆从 $O(M^4)$ 减少到 $O(M^2)$。这可以实现更高的token密度和更大的感受野，而不会增加记忆开销。 训练记忆在窗口大小上是平坦的：在 $32\times32$ 窗口中，FLASHSWIN-T 仅需要 $12.4$\,GB，与 $8\times8$ 相同，而 SwinV2/V1-T 的 $70/90$\,GB 则没有变化。然而，将 FlashAttention 直接应用于 Swin 会产生一种权衡：绕过分数矩阵会排除 Swin 的附加相对位置偏差，放弃空间信息以换取记忆效率。 FLASHSWIN 将位置信息恢复为窗口本地可学习的 2D RoPE，使大窗口和密集的token网格既经济又准确。在匹配的规模上，FLASHSWIN-T 的性能优于 Swin 变体。通过密集的token和宽窗口 ($p=2,M=32$)，相同的 Tiny 模型达到了 $84.1\%$ ImageNet-1K、$44.1$ COCO box AP 和 $47.28$ ADE20K mIoU——在 $M=16$ 上分别获得了 $+1.3$、$+5.1$ 和 $+1.82$ 相对于 SwinV2-T 的增益。在固定的 $M=32$ 下，将 patch 大小减半会产生比 mIoU 更大的边界质量增益，大致为 $3\times$。