Source-linked AI summary
Fast Attention Requires Bounded Entries
Josh Alman, Zhao Song
TL;DR
Attention computation is a quadratic-time bottleneck when algorithms explicitly form its n × n matrix, motivating the search for faster implicit methods. The paper studies how input-entry magnitudes affect approximate attention and uses polynomial approximation and fine-grained reductions to prove upper and lower bounds. For d = O(log n), it gives an almost-linear-time algorithm below the √log n scale and, assuming SETH, rules out truly subquadratic time at that scale.
Problem
Explicitly computing the n × n attention matrix takes quadratic time, so the paper asks when approximate attention can be computed faster by using the matrix implicitly.
Source: §1 · pp. 2–3Method
The paper relates attention to KDE, using polynomial approximation techniques for its algorithm and a fine-grained ANN reduction for its lower bound.
Source: §5 · p. 17Results
For d = O(log n), the algorithm runs in n^1+o(1) time when B = o(√log n) with polynomial additive error; assuming SETH, no O(n^2−q)-time algorithm exists at d = C log n and B = Cβ√log n.
Source: §1.1 · p. 3, §4.1 · p. 12Takeaways & Limitations
The paper concludes that attention-computation complexity depends critically on input-entry magnitude and that KDE techniques can be adapted to this setting.
Source: §5 · p. 17The lower-bound result is conditional on SETH.
Source: §4.1 · p. 12
Abstract
from arXiv · showhide
In modern machine learning, inner product attention computation is a fundamental task for training large language models such as Transformer, GPT-1, BERT, GPT-2, GPT-3 and ChatGPT. Formally, in this problem, one is given as input three matrices $Q, K, V \in [-B,B]^{n \times d}$, and the goal is to construct the matrix $\mathrm{Att}(Q,K,V) := \mathrm{diag}(A {\bf 1}_n)^{-1} A V \in \mathbb{R}^{n \times d}$, where $A = \exp(QK^\top/d)$ is the `attention matrix', and $\exp$ is applied entry-wise. Straightforward methods for this problem explicitly compute the $n \times n$ attention matrix $A$, and hence require time $\Omega(n^2)$ even when $d = n^{o(1)}$ is small. In this paper, we investigate whether faster algorithms are possible by implicitly making use of the matrix $A$. We present two results, showing that there is a sharp transition at $B = \Theta(\sqrt{\log n})$. $\bullet$ If $d = O(\log n)$ and $B = o(\sqrt{\log n})$, there is an $n^{1+o(1)}$ time algorithm to approximate $\mathrm{Att}(Q,K,V)$ up to $1/\mathrm{poly}(n)$ additive error. $\bullet$ If $d = O(\log n)$ and $B = \Theta (\sqrt{\log n})$, assuming the Strong Exponential Time Hypothesis from fine-grained complexity theory, it is impossible to approximate $\mathrm{Att}(Q,K,V)$ up to $1/\mathrm{poly}(n)$ additive error in truly subquadratic time $n^{2 - \Omega(1)}$. This gives a theoretical explanation for the phenomenon observed in practice that attention computation is much more efficient when the input matrices have smaller entries.
1 Introduction
The paper asks when attention can be computed faster than the quadratic time required by explicitly forming its n × n attention matrix. It identifies a transition in input-entry magnitude and develops polynomial-approximation and ANN-reduction techniques to establish upper and lower bounds.
1 Introduction
- The task is motivated by attention computations being a major bottleneck for LLM operations on long sequences.The paper focuses on approximate attention because reasonable precision guarantees can suffice during inference.Source: §1 · pp. 2–3
- Explicitly forming the n × n attention matrix costs n^2+o(1) time, motivating algorithms that use the matrix implicitly.The input matrices have only n^1+o(1) entries when d = n^o(1).Source: §1 · p. 2
1.1 Our Results
- For d = O(log n), attention approximation with polynomial additive error takes n^1+o(1) time when B = o(√log n), while under SETH no truly subquadratic algorithm exists at B = Θ(√log n).The authors describe the transition between these regimes as tight.Source: §1.1 · p. 3
- The results offer a theoretical explanation for why attention computations can be faster when input entries are bounded or represented with few bits.The paper connects this phenomenon to the practical effectiveness of quantization and low-degree polynomial approximation.Source: §1.1 · pp. 3–4
1.2 Technique Overview
- The algorithmic upper bound adapts the polynomial method: approximate the exponential with a polynomial to obtain a low-rank approximation of the attention matrix.The approach draws on polynomial techniques used for Gaussian KDE.Source: §1.2 · p. 4
- The lower bound reduces Approximate Nearest Neighbor search to attention computation by encoding vectors in Q and K, so nearest neighbors correspond to large attention-matrix entries.The proof modifies input vectors to approximately balance rows and address the output's rescaling by D^-1.Source: §1.2 · p. 5
2 Preliminaries
The preliminaries specify the computational model and establish polynomial approximations to the exponential function. They also show how an additive approximation can be rescaled to give a relative-error approximation.
2 Preliminaries
- The algorithms use the real-RAM model, with real-number arithmetic treated as constant time.The paper uses the straightforward matrix-multiplication bound; faster theoretical multiplication could improve only low-order running-time terms under its parameter settings.Source: §2 · p. 5
2.1 Additive Error for Polynomial Approximation
- For B > 1 and 0 < ǫ < 0.1, the lemma specifies a polynomial P: R→R and reports |P(x) − exp(x)| < ǫ.Its rational coefficients can be computed efficiently in poly(g) time and have poly(g)-bit integer numerators and denominators.Source: §2.1 · p. 6
2.2 From Additive Error to Relative Error
- Rescaling an additive approximation on [0, 2B] yields a polynomial whose error on [-B, B] is less than ǫ · exp(x).The rescaling shifts the input by B and divides the polynomial output by exp(B).Source: §2.2 · p. 6
3 Attention Algorithm
The attention algorithm uses polynomial approximations of the exponential to approximate attention computations. Its core construction turns a low-degree polynomial applied entry-wise to a low-rank matrix into a computable low-rank factorization.
3 Attention Algorithm
- The algorithm uses polynomial approximations of the exponential function to approximately perform attention computations.The section develops low-rank approximation, bounds the relevant entries, and controls error propagation through the rescaling and output.Source: §3 · p. 6
3.1 Matrix Low-Rank Approximation
- The low-rank approximation framework represents an approximation to a matrix using two n × r factor matrices.The definition requires the factorization error to meet the specified accuracy parameter.Source: §3.1 · p. 7
3.2 From Low Degree Polynomials to Low Rank Matrices
- If M = XY^T and P is a degree-g polynomial, the entry-wise matrix P(M) can be factored as U1U2^T with rank parameter r.The factor matrices can be constructed from X and Y in O(nrg) time.Source: §3.2 · pp. 7–8
3.3 Matrix QK⊤Has Bounded Entries
Lemma 3.3 concerns bounding entries when Q and K have bounded entries.
3.3 Matrix QK⊤Has Bounded Entries
- Lemma 3.3 considers Q and K with entry magnitudes bounded by B, under the condition B ≥ 1.The supplied passage identifies the bounded-entry result but does not include its bound.Source: §3.3 · p. 8
3.4 Key Lemma
The key lemma gives an efficiently computable low-rank approximation to the attention matrix, using a polynomial in QK^T/d.
3.4 Key Lemma
- The attention matrix has a low-rank approximation that is easy to compute, despite potentially having full rank.The approximation is represented using matrices U1 and U2 computable in O(n · r) time.Source: §3.4 · pp. 8–9
- Lemma 3.4 provides an (ε, r)-approximation to A, with the approximation parameters bounded by quantities g and r.The supplied passages omit the bounds on g and r.Source: §3.4 · pp. 8–9
- The proof applies a degree-g polynomial to M = QK^T/d and uses Lemma 3.2 to compute the matrices defining the approximation.The polynomial approximation is described as P(M).Source: §3.4 · p. 9
3.5 From A to D
Lemma 3.5 defines diagonal matrices from the row sums of A and its approximation, and relates their diagonal entries.
3.5 From A to D
- The matrices D and eD are defined as the diagonal matrices of the row sums of A and eA, respectively.The lemma considers a positive-entry matrix A and an approximation parameter εA in (0, 0.1).Source: §3.5 · p. 9
- The passage defines D = diag(A1n) and eD = diag(eA1n), and states that for every i in [n], the ensuing relation has right-hand side ǫA · Di.Source: §3.5 · p. 9
3.6 From A and D to Attention Matrix
The proof bounds attention-output error by separating the effects of approximating the normalization and the attention matrix.
3.6 From A and D to Attention Matrix
- The argument assumes positive entries in A and eA and bounds the value matrix V entrywise by B.The supplied passages identify these assumptions but omit portions of the inequalities used to establish the error bound.Source: §3.6 · pp. 9–10
- The attention-output error is decomposed into a normalization term and a term comparing the approximate and exact attention matrices.The proof states that it bounds these two terms separately.Source: §3.6 · p. 10
3.7 Main Upper Bound
The section gives an algorithm for approximate attention computation and states its running time in terms of matrix multiplication and additional parameter-dependent work.
3.7 Main Upper Bound
- The theorem specifies positive integer parameters g and r for the algorithm's stated bound.The supplied theorem passage begins its parameter definitions but does not provide their full expressions.Source: §3.7 · p. 11
- The algorithm solves AAttC(n, d, B, ε) in O(Tmat(n, r, d) + nrg) time.The theorem identifies this as Algorithm 1's running time.Source: §3.7 · p. 11
- The proof bases the algorithm's running time and correctness on Lemma 3.4 and Lemmas 3.5–3.6, respectively.These lemmas establish the step costs and correctness used by the theorem.Source: §3.7 · p. 11
3.8 Proof of Theorem 1.4
For d = O(log n), B = o(√log n), and inverse-polynomial additive error, the proof obtains an n^1+o(1)-time approximation algorithm.
3.8 Proof of Theorem 1.4
- n^1+o(1) time suffices for AAttC when d = O(log n), B = o(√log n), and εa = 1/poly(n).This is the formal upper-bound statement for Theorem 1.4.Source: §3.8 · p. 11
- The proof selects parameters in Theorem 3.7 under the stated dimension, entry-magnitude, and error conditions.The supplied proof fragments identify ε = 1/poly(n) and B = o(√log n) among the substitutions.Source: §3.8 · p. 11
- The resulting bound is n^1+o(1) because d, r, and g are each bounded by n^o(1).This concludes the running-time calculation.Source: §3.8 · p. 12
4 Hardness
The hardness argument assumes SETH, invokes known hardness for approximate nearest-neighbor search, and reduces that problem to approximate attention computation. The reduction handles small distance thresholds by brute force and otherwise constructs an attention instance.
4.1 Fine-Grained Hypotheses
- SETH posits that for every ε > 0, there is a positive integer k ≥ 3 such that k-SAT on formulas with n variables cannot be solved in O(2^(1−ε)n) time, even by a randomized algorithm.Source: §4.1 · p. 12
4.2 Nearest Neighbor Search
- Under SETH, (1 + ε)-approximate Hamming nearest-neighbor search in dimension C log n requires Ω(n^2−q) time for every q > 0 and suitable ε and C.This is the hardness result used as the starting point for the attention reduction.Source: §4.2 · p. 12
4.3 Hardness Result
- Assuming SETH, approximate attention computation with logarithmic dimension and the theorem's specified parameter regime requires Ω(n^2−q) time.The formal statement gives d = C log n and additive error εa = n^−Cα; the supplied passage truncates the full entry-bound expression.Source: §4.3 · p. 13
4 Hardness
- The hardness section establishes the lower bound by combining the ANN hardness result with a reduction from Gap-ANN to approximate attention computation.The paper states that these two results imply its attention-computation hardness theorem.Source: §4.2 · p. 12, §4.3 · p. 13
4.3 Hardness Result
- If t < c log n, the reduction solves Gap-ANN by enumerating vectors within Hamming distance t of each input vector and checking a lookup table.This case takes O(n^2−Cγ) time in the reduction.Source: §4.3 · pp. 13–14
- For t ≥ c log n, the reduction constructs an attention instance from the Gap-ANN vectors.The reduction lemma converts an attention algorithm's running time T into a Gap-ANN algorithm running in O(T + n^2−Cγ) time.Source: §4.3 · pp. 13–14
5 Conclusion
The conclusion emphasizes that attention-computation speed depends on input-entry magnitude and relates the paper's techniques to kernel density estimation (KDE).
5 Conclusion
- Attention-computation speed depends critically on the magnitude of the input-matrix entries.The paper identifies B as the relevant magnitude parameter.Source: §5 · p. 17
- The paper uses similarities between attention computation and KDE to apply known KDE techniques to attention.It points to further transfer of KDE techniques as a direction for future work.Source: §5 · p. 17
- The authors suggest exploring what other techniques from the well-studied KDE problem can be applied to attention computation.Source: §5 · p. 17