Transformers have achieved state-of-the-art performance across various tasks, but suffer from a notable quadratic complexity in sequence length due to the attention mechanism. In this talk, I will share our work on MonarchAttention – a novel approach to sub-quadratic attention approximation via Monarch matrices, an expressive class of structured matrices.
Based on the variational form of softmax, we describe an efficient optimization-based algorithm to compute an approximate projection of softmax attention onto the class of Monarch matrices with Θ(N √N d) computational complexity and Θ(N d) memory and IO complexity. We show with experiments that unlike previous approaches, MonarchAttention is both (1) transferable, yielding minimal performance loss with no additional training, even when replacing every attention layer of the transformer, and (2) hardware-efficient, utilizing the highest-throughput tensor core units on modern GPUs. With optimized kernels, MonarchAttention achieves substantial speed-ups in wall clock time over FlashAttention-2: 1.4× for shorter sequences (N = 256) and 8.2× for longer sequences (N = 16K).
We demonstrate the quality of MonarchAttention on diverse tasks and architectures in vision and language problems, showing that it flexibly and accurately approximates softmax attention in a variety of contexts. This work is with Can Yaras, Alec Xu, Pierre Abillama, and Changwoo Lee. Our code is available at https://github.com/cjyaras/monarch-attention.