ELFATT: Efficient Linear Fast Attention for Vision Transformers

TL;DR

ELFATT introduces a novel linear attention mechanism that significantly speeds up vision transformers with minimal performance loss, suitable for high-resolution tasks and resource-constrained environments.

Contribution

The paper proposes ELFATT, a new linear attention method that reduces memory and computational complexity while maintaining high performance, compatible with FlashAttention-2 and applicable to various tasks.

Findings

4-7x speedup over vanilla attention in high-res vision tasks

2-3x speedup with FlashAttention-2 acceleration

Effective in non-vision long-range tasks and on edge GPUs

Abstract

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.