Subquadratic Algorithms and Hardness for Attention with Any Temperature
This work provides a theoretical characterization of the computational limits of Attention for all temperatures, which is important for researchers and practitioners working on efficient Transformer architectures. It offers a new subquadratic algorithm for a specific regime and establishes hardness results for others.
The paper investigates the computational complexity of Attention in Transformer models, which typically scales quadratically with context length. For constant head dimension $d=O(1)$, they present the first subquadratic algorithm with a runtime of $\tilde{O}(n^{2 - 1/d} \cdot \mathrm{polylog}(B))$, even for large entry sizes $B$. They also provide hardness results, showing that substantial improvements are unlikely under $\mathsf{SETH}$ for $d = 2^{\Theta(\log^* n)}$ and that the standard algorithm is optimal for $d = \mathrm{poly}(n)$.
Despite the popularity of the Transformer architecture, the standard algorithm for computing Attention suffers from quadratic time complexity in context length $n$. Alman and Song [NeurIPS 2023] showed that when the head dimension $d = Θ(\log n)$, subquadratic Attention is possible if and only if the inputs have small entries bounded by $B = o(\sqrt{\log n})$ in absolute values, under the Strong Exponential Time Hypothesis ($\mathsf{SETH}$). Equivalently, subquadratic Attention is possible if and only if the softmax is applied with high temperature for $d=Θ(\log n)$. Running times of these algorithms depend exponentially on $B$ and thus they do not lead to even a polynomial-time algorithm outside the specific range of $B$. This naturally leads to the question: when can Attention be computed efficiently without strong assumptions on temperature? Are there fast attention algorithms that scale polylogarithmically with entry size $B$? In this work, we resolve this question and characterize when fast Attention for arbitrary temperatures is possible. First, for all constant $d = O(1)$, we give the first subquadratic $\tilde{O}(n^{2 - 1/d} \cdot \mathrm{polylog}(B))$ time algorithm for Attention with large $B$. Our result holds even for matrices with large head dimension if they have low rank. In this regime, we also give a similar running time for Attention gradient computation, and therefore for the full LLM training process. Furthermore, we show that any substantial improvement on our algorithm is unlikely. In particular, we show that even when $d = 2^{Θ(\log^* n)}$, Attention requires $n^{2 - o(1)}$ time under $\mathsf{SETH}$. Finally, in the regime where $d = \mathrm{poly}(n)$, we show that the standard algorithm is optimal under popular fine-grained complexity assumptions.