提出一种线性快速注意力机制,高效处理长序列视觉任务。
ELFATT: Efficient Linear Fast Attention for Vision Transformers
- 设计线性复杂度的ELFATT,避免传统注意力的二次计算瓶颈。
- 在高分辨率视觉任务中比标准注意力快4-7倍,性能不降。
- 适配边缘设备与扩散模型,实测加速1.6-3倍,兼容FlashAttention-2。
注意力机制是Transformer在各类机器学习任务中成功的关键。然而,基于softmax的原始注意力机制随序列长度呈二次复杂度,成为长序列任务(如视觉任务)应用的主要瓶颈。尽管已有多种高效线性注意力机制被提出,但通常需牺牲性能以换取高效率。此外,内存高效的FlashAttention-1-3仍存在二次计算复杂度,尚有优化空间。本文提出一种新型高效线性快速注意力(ELFATT),在实现低内存输入输出操作、线性计算复杂度的同时保持高性能。在高分辨率视觉任务中,ELFATT相比原始softmax注意力实现4-7倍加速,且性能无损。该方法兼容FlashAttention-2,在其加速下仍可提供2-3倍速度提升。即使在非视觉类长序列任务(long-range arena)中,ELFATT亦表现领先,较FlashAttention-2提速1.2-2.3倍。在边缘GPU上,无论5W至60W不同功耗模式,均实现1.6-2.0倍加速。此外,ELFATT可直接用于增强和加速扩散任务,无需重新训练。
原文摘要 · Abstract (English)
The attention mechanism is the key to the success of transformers in different machine learning tasks. However, the quadratic complexity with respect to the sequence length of the vanilla softmax-based attention mechanism becomes the major bottleneck for the application of long sequence tasks, such as vision tasks. Although various efficient linear attention mechanisms have been proposed, they need to sacrifice performance to achieve high efficiency. What's more, memory-efficient methods, such as FlashAttention-1-3, still have quadratic computation complexity which can be further improved. In this paper, we propose a novel efficient linear fast attention (ELFATT) mechanism to achieve low memory input/output operations, linear computational complexity, and high performance at the same time. ELFATT offers 4-7x speedups over the vanilla softmax-based attention mechanism in high-resolution vision tasks without losing performance. ELFATT is FlashAttention friendly. Using FlashAttention-2 acceleration, ELFATT still offers 2-3x speedups over the vanilla softmax-based attention mechanism on high-resolution vision tasks without losing performance. Even in some non-vision tasks of long-range arena, ELFATT still achieves leading performance and offers 1.2-2.3x speedups over FlashAttention-2. Even on edge GPUs, ELFATT still offers 1.6x to 2.0x speedups compared to state-of-the-art attention mechanisms in various power modes from 5W to 60W. Furthermore, ELFATT can be used to enhance and accelerate diffusion tasks directly without training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。