Support Basis: Fast Attention Beyond Bounded Entries

Maryam Aliakbarpour, Vladimir Braverman, Junze Yin, Haochen Zhang

TL;DR: In this paper, we introduce support-basis decomposition as a new framework for fast attention approximation in Transformers.

Abstract:

Large language models (LLMs) have demonstrated remarkable performance across a wide range of tasks. However, the quadratic complexity of softmax attention remains a central bottleneck that limits their scalability. Alman and Song (NeurIPS 2023a; NeurIPS 2024a) proposed sub-quadratic time algorithms for attention inference and training, respectively, but they rely on the restrictive **bounded-entry assumption**. We show that this assumption rarely holds in practice, which significantly limits their applicability to modern LLMs.

In this paper, we introduce **support-basis decomposition**, a new technique for accurate and efficient attention inference and training **without** the bounded-entry assumption. We empirically show that the entries of the query and key matrices exhibit sub-Gaussian behavior. Leveraging this widely observed property, we perform exact computation on sparse components and polynomial approximation on dense components. Without relying on restrictive assumptions, we theoretically show that our algorithm achieves sub-quadratic runtime while matching the approximation error of prior work, and we empirically validate its computational efficiency and downstream task performance. We further generalize our method to a multi-threshold setting that eliminates all distributional assumptions, providing the first theoretical justification for the empirical success of polynomial attention. Moreover, we show that softmax attention can be closely approximated by multiple polynomial attentions with significantly smaller $\ell_p$ error.

Code Dataset Promise: Yes

Code Dataset Url: https://github.com/yinj66/support_basis

Format Confirmation: I agree that I have read and followed the formatting instructions for the camera ready version.

Submission Number: 155