Subquadratic Algorithms and Hardness for Attention with Any Temperature

TL;DR

The paper addresses how attention in Transformers can be computed faster than the standard quadratic time barrier for a broad range of temperatures and head dimensions. It introduces a subquadratic algorithm for constant head dimension by removing irrelevant keys, approximating the exponential via low-degree polynomials on a bounded interval, and employing simplex range-searching in higher dimensions, achieving $\tilde{O}(n^{2-1/d})$ time (with polylog factors in $B$ and $1/\varepsilon$). It extends to low-rank cases and derives a subquadratic gradient computation, enabling faster end-to-end training in this regime. Complementarily, the work proves conditional lower bounds: under SETH, Attention requires $n^{2-o(1)}$ time even for $d = 2^{\Omega(\log^* n)}$ and polynomial $B$, and for large $d = \mathrm{poly}(n)$, the standard algorithm is near-optimal via generalized OV hypotheses. Together, these results delineate the precise trade-offs between head dimension, entry size, and temperature that govern the feasibility of truly subquadratic Attention.

Abstract

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.

Figures (2)

• Figure 1: Rounding based algorithm for $1$-dimensional Attention illustrated for $q_i = 1$. Each point is placed at $k_j$ and has value $v_{j}$. Points (e.g. $k_1$) such that $q_i k_j < q_i k_{\max} - \log(n/\varepsilon)$ are irrelevant and discarded (in this example $q_i k_{\max} - \log(n/\varepsilon) = 1$). Relevant points with similar $k_j$ (e.g. $\{k_{2}, k_{3}\}$ or $\{k_{6}, k_{7}, k_{8}\}$) are grouped together and assigned the same (rounded) key $\overline{k}$. The width of each region is $\log(1 + \varepsilon)$ (in this example $\log(1 + \varepsilon) = 3$). The algorithm outputs $\sum \overline{p}_{j} v_{j}$ where $\overline{p}_{j} \propto \exp(\overline{k}_{j})$.
• Figure 2: Polynomial method algorithm for $d$-dimensional Attention illustrated for $q_{i} = (2, 1)$. Relevant points are in the shaded blue region. Irrelevant points are in the white region. Weights are omitted for clarity.

Theorems & Definitions (53)

• Definition 1.0: Attention
• Definition 1.1: Approximate Attention Computation $\mathsf{AttC}(n, d, B, \varepsilon)$
• Theorem 1.2
• Theorem 1.3
• Theorem 1.4: Informal \ref{['thm:backwards-pass-alg']}
• Theorem 1.5: Informal \ref{['thm:super-const-d-lb']}
• Theorem 1.6: Informal \ref{['thm:poly-d-lb']}
• Theorem 2.2: williams2004ovc
• Theorem 2.3: chen2018maxip
• Theorem 3.1