Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers

The self-attention mechanism is the key to the success of transformers in recent Large Language Models (LLMs). However, the quadratic computational cost O(n2) in the input sequence length n is a notorious obstacle for further improvement and scalability in longer contexts. In this work, we leverage the convolution-like structure of attention matrices to develop an efficient approximation method for attention computation using convolution matrices. We propose a mathsfconv basis system, analogous to the rank basis, and show that any lower triangular matrix can always be decomposed as a sum of structured convolution matrices in this basis. We then design a fast algorithm to approximate the attention matrix via a sum of such k convolution matrices. This allows us to compute the attention it inference via Fast Fourier Transforms (FFT) in O(knd log n) time, where d is the hidden dimension, and thus achieve almost linear time n1+o(1) in the practical scenario where kd = no(1). Furthermore, the attention it training forward and it backward gradient can be computed in n1+o(1) as well. We provide theoretical guarantees on the run time and approximation error and conduct preliminary experiments to evaluate its effectiveness. We hope our new paradigm for accelerating attention computation in transformer models can help their application to longer contexts.