Fast Attention Requires Bounded Entries
Authors: Josh Alman
Zhao Song
Venue: NeurIPS 2023 poster (2023)
Keywords: fast attention computation
algorithm
hardness

Abstract: 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 \emph{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.

Fast Attention Requires Bounded Entries
Josh Alman'
Zhao Song
Abstract
In modern machine learning, inner product attention computation is 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 input three matrices Q, K, E [-B, B)nxd and the goal is to construct the matrix Att(Q, K, V) == diag( (A1,) -' e Raxd, where A = exp(QKT /d) is the 'attention matrix', and exp is applied entry-wise. Straightforward methods for this problem explicitly compute the n x n attention matrix A, and hence require time ?(n²) even when d = o1 small.
In this paper, we investigate whether faster algorithms are possible by implicitly making use of the matrix A. present two results, showing that there sharp transition at B = @(vlog n.)
= Olog n) and = ol)log,), there is an ++1 time algorithm to approximate Att(Q, K, V) up 1/poly(n) additive error
d O(log n) and Hypothesis from fine one-grained complexity theory, it is impossible to approxi- mate Att(Q, K, V) up to 1/poly(n) additive error in truly subquadratic time 2-11(1)
((v (vllog assuming the Strong Exponential Time
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.
Introduction
Large language models (LLMs) such as Transformer [VSP+17], BERT [DCLT18), GPT-3 [BMR+20), 20], PaLM [CND+22], and OPT [2RG+22] can process natural language more effectively than smaller models or traditional algorithms. This means that they can understand and generate complex and nuanced language, which can be useful for variety of tasks such language translation, question answering, and sentiment analysis. LLMs can also be adapted to multiple purposes without needing be retained from scratch. Their power particularly exemplified by the recent success of ChatGPT, a chat software by OpenAl built on top of GPT-3 [Ope22].
The key technical backbone of LLMs is the attention matrix [VSP+ 17, RNS 18, DCLT18, RWC 19, BMR 20]. An attention matrix is a square matrix whose rows and columns correspond to words or "tokens", and whose entries correspond to the correlations between these tokens in natural text. The attention matrix is then used to calculate the importance of each input token in sequence when producing output. In attention mechanism, each input token is given weight score, which reflects its importance or relevance to the current output being generated. These scores are calculated based on comparison between the current output state and the input states, using similarity function.
joohhca columbia. edu. Columbia University. "zzonggadobeecomm com. Adobe Research.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

More formally, the attention matrix is defined as follows. Let Q E Rnxd be the matrix of query tokens, and K E Raxd be the matrix of key tokens. (We focus here on the case when d = n (1) so d n.) The attention matrix is an the input tokens in the sequence. Each entry in the matrix represents the attention weight or score between a particular input token (query token Q) and a particular output token (key token K). The diagonal entries of the matrix represent alf-aattenioon scores, which measure the importance of each token with respect to itself.
matrix A where the rows and columns correspond to
The major bottleneck to speeding up LLM operations (in the case of modeling long sequences with large n) is the time to perform attention matrix computations [VSP+17, RNS* 18, DCLT18, RWC+ 19, BMR+ 20, KKI.20J. These computations ask us to multiply the attention matrix A with another value token matrix V E I
We formally define Attention computation as follows. Throughout this paper, we write exp to denote the entry-wise exponential for matrices.
Definition 1.1 (Exact Attention Computation EAttC(n, d)). Given three matrices Q, K, V E 12axd output the n x d matrix Att(Q, K, V) defined by
Au(Q, K,) DV
where A Ruxn and diagonal matrix D Ruxn are defined
A = exp(QK" /d); and D (iag(AA,.)
The straightforward algorithm for this problem computes the matrix A and then performs the multiplications D-'AV AA, in time n²²a((1) Since A is an n matrix with n² entries, it is impossible to improve on this much while explicitly computing the matrix A. However, the input to the problem is not A, but rather the three matrices Q, K, V which each have only n' (+)) entries. An algorithm which only implicitly makes use of A, without explicitly computing all its entries, could hope to run in almost linear time!
In this paper, we investigate the possibility of accelerating attention computations in this way. The two main questions we address are:
Q1. When can we perform attention computations in almost linear time n -
Q2. When can we prove that subquadratic-time algorithms for attention computations are impossible?
In most LLMs, suffices approximately perform attention computations throughout the infer- ence process as long as there are reasonable precision guarantees [CGRS19, KKL20, WLK+ DKOD20, KVPF20, 21, CDL + 22, LWD + 23, ZSZ Z... We therefore focus here on approxi- mate attention computation, which can potentially be performed even faster than exact computation. Mathematicallyy, we define the approximate version of EAttC as follows.
Definition 1.2 (Approximate Attention Computation AAttC(n, d, B, Let > 0 and B 0 be parameters. Given three matrices Q, K, V E Raxd, with the guarantees that ||Q] < B, || loo < B, and [|V] S B, output a matrix T E gaxd which is approximately equal to DVA meaning,
IT D D-'AVIlle S Ca. Here, for matrix M E Ruxn we write || 11 = maxij |M,i
Again, the straightforward algorithm for this problem runs in time O(n2d) S + but the input size is only ((d)) < to()), Our goal is to investigate when faster algorithms are possible in terms of the parameters d, B, and 6a-
1.1 Our Results
We focus on the natural setting where d OOog n)) (the setting where we model long sequences) and ca = 1/ poly(n) (low enough error so that attention computations over entire network can be combined). Our main results show that whether or not there is fast algorithm for AAttC critically depends B, the magnitudes of the entries in the input matrices.

---

We first show lower bound, that when B 2 (Vylogn), it is impossible to design a truly subquadratic-in algorithm. Our lower bound makes use of the Strong Exponential Time Hy- pothesis (SETH) [IPO1], a popular conjecture [Wil18] from the area of fine-grained complexity regarding the time required to solve k-SAT. (See Section below where we discuss SETH in more detail.)
Theorem 1.3 (Lower bound, informal version of Theorem 4.6), Assuming SETH, for every > O, there are constants C, Ca. C2 > such that: there is no (n²-1) time algorithm for the problem AAttC(n, d = C' logn, B Covlog ni =n
Our second complementary result is a new algorithm, showing that when B < o(ylogn), the problem can be solved very efficiently, in almost linear time.
Theorem 1.4 (Upper bound, informal version of Theorem 3.8). There is an algorithm (Algorithm 1) that solves AAttC(n, d = Ollog r.), B of o(ylogn), , Ea =1/ =1/polyy poly(n)) in time n +o(1)
Our Theorems 1.3 and 1.4 show that the attention computation problem AAttC exhibits very tight transition at B = O /log ri) from almost linear time to trivial quadratic time. When B < of Vlogn) smaller, the problem can be solved in almost linear time I+o(()) in the input size, using our algorithm for Theorem 1.4. When B 2 S2( VTog n.) is greater, our algorithm from Theorem 1.4 no longer applies, and furthermore our lower bound from Theorem 1.3 shows that it is impossible to solve the problem in truly subquadratic time, matter what algorithmic techniques one uses (assuming SETH).
It has been observed in LLM implementations in practice that computations are much faster when one assumes that the matrix entries are bounded can be well-approximated using small number of bits (see, e.g., IZBIW19, Section 2] and [KVPF20, Section 3.2.1J). Our work can be viewed as giving theoretical explanation for this phenomenon, and helping explain why techniques like quantization [ZBIW 19] and low-degree polynomial approximation [KVPF20] have been so effective practice.
Related Work.
A recent work by Zandieh, Han, Daliri, and Karbasi [ZHDK23] was the first to give an algorithm with provable guarantees for attention approximation. Their algorithm makes use of locality sensitive hashing (LSH) techniques [CKNS20] which, algorithm for Theorem 1.4 which uses the polynomial method [ACSS20, AA22].
will discuss next, is quite different from our
In the case when d = ologg" n). they achieve a running time of roughly ((n--77-d/e where e, is relative error parameter (which is similar, though not exactly the same, as our from Definition 1.2). In particular, their algorithm applies for larger d than ours (we require d O(log n)), but we achieve almost linear time n 1+((() (whereas their running time is bounded below by 1(nl-17)), and our algorithm can handle any polynomial error ea = 1/ poly(n) (whereas they require e, 2 (2) to not increase the running time by polynomial factor).
It is natural to wonder whether further improvements with those of [ZHDK23). However, our lower bound of Theorem 1.3 shows that our algorithm of Theorem 44 is already essentially tight and cannot be substantially improved.
possible by combining our techniques
Another recent work by Keles, Wijewardena, and Hedge [KWH23] was the first to prove lower bound for attention computation assuming SETH. They prove, among other results, that AAttC cannot be solved in truly subquadratic time in the case when d = w(og n.). Our Theorem 1.3 improves their result to also hold for d = Olog nt), and to show how the complexity changes with the magnitude of entries (which not studied by [KWH23]). As we discuss more shortly, both our lower bound proof and [KWH23] use the high-level technique of [BIS17]. although our more fine-grat fine-grained analysis of the parameters d, requires a more intricate analysis and the use of other techniques from fine- fine-graines grained complexity related to approximate nearest neighbor search [Rubl8 18] and the polynomial method [AA22]]
1.2 Technique Overview
Our high-level approach is to make use of similarities between attention computation and other computational problems related to Kernel Density Estimation (KDE). Such relationship was
3

---

investigated by recent work [TBY+ 19, ZHDK23]. In particular, [ZHDK23] was inspired to apply LSH techniques attention computation because of the prevalence of LSH KDE algorithms [CS17. BCIS18, CS19, CKNS20J. The main conceptual idea behind our results is that different techniques from the KDE literature, other than LSH, can be modified to apply this setting and yield tight algoriths and lower bounds.
To design our algorithm for Theorem 1.3, we instead build off of a different line of work on KDE which makes use of the 'polynomial method in algorithm design'. Suppose M E Ruxn is a matrix, f: R - R is function, and let (MV denote the matrix one gets by applying f entry-wise The polynomial method is a technique for finding low-rank approximations of f(M). ). It shows that if M has low rank, and if can be approximated by a low-degree polynomial. then the matrix fMM is very close to low-rank matrix whose low-rank decomposition can be computed efficiently.
M.
To use this to solve AAttC, we make use of recent result which bounds the degree required to approximate the exponential function by a polynomial [AA22] in order to find a low-rank approxima- tion of the attention matrix A. Prior work [ACSS20, ACM 20, AA22] applied these polynomials in a similar way to solve the Gaussian KDE problem; our main observation is that by an appropriate rescaling, this approach can be modified apply AAttC well.
The proof of our lower bound Theorem 1.3 builds off of another line of work on the fine- ane-grained complexity of KDE problems [BIS17, ACSS20, AA22]. The main idea is to give fine-grained grained reduction from the well-studied problem of Approximate Nearest Neighbor search ANN. In ANN, one is given input n vectors of dimension d, and error parameter
O, and the goal is to find pair of vectors whose distance is at most (1 + <) times the minimum distance between any pair of the vectors. The straightforward algorithm for ANN runs in quadratic time, and it is known that it is impossible to solve ANN in truly subquadratic time assuming SETH [Rub18].
In order to prove our lower bound, we show that AAttC can be used to solve ANN. The key idea is that, if the matrices Q and K from AAttC are formed by concatenating the input vectors to the ANN problem, then the nearest neighbor vectors correspond to the largest entries of the attention matrix A. not immediately clear that AAttC can be used to detect large entries of A, since the output is rescaled by the matrix D 1, but show that this can be overcome with some modifications to the input vectors which approximately balance the rows of A. Prior work [BIS17, ACSS20, AA22] used very similar approach give lower bounds for KDE problems, although KDE doesn't involve any rescaling factors.
Roadmap.
In Section introduce relevant notation and tools from prior work. In Section 3, we present and analyze our attention algorithm. In Section 4, we prove our fine-grained attention lower bound. In Section 5, we provide a conclusion for this paper.
2 Preliminaries
We work in the standard real-RAM model and assume arithmetic operations on real numbers can be performed in constant time in our algorithms.
We use Tmar(a. b, c) to denote the time to multiply an a x matrix with another b x matrix. In fact, will only make of the straightforward, practical bound Taar (a, b,c) O(ab)). In principle, fast theoretical matrix multiplication algorithms could be used instead to improve this bound and speed up our algorithms here (in exchange for making them less practical). That said, because of our parameter settings', 3 we will see that faster matrix multiplication could only improve low-order terms our running times.
For any positive integer, we use [n] to denote set {1,2, ... .n).
For matrix M, we write || to denote its fox norm, i.e., M| ||x := maxi,j [Mag]- For matrix M, we use T to denote its transpose.
We use 1, to denote length-n vector whose entries are all 1s. We use O,, to denote length-n vector whose entries are all Os.
We will make use of Tmaa (n, o1)) (0))), which can be solved straightforwardly in time n=+t((1), and which cannot be solved much faster since it has input size n I+o(1)

---

For any matrix A E Rux", we use exp(A) E R=x* to denote the matrix where exp(A)i, exp(A.,). In other words, all the exp() operators in this paper are applied entry-wise to matrices. In particular, we will not use matrix exponentials in this paper.
For vector E R", use || denote its number of non-zero entries, denote its & norm, i..., | == Ei=1 | and we use || to denote its lg norm, i.c., || == (E m3 1/2 For vector we use x to denote its transpose.
||||| to
2.1
Additive Error for Polynomial Approximation
Our algorithm for attention computation will critically make use of polynomial approximation for the exponential function. In particular, we use the following tight construction from previous work [AA22].
Lemma 2.1 ([AA22]). Let B > 1 and let (0,0.1). There is polynomial P IR + R of degree =
θ max Log(log(1e)//B)
log(1/c)
B)) such that for all E 10, B), we have
PP(x) exp()=) <.
Moreover, P can be computed efficiently: its coefficients are rational numbers with poly(g)-br integer numerators and denominators which can be computed in poly(g) time.
2.2 From Additive Error to Relative Error
We note that our setting, Lemma 2.1 can be used to give relative error approximation as well:
Corollary 2.2. Let 1 and let c e 0.,0.1). There is a polynomial P R R of degree log(l/e) = (max lon(log(1/e)/B] B}) such that for all a E [-B, B}. we have
IP(z) - exp(x)) < exp(x).
Proof. By Lemma 2.1, there is polynomial Q RR + R of degree g 0(1 log(1/e)
(ng(lo[(1/e)// B))
such that, for all y E [0, 2B) we have [Q(y) - exp(y) y)) < e. Our desired polynomial is the rescaled P(z) = Q(x + B exp(B). Indeed, for any x E [- B), have exp(x) 2 exp(-B). and
P(x)) exp(x)) = 1Q(x + B)/ exp B) exp(x))
1Q(x B) exp(x + /3)l exp(B)
I e/ eXp(B)
exp(x),
as desired.
3 Attention Algorithm
In this section, we show how to use polynomial approximations for the exponential function order approximately perform attention computations. In Section 3.1, define the type of low-rank matrix approximation which we will use. In Section 3.2, we show how polynomial approximations can give rise to such low-rank matrix approximations. In Section 3.3, bound the entries of the matrix QK E Rux" (before converting it to the attention matrix) to confirm that our polynomial approximation applies. In Section 3.4, we state our main technique for approximating the attention matrix. In Section 3.5, we show how to control the eтrоr propagation from A to the rescaling matrix D. In Section 3.6, we further explain how to control the error propagation from D and the resulting attention matrix. Finally, in Section 3.7, we conclude our general algorithm, and in Section 3.8, we appropriately select the parameters achieve almost linear time.
3.1 Matrix Low-Rank Approximation
Definition 3.1. Let r 2 1 denote a positive integer. Let e E (0,0.1) denote an accuracy parameter. Given matrix A R2x= say À E Bax" is an (e, approxaaptoi of A if
5

---

Ā = U1 U2 for some matrices U1, U2 E R (t.e., A has rank at most r), and
A., A e. A1, for all (i.j) E [n]?
3.2 From Low Degree Polynomials to Low Rank Matrices
Lemma 3.2. Let M XYT Ruxn denote matrix with X,Y Raxd Let P(x) denote degree-g polynomial, and define (3(gtd)).
2g
There is an algorithm that runs O time and, given as input the matrix X, Y, constructs matrices U1, U2 Ruxr such that P(M) U.U. (Here, P(M) denotes the entry-wise application of P to M.)
Due to space limitation, we defer the proof of Lemma 3.2 to Appendix A.
3.3 Matrix QK Has Bounded Entries
Lemma 3.3 (Bounded entry). Suppose B 2 1 and matrices Q, K Raxd have ||01l I B and
| S B. Then, we have
IIQK" /d BB²
Proof. For each (i,j) [n]] x [m], we have
d =
I(QK )as
Qu,K5,l
1=1
$ d 110 |0 ||||
I d B²,
as desired.
3.4 Key Lemma
Our key lemma shows that, even though the attention matrix A may have full rank, it has low-rank approximation that is easy to compute:
Lemma 3.4. Suppose Q, K E Ra with [1Q1 S B, and ||]] S. B. Let A := exp(QK] /d)
Raxx For accuracy parameter E (0, 1), there is a positive integer g bounded above by
g = 0 max {
(og(lo((1/c)BB3 B²}
log(1/e)
and positive integer bounded above by
(2(g + +d)
2g
such that: There is matrix À 8 that is an (e, r) approxximmai (Definition 3.1) of A Raxn Furthermore, the matrices U1 and U2 defining À can be computed in O(n r) time.
Proof. Let M == =QKTT /d. From Lemma 3.3, we know that || |o S B2, Thus, applying Corol- lary 2.2 (with bound B² its entries), there is degree-g polynomial P such that the matrix A (See the definition of (c, appapoooximiioonn in Defini
Ā P(M) is an (c, approxxmaattiiioo tion 3.1.) We can then compute U1, U2 using Lemma 3.2, which gives the bound
2(g + d) 2g
This completes the proof.

---

3.5 From A to D
Lemma 3.5. Let A E R be any matrix whose entires are all positive and e (0,0.1) be any parameter. Let [Raxu be any matrix such that, for all (i, E [m]] x [n], we have
| A S..ssEA EA Aij-
Define the matrices D, D E R by D = diag and D = diag( (A1n). Then, for all i E [n] we have
D Da,al EA Di.
Due to space limitation, we defer the proof of Lemma 3.5 into Appendix A.
3.6 From and D Attention Matrix
Lemma 3.6. Let CD E (0,0.1 and B > be parameters, and let V E Raxd denote a matrix with | < B. Let A E _xx= be any matrix whose entires are all positive, and let A E Ruxn be a matrix such that, for all (i,5) E [n] x [n] we have
[
A.pl EA Aij-
Let D, D E Raxx be any diagonal matrices with positive entries on their diagonals, with the property that, for all E [n], we have
1 - Dial SDD DD.4.
Then, we have
D- DD-AAVIlo I (EA EE) B.
Due to space limitation, we delay the proof of Lemma 3.6 to Appendix A.
3.7 Main Upper Bound
Theorem 3.7. For positive integers n, and real parameters c > O and B > there are positive log(l/e)
integers = (maxx log(lo((1/e)/ B25 B²}) and r (2(g+d)) such that: There is an algorithm 2d
(Algorithm 1) that runs in O mat (n, d) + nrg) time solve AAttC(n, d, B, (Definition 1.2).
Proof. The running time of each step shown in Algorithm 1; its running time follows from Lemma 3.4. Its correctness follows from Lemma 3.5 and Lemma 3.6.
Algorithm 1 Our Algorithm
1: procedure POLYATTENTION((? E Bxd K E Rexd V E paxd n, d, B, > Theorem 1.4 is the accuracy output
2:
3:
> ||Q [ [V < B
4:
Omaxx Tog(1/e) r <- 7-(2(00t)) (2(
log(log(1 (1g((///))/ B²))
5:
6:
Construct U1, U2 E Raxr via Lemma 3.4
D O(rr)) time
7:
T - U1 (U-] 1n)
O time
8:
D-1 = diag(w
D O(n) time
9:
Compute U. V E Bexd
Takes Tiaas (r, n, d) time
Compute U, (UT V)
Tmaz (n,, r, d) time
T D-1. (U, (U
((d)) time
return T
D T E R Baxd
13: end procedure
7

---

3.8 Proof of Theorem 1.4
Theorem 3.8 (Upper bound, formal statement of Theorem 1.4). AAttC(n, Ollogn), n of (llog 1)), E 1/ poly(n)) can be solved in time mast (n, n((1) d) n nttll))
Proof. If we select the parameters
B = ol nl
c 6= 1/ poly(n), d= Ologn)
Theorem 3.7, then we see that
log(1/e)
g 0 (max
log (log(( (1/c)/ (/2) BB3)
log(n)
(max
log(log( (logg(n))/33)) B²
log v
(maxx
0 log ologn))
log log
o(log m),
where the second step follows from = 1/ /olly(n) and the third step follows from B = of Vlog n). Since g o(log n.), let write
(log f for some (1). We thus have that
r= (2 + g) 0) 2g
c(d + g) 5)
29
20gglog((ogn)//)) < 20 log(s/f) 20(log < n1)
<
The second step follows from the generic bound (*) S (ea/b)o for s S a, and the third step uses that d = (logn)
Since d, 9 are all bounded by (1), our final running time is n ( as desired.
L
4 Hardness
In this section, we prove our fine-grained lower bound for attention computation. In Section 4.1, we state the Strong Exponential Time Hypothesis (SETH), the main hardness assumption we will use. In Section 4.2, we define the approximate nearest neighbor search problem, and its known hardness assuming SETH. Finally, in Section 4.3, we give a reduction from approximate nearest neighbor search to attention computation, which implies our hardness result.
4.1 Fine-Grained Hypotheses
The Strong Exponential Time Hypothesis (SETH) was introduced by Impagliazzo and Paturi [IPO1] over 20 years ago. It is a strengthening of the P * NP conjecture, which asserts that our current best SAT algorithms are roughly optimal:
Hypothesis 4.1 (Strong Exponential Time Hypothesis (SETH)). For every > 0 there is a positive integer k such that k-SAT on formulas with variables cannot be solved in ((2(1- time, even by a randomized algorithm.
SETH is popular conjecture which has been used prove fine-grained lower bounds for wide variety algorithmic problems. See, for instance, the survey [Wil18].
4.2 Nearest Neighbor Search
We will make use of known relationship between SETH and approximate nearest neighbor search.
Definition 4.2 (Approximate Hamming Nearest Neighbor Search (ANN)). For a parameter
in the (1 + Appprrxxmmate Hamming Nearest Neighbor Search problem for n vectors of dimension
we are given as input two sets A, B C (0, 1)d with A = B = n, and our goal to find an a" e A
and b* E B satisfying ||a" - Iloo < (1 + min.e A B ||a bllo-

---

(This is sometimes called the 'bichromatic' ANN problem, and monochromatic version has also been studied; see, for instance, [SM19].) Rubinstein [Rub18] showed that for certain parameters. is impossible to substantially improve on the straightforward quadratic-tima algorithm for ANN assuming SETH:
Lemma 4.3 ([Rub18J). Assuming SETH, for every > O, there are E (0, and C'> > 0 such that (1 Approxxmmmtt Hamming Nearest Neighbor Search in dimension d = C' log n requires (2n ((22-9) time.
Remark 4.4. We may assume that 4.3 holds even in the special case where each input vector from A and B has half its entries equal to 0 and half equal to 1. Indeed, for any vector a {0, 10,1), we can construct a new vector ā E {0,1] 1]22d given by a = [a_ a Here a E {0, fa is the binary complement of vector a, i.e., ā = 1 - a, for all i E [d]. Thus. a = d. We can similarly construct a new vector b E {0, 1] for each b E B. After this transformation, for any E A and b E B, we have l|ã bllo- 2 ||a bllo. it suffices find an approximate nearest neighbor among these transformed vectors.
For convenience in our the analysis, we define a gap version of approximate nearest neighbor search problem Gap-ANN(ni, d, e)..
Definition 4.5 (Gap approximate nearest neighbor search Gap-ANN(mn, d, (,)) Let d denote two positive integers. Lei i > 0 denote a threshold parameter. Let denote a accuracy parameter. Given two sets of points A = {a1,... (n) C {0, 1] and B (b1, an) C {0, 1] For each i E [n]. we need to distinguish the following two cases
* Case 1. There exists j E [n] such that |a; - llo < t.
* Case 2. For all j E [m]] we have ||a, by1l (1 E)
An algorithm for Gap-ANN(rr d, it e) can be called log log(nd) times binary search for the answer to ANN, so Lemma 4.3 holds as well for Gap-ANN(n, d, t)..
4.3 Hardness Result
In the remainder of this section, we prove our lower bound for attention computation:
Theorem 4.6 (Main Result, formal version of Theorem 1.3). Assuming SETH, for every sufficiently small q > O, there are constants C > 0 and Ca 0 and C3 >> 1 such that Approximate Attention Computation AAttC (Definition for parameters (n, d C' log B Coylog n, Ea = Co) requires 2-) time.
Proof. This follows from combining Lemma 4.3 (hardness for approximation nearest neighbor search) and Lemma 4.7 (a reduction from approximate nearest neighbor search to approximate attention computation) which we prove below.
Lemma 4.7. For any constant C2 (0,0.1): For every c 0 and C 0, there exist constants Ca > and C. > O and such that, if AAttC (Definition 1.2) for parameters (2n, d 2C logn, B = C. Civlogn,
be solved in time T, then Gap-ANN(n,, d C nggnt,, t,c) (Defini- tion 4.5) can be solved in time (TT + C.)
nogn,c,, C, CCa)
Due to space limitation, we defer the proof of Lemma 4.7 to Appendix B.
5 Conclusion
In this work, we showed that how quickly one can perform attention computation depends critically on the magnitude, B, of the entries of the input matrices. Our main idea was to make use of similarities between attention computation and KDE, and show how many known techniques for KDE can also be used in this setting. Since KDE is very well-studied problem, it would be exciting to see what other techniques can applied to attention computation as well. One limitation of our lower bound result is, it is conditional lower bound which is based on well-known conjecture SETH in the area of complexity. would be interesting to show unconditional lower bound for future work. As far are aware, our work does not have negative socictal impacts.

---

Acknowledgements Josh Alman was partly supported by a grant from the Simons Foundation (Grant Number 825870 JA). The authors would like to thank Beidi Chen for helpful discussions related to L.,,,, and Feyza Duman, Chinmay Hegde, and Piotr Indyk for helpful comments on an earlier draft. The authors would like to appreciate very constructable feedbacks for NeurIPS 2023 Reviewers. The authors would like to thanks Lichen Zhang and Ruizhe Zhang for useful and helpful suggestions about proof-reading the paper.
References
[AA22] Amol Aggarwal and Josh Alman. Optimal-degree polynomial approximations for exponentials and gaussian kernel density estimation. In 37th Computational Complexity Conference (CCC 2022). Schloss Degssull-Lciinniz-Zenttrr für Informatik, 2022.
[ACM*20) Josh Alman, Timothy Chu, Gary Miller, Shyam Narayanan, Mark Sellke, and Zhao Song. Metric transforms and low rank matrices via representation theory of the real hyperrectangle. arXiv preprint arXiv:2011.1 2020.
[ACSS20] Josh Alman, Timothy Chu, Aaron Schild, and Zhao Song. Algorithms and hardness for linear algebra on geometric graphs. In 2020 IEEE 61st Annual Symposium on Foundations of Computer Science (FOCS), pages 541-552, IEEE, 2020.
[AS23] Josh Alman and Zhao Song. How to capture higher-order correlations? generalizing matrix softmax attention to kronecker computation. arXiv preprint 2rXiv:2310.04066 2023.
[BCIS18] Arturs Backurs, Moses Charikar, Piotr Indyk, and Paris Siminelakis. Efficient density evaluation for smooth kernels. In 2018 IEEE 59th Annual Symposium Foundations of Computer Science (FOCS), pages 615-626. IEEE, 2018.
[BIS17] Arturs Backurs, Piotr Indyk, and Ludwig Schmidt. On the fine-grained complexity of empirical risk minimization: Kernel methods and neural networks. Advances Neural Information Processing Systems, 30, 2017.
[BMR 20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared 1) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems (NeurIPS), 33:1877-1901, 2020.
[CDL+22) Beidi Chen, Tri Dao, Kaizhao Liang. Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. In International Conference on Learning Representations (ICLR), 2022.
[CDW+21] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scat- terbrain: Unifying sparse and low-rank attention. Advances in Neural Information Processing Systems (NeurIPS), 14411411-17 2021.
[CGRS19] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: /904. 10509, 2019.
[CKNS20] Moses Charikar, Michael Kapralov, Navid Nouri, and Paris Siminelakis. Kernel density estimation through density constrained near neighbor search. In 2020 IEEE 61 Annual Symposium on Foundations of Computer Science (FOCS), pages 172-183. IEEE, 2020.
[CND+22] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint 121331, 2022.
[CS17] Moses Charikar and Paris Siminelakis. massinng-base-essine high dimensions. 20/7 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS), pages 1032-1043. IEEE, 2017.
for kernel density in
10

---

[CS19] Moses Charikar and Paris Siminelakis. Multi-resolution hashing for fast pairwise summations. In 2019 IEEE 60th Annual Symposium on Foundations of Computer Science (FOCS), pages 769-792. IEEE, 2019.
[DCLT18] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- training of deep bidirectional transformers for language understanding. arXiv preprint vrXiv:1810.04800 2018.
[DKOD20] Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf- efficient attention using asymmetric clustering. Advances in Neural Information Pro- cessing Systems (NeurIPS), 33:6476-6489, 2020.
[DLS23] Yichuan Deng, Zhihang Li, and Zhao Song. Attention scheme inspired softmax regres- sion. arXiv preprint arXiv:2304. 10411, 2023.
[DSZ23] Yichuan Deng. Zhao Song, and Tianyi Zhou. Superiority of softmax: Unveiling the performance edge over linear attention. arXiv preprint arXiv:2310.11685, 2023.
[GSWY23] Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. fast optimization view: Refor- mulating single layer attention in Ilm based tensor and svm trick, and solving it in matrix multiplication time. arXiv preprint arXiv:2309.07411 2023.
[GSY23] Yeqi Gao, Zhao Song, and Xin Yang. Differentially private attention computation. arXiv preprint arXiv:2305.0417 04701, 2023.
[GSYZ23] Yeqi Gao, Zhao Song, Xin Yang, and Ruizhe Zhang. Fast quantum algorithm for attention computation. arXiv preprint arXiv:2307.08045, 2023.
[HJK+23] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. arXiv preprint arXiv:2310.05869, 2023.
[IPO1] Russell Impagliazzo and Ramamohan Paturi. On the complexity of k-sat, Journal of Computer and System Sciences, 62(2):367-375, 2001.
[KKL200] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient trans- former. arXiv preprint arXiv:2001.0444 2020.
[KMZ23] Praneeth Kacham, Vahab Mirrokni. and Peilin Zhong. Polysketchformer: Fast trans- formers via sketches for polynomial kernels. arXiv preprint arXiv:2310.01655, 2023.
[KVPF20] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Trans- formers mns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning (ICLR), pages 3156-5165. PMLR, 2020.
[KWH23] Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In International Conference on Algorithmic Learning Theory, pages 597-619. PMLR, 2023.
[LWD*23] Zichang Liu, Jue Wang. Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song. Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient llms at inference time. International Conference on Machine Learning, pages 22137-22176, PMLR, 2023.
[Ope22] OpenAl. Chatgpt: Optimizing language models for dialogue. In OpenAl Blog.
https://openai.coom/ http://pppnniicoomm/bo//ccaggg/ Nov 2022.
IRNS+ 18] Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, al. Improving language understanding by generative pre-training. 2018
[Rub18] Aviad Rubinstein. Hardness of approximate nearest neighbor search. In Proceedings of the 50th annual ACM SIGACT symposium on theory of computing (STOC), pages 1260-1268, 2018.
11

---

[RWC* 19] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[SM19] Karthik C. S. and Pasin Manurangsi. On closest pair euclidean metric: Monochromatic is hard as bichromatic. 10th Innovations in Theoretical Computer Science (ITCS), 2019.
[TBY+19) 191 Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Rus- lan Salakhutdinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 4344-4353, 2019.
[VSP+17] 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is you need. Advances in neural information processing systems (NeurIPS), 30, 2017.
[Wil18] Virginia Vassilevska Williams. On some fine-grained questions in algorithms and complexity, In Proceedings of the international congress of mathematicians: Rio de janeiro 2018, pages 3447-3487. World Scientific, 2018.
[WLK+20] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04760 2020.
[ZBIW19] Ofrr 7afrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing- -NeurIPPS Edition (EMC2-NIPS), pages 36-39. IEEE, 2019.
[ZHDK23] Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In ICMI., arXiv preprint arX1v:2302.02551, 2023.
[ZRG+22] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint 2205.02000000066.. 2022.
[ZSZ+23] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H _2 o: Heavy-hitter oracle for efficient generative inference of large language models. In NewrIPS. arXiv preprint arXiv:2306. 14048, 2023.
12

---

Appendix
Roamdap.
We provide the missing proof's of our upper bound (algorithm) result in Section A. We provide the missing proofs of our lower bound (hardness) result in Section B.
A Missing Proofs for Upper Bound
A.1 Proof of Lemma 3.2
Lemma A.1 (Restatement of Lemma 3.2). Let = XX E ggxx denote matrix with X, Y Raxd Let (z)) denote degree-g polynomial, and define (21g+d)).
2g
There is algorithm that in (nrg) time and, given input the matrix X,Y X,Y, constructs matrices U1, U2 E Ruxr such that P(M) U.U. (Here, P(M) denotes the entry-wise application of P to M.)
Proof. Let P(z) denote the degree-g polynomial. Expand it in terms of its coefficients
P(x) Σ x².
i=0
Consider the function K Rd x Rdd + R defined by, for Rd.
K(u,v) == P((u, u)).
K is a degree-2g polynomial in the 2d entries u1, Udd U1 va of the vectors v. Define the set of its variables,
V = [ui."
ud. U1. Fa
Let FF denote the set of functions
F 1.V 10,1, 2gg Σs(v) 2g
vEV
We can count that
2d + 2g
F
2g
Hence, there are coefficients o for each E F such that
K(u, v)
C
(())
IEF
3EV
Define
Vú = {ur," :. ; ua)
and
=- V\Va.
We define du IRdd - [F by, for any F,
du (ui, ... ua)) = Ct
u
((w))
e.EV,
Similarly, we define d. Rdd + IR by, for any t t F,
(0 va):
II
(((4)
,,V,
13

---

Thus, we have
K(u, v) = (du(u), @o(u)).
For i E [n], let X, Rd denote the i-th row of X, and let Y; R denote the i-th row of Y. Our algorithm can thus construct
the matrix U1 E |R (5] whose i-th row is the vector ou (2;) for i E [n]], and
the matrix U/ "= whose i-th row is the vectors (y:) for e [n]]
Each entry of these matrices be constructed by multiplying together at most variables, these n X r matrices can be constructed in time ((rrg) as desired.
A.2 Proof of Lemma 3.5
Lemma A.2 (Restatement of Lemma 3.5). Let A E Raxx be any matrix whose entires are all positive and ca (0,).1) be any parameter. Ler RRxx be any matrix such that, for all (i. E [m] [m].
we have
Au. A S CA Au,-
Define the matrices D, E Ruxn by D = diag(Aln) and D = diag( (A1m). Then, for all i [n]
we have
D Dr.el SEA Des-
Proof. We calculate that
Ds.5 Da.l = Σ
Ai.
Aiil
n
A
Av.l
j=1
R
<
j=1
EA
D..
where the second step follows from triangle inequality. This completes the proof.
A.3 Proof of Lemma 3.6
Lemma A.3 (Restatement of Lemma 3.6). Let EA, E (0,0.1) and B > 1 be parameters, and let E nxd denote a matrix with |VV I B. Let A E Ruxn be any matrix whose entires are all positive, and let À E BRxn be a matrix such that, for all (i,j) [m] x [m] we have
A - A E Auj-
Let D, D E Raxx be any diagonal matrices with positive entries on their diagonals, with the property that, for all E [r]]. we have
1is - Di.el Sen en.... Di.i-
Then, we have
- Āv -D1 D-IAVlooo - (EA+ B.
14

---

Proof. We have
DV - D A & 1 ĀV - DD- -V -1 ĀV - -D-AAVV (1)
We now bound each of these two terms separately.
First, for each (i E [x]] x [d].
( ĀV D-1An)
> (D -D,)) A.. Vi.5l
1=1
1D - -D7.) Aud |V
1=1
Dr. DEs A., -|V]
Di,Dii
1=1
| AAull |V
S(D'
1=1
Σ Dj Auu |
=
ED |V
S <D
(2)
where the second step follows from the triangle inequality, the forth step follows from I(D;, Da,)/Daal S ED. the fifth step follows from D and A., > 0, and the last step follows from assumption on V.
Second, for each (i, E [n] x [d].
I(D Äv - Dー AV)
D (Art - AAr.) Visl
1=1
Σ
D 1(A,, - A..)l -IVIl
n
Σ D.d 1(A,, A.,)l -
Σ Dj. EAAii B
cA B,
where the second step follows from triangle inequality, the third step follows from D.. 0, the forth step follows from A Ar,l S CA* A and the last step follows from definition of D The result follows by combining Eq. (1), and two inequalities (Eq. (2) and Eq. (3)).
B Missing Proofs for Lower Bound
Lemma B.1 (Restatement of Lemma 4.7). For any constant C, (0, 0.1): For every > 0 and C O, there exist constants C. O and C2 > and such that, if AAttC (Definition 1.2) for parameters (2n, = 2C' log B = Cs Vlg C. nCa) can be solved in time T, then Gap-ANN(n, C log n, t, c) (Definition 4.5) can be solved in time ((T + n²-c.).
Proof. We give an algorithm with the stated running time for Gap Bap-ANN(n,, d = C'log flognnt, (,e.. Let e >> 0 be parameter we will choose later (it will a function of C and C.). Our algorithm will proceed to one of two cases depending on the value of t. Ift < clog n, then we will use one algorithm
15

---

which runs in time Onn2-C+). Otherwise, if 2 e clog n. we will use another algorithm which runs time O
Case 1: << <logg
Let a1,-.., An. bi, bn {0. 1)dh be the input vectors Gap Gap-ANN, and let E [0, d] denote the target distance. Recall that C' logn.
In this f < clog n case, we will simply brute-force for the answer in the following way: We first store the vectors ....... bn lookup table, then for each i [n], we iterate over all vectors b' E {0. 1) which have Hamming distance at most from and check whether b' is in the lookup table. This determines whether there is á B at distance most & from ag, as desired.
For each i E [m)], we need to iterate over (+) choices for the vector b' so the total running time will be On ())). By standard bounds binomial coefficients, know that
d S n (Clog n (
clog
n (C+((C
for some function f Rzo Rx0 Rso with the property that, for any fixed C' >> 0, we have
lim f(C,c)=
We can thus pick sufficiently small constant c depending only C2 and C such that (CC, 1 C, and this entire brute-force takes (n time. Case 2: t 2 clogn.
Let a..... On b.. bn E {0, 1j4 denote the input of Gap-ANN(n,, d, t, (Definition 4.5), and recall from Remark 4.4 that may assume each has half its entries 0 and half its entries 1. We will explain how to construct an Attention matrix using this instance.
Let Co be such that
t
Co log
(4)
Let 3 > 0 and d 2 d denote parameters we will choose later (see Eq. (9) and Eq. (6), respectively).
Define > by
T :: cxp(3/2).
Intuitively, our goal picking these parameters that 7 will be an upper bound on entries of the attention matrix, i.e., we will have:
z max (ap(B(a,,,)/0))
ke (E(])
We will make use of an algorithm for the AAttC d, ea) problem, for the following parameters:
2n, d 2d,
(6)
B = where C::= /40C/(Coe),
= n C. where C. = 1
(8)
Furthermore, set
B = B
(9)
16

---

We define Q Rixd and K E R
0
d2
bg
T
0
Q: QB=VB-
2
and K = /B.
0
d
0g
I
0
1/
1
Since each entry of Q and K either vB or 0, it follows that
|0 S V = B
| S v /= = B
|IQK / I 3.d
d = 3 = B².
Using the above construction of Q E R R and K Rudd we note that
A = exp(QK RRx
given by
expp(((n, bi)/d) exp(B(a2,bi)/d)
exp(3(a1, b2) exp(3(a2, b2) /2)
exp(3(an, bn) exp(3(a2, bm)/dl
exp(3(an,,b))
(xp(3(an, (a),,,nn)//) b2) exp(O)
exp(3(an, bn) exp(0)
4
exp(0)
T
exp(O)
exp(O)
exp(O)
exp(O)
exp(O)
exp(0)
(Note that we do not explicitly compute all the entries of our algorithm; we will make use only through calling algorithm for the Attention problem.)
For each (i,1) E [m]] [n]. we know that
Ars exp(3(a,,b))
S exp(31lasll 1bblloo d/d)
exp(3)
(10)
where the first step follows from definition of A Rüxa the second step follows from (a, b,) < |as]loe |رو | d, the third step follows from d < (see Eq. (6)), and the last step follows from definition of 7 (see Eq. (5)).
On the other hand, we know that that for each (i,j) E [n]] x [m].
As, 20
(11)
since it an exponential of an entry of QK /d.
Using Eq. (10) and Eq. (11), combined with our expression for A, it thus follows that
nr < (A1x); & 2nt, Vi E [7].
17

---

Since D,, = (A1g)s, thus we know that
nt Dis S 2nt,
E [
Choose the vector v E I² defined as
1 On
v =
We define as
8:= 2
(12)
We show that f 2 E, follows:
i exp(0.255 t/d) B)
exp(-0.753 0.253t/d
exp(-0.753 0.25,0,
6n
exp( exp(-0.75p 0.253Ca/C log
exp( exp(-0.750f - 0 0,250)(Co//C) logn logn)
2 n C.
Cas
where the second step follows from simple algebra, the third step follows from / = Co log n (Eq. (4)) and C Cllogn (assumption in Lemma statement), the second step follows from choice of (Eq. (7)), and the sixth step follows from choice of C. (Eq. (8)), and the last step follows from Eq. (8).
Since i. 2 Ea. if run algorithm for Approximation Attention Computation (Definition 1.2) AAttC(ã, d, B, ea), where we pick to be matrix with one row a and the 0, we can output vector E Ra such that, for all i [i].
1/4 - (D-1At)) <
Note that using Remark 4.4, we have
||a,1//d== = 0.5, Vi Ei [n], |b, = 0.5,
Vj E [n].
Therefore, for any (i,5) E [m] [m].
d (a;, by) =
(1|ail13 + 1|b, - |la by b)112)
2d
1
(0.5d + 0.5d - |la; - by b1))
2d
= 0.5 - 0.5||a: - b/1//d.
where the second step follows from ||a: l12 = |b6, = d/2, and the last step follows from simple algebra.
Recall that our goal is to determine, for each E [n]. whether there is a E [n] such that ||a; - by < t, or whether ||a; - b 2 (1 + E. for all j [n]. We will show next that we can distinguish these two cases by seeing whether uy is greater than or less than the value To = 27.
Case 2a.
18

---

If there exists an (i.5) E [m]] x [na] such that ||a; - -byl2 - t, then
b/a,, = 0.5 ( b,) 2 0.25 3(1 t/d),
where the first step follows from 2d = (see Eq. (6)). This means that
Uj 2 exp(0.25.3(1 - t/d))/(2nr) -
31 -i
27
To
where the second step follows from the definition of (see Eq. (12)), and the last step follows from the definition of to-
Case 2b.
If for all (i,j) E [n] x [n], we have |a, -o--b byl12 > t(1 + e). this implies
3/a b,)/d 0.25,3 (1 t(1 + c)/d).
Then, for all i E [ns].
u < (n -00-00000200011--000000)mmm0m +7
(1
= -1(()))((mm (n/ exp(0:258c//d)) +
-
3ē (n/ oxp(0.25.5ctdd)) +
1
< <
+
4
2f
to.
where the third step follows from definition of f (see Eq. (12)), the forth step follows from the calculation in Eq. (13) below, and the last step follows from i = fo/
Finally, b our choice of (3 and I, we can see that
exp(0.25,et/a d) = exp (0.2000033e log /d)
(xp((..2.Bea//C))
exp(10logm)
4n,
(13)
where the first step follows / Co log n (Eq. (4)), the second step follows from d = C' C'log n, and the third step follows from 3 = B2 (Eq. (9)) and the choice B (Eq. (7)).
Recent Followup Work During the preparation of camera ready, are aware of a number of new related work [KMZ23, HJK 23, AS23, DLS23, GSWY23, DSZ23, GSY23] to this paper. This work studies softmax attention, and [KMZ23] shows speed up the polynomial based attention unit via sketching techniques. [HJK* 23] introduces hyperattention, which able to handle long- context attention in near-linear time. This work is mainly focusing classical attention computation, [AS23] consider a more general tensor version computation of attention scheme. This work studies the inference computation, the recent work [DLS233 GSWY23] also considers how speedup the training process which involves the computation of gradient and hessian. Inspired by hardness result of this work, [DSZ23] shows that there is binary classification tasks (involves two datasets) that softmax attention can distinguish but linear attention cannot. [GSY23] provides result for computing the attention matrix differentially privately. [GSYZ23] introduce quantum algorithm for computing attention matrix.
19


Reviews:
[{"review_id": "kKBzQwtbmC", "reviewer": "Reviewer_ioBT", "summary": "The attention mechanism, whose runtime complexity is approximately quadratic in the sequence length $n$, is the main bottleneck to LLM inference. It is thus quite reasonable to study whether this can be *approximated* with sufficient accuracy in subquadratic time. In particular, for this paper the authors define and study the problem of approximate attention computation (AAttC), wherein one seeks a result that is accurate within a pre-specified entrywise error tolerance $\\epsilon_a$.\n\nUnder what I think are reasonable assumptions, the authors find that whether AAttC can run in subquadratic time hinges on the bound $B$ on the magnitudes of the entries of $Q$/$K$/$V$. This is slightly surprising and, in my opinion, a neat story for the paper. This conclusion consists in two results:\n\n1. Theorem 3.8: an algorithm that computes AAttC in near-linear time provided that $B = o(\\sqrt{\\log n})$. The high-level idea is\n   - Approximate the $\\exp$ in $\\exp(QK^T/d)$ with a polynomial $P$ (Lemma 2.1 from [AA22]).\n   - Compute a low-rank decomposition $U_1 U_2^T = P(QK^T/d)$ directly in linear time based on knowledge of $P$ and the entries of $Q$ and $K$ (Lemma 3.2).\n   \n    These lemmas are the main workhorses of Theorem 3.8, but effort is also made stitching everything together. The crucial thread is that $B^2$ determines the degree of $P$, which in turn determines the rank of the decomposition $U_1 U_2^T$. $B$ thus needs to grow slowly in $n$ to make the algorithm subquadratic.\n2. Theorem 4.6, which shows that if $B = \\Theta(\\log n)$, there is no subquadratic-time algorithm for AAttC (assuming the strong exponential time hypothesis, SETH). This result is based on a reduction of approximate nearest neighbor search to AAttC.", "questions": "I am curious to hear how you think your results might lead to practical speedups in attention computation for future work.", "limitations": "The authors do highlight their main limitation: Theorem 4.6's reliance on the strong exponential time hypothesis (SETH). As far as I can tell, no important limitations were omitted.", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Although past work has produced some closely related results (ZHDK23, KWH23), the authors' focus on $B$, the entrywise bound, is original and insightful. Furthermore, the proof techniques in this paper bring up interesting connections to the KDE literature which may be interesting for the community.\n- They provide a near-linear time attention algorithm based on $B = o(\\sqrt{\\log n})$, and show that the rate at which $B$ scales cannot be improved upon. In some sense their work is a complete description of the problem, conditional on their (reasonable) assumptions on how other quantities scale with $n$ ($d = O(\\log n), \\epsilon_a = 1/\\text{poly}(n)$). This is quite satisfying.\n- For a theoretical paper, this is well-written and easy to understand. The authors do a great job of motivating, roadmapping, and explaining their work.\n- This work offers a key insight about an important problem.\n", "weaknesses": "This is a pair of complexity results. It is not clear to me how the insights provided here can be operationalized to improve AAttC in a practical setting. (The authors do highlight practical examples in ZBIW19 and KVPF20, but the connections between those works and the theory developed here do not seem particularly deep).\n\nTypos\n- L117: Hedge -> Hegde\n- L133: algoriths -> algorithms\n- L134: I believe should read \"To design our algorithm for *Theorem 1.4*\" not Theorem 1.3.\n- L242: entires -> entries\n- L232, L257, Algorithm 1 L4: I believe $B$ in $\\frac{\\log(1/\\epsilon)}{\\log(\\log(1/\\epsilon)/B)}$ was meant to be $B^2$\n- Appendix L434: roamdap -> roadmap\n- Appendix L453: $v_i$ -> $u_i$ under the product sign\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful review. We really appreciate the reviewer for pointing out these typos, we will fix them in the next version. (We remark that none of the typos impact the correctness of the paper.)"}, {"review_id": "AXfxKGeCHo", "reviewer": "Reviewer_Ze7L", "summary": "In this work, the author studied the problem of how to efficiently compute or approximate the attention matrix. The major contribution can be summarized as: \n\n1. they found for the case when $d = O(\\log n)$ and entries of the input matrices are uniformly bounded by $o(\\sqrt{\\log n})$, then an algorithm that approximates the attention matrix can be constructed with sub-quadratic complexity, the key idea of the approximation is based on a polynomial approximation of the $e^x$ that was introduced by another work [AA22].\n\n2. when $d = O(\\log n)$, and the entries of the input matrices are of the order $\\Theta(\\sqrt{\\log n})$, then under the Strong Exponential Time Hypothesis, it is impossible to approximate the attention matrix with small error and sub-quadratic complexity.\n\n\n ", "questions": "Please see some questions I mentioned in the weakness section.\n\nAside from that, can the author explain the claim \"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\". Why the results in this work explain this in practice?", "limitations": "Please see the weakness section.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. I think in this work, the author did a good job explain the basic intuitions and ideas behind their approach. Especially, breaking down the proof steps into smaller parts improved the readability of this work. \n\n2. On the other hand, by elaborating the related works and contributions, the reviewer can have a better understanding of the contribution of this work.", "weaknesses": "There are several weakness of this work:\n\n1. The major contribution of this work is not emphasized enough or minor. For example, the technical tools used in this work are developed by previous works to tackle similar problems. As a theory paper, I think the author at least needs to address what technical contribution is made in this work. To me, it seems that the major results are straightforward applications of some closely related recent works.\n\n2. The lack of simulation results is another weakness. As the author is trying to address some practical/computational aspects of the problem, experimental results are essential to validate the theory and claims.\n\n3. Some of the claims in the theorem are ambiguous. For instance, in theorem 4.6, there are many constants $C_a, C_{\\beta}, C, q$ used. What are those constants? Numerical constants? What are the relationships between them? Statements like this is not rigorous at all for theory papers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful review.\n\n1. We discuss our technical contributions and how they compare with prior work in sections 1.1 and 1.2 of our paper. Perhaps the biggest technical insight is that a number of tools from the literature on KDE, which were not previously known to be applicable to Attention, can in fact be applied in this setting. (Indeed, prior work like [ZHDK23] and [KWH23] knew about the KDE connection but was not able to prove such strong results.) We also believe that the details of using these techniques in this setting are non-trivial. For instance, we initially thought that proving a lower bound in this way would be impossible because of the normalizing factor D^{-1} which appears in Attention computation (Definition 1.1 in our paper): previous lower bound reductions for KDE requiring finding large entries of the output, but the normalizing factor ensures that all the entries are roughly equal to each other.\n\n2 and Questions: please see the general remark to reviewers above.\n\n3: All our theorems and lemmas are fully formal, precise and rigorous mathematical statements. For instance, for the mentioned Theorem 4.6, we give a fully quantified logical statement (for all q, there exist constants C, C_alpha, C_beta, such that …) and all the variables in the statement are captured by the quantifications.\n\nThe dependance of the constants C, C_alpha, C_beta on q is an interesting question which touches on a major open challenge in fine-grained complexity theory. The statement of the Strong Exponential Time Hypothesis [IP01] says that for every eps > 0, there exists a k such that k-SAT cannot be solved in time 2^{(1-eps)n} (see Hypothesis 4.1 in our paper). The relationship between eps and k is a major open problem; the current best SAT algorithms achieve eps = 1/O(log k), although it is generally believed that there should be better SAT algorithms that improve on this dependence. This unknown propagates to yield unclear relationships between constants in most fine-grained complexity results. For instance, it also appears in hardness for approximate nearest neighbors [Rub18] (see Lemma 4.3 in our paper), and thus in turn appears in our results which build off of [IP01] and [Rub18].\n\nThat said, some relationships between the constants do appear in our proofs; see, e.g., line 518 in our supplementary material. We don’t think the details of these constants are particularly salient, and we chose not to emphasize them in our theorem statement to make it more readable.\n\n(The references [IP01] and [Rub18] are cited in our submission, but we copy them here below for convenience.)\n[IP01] Russell Impagliazzo and Ramamohan Paturi. On the complexity of k-sat. Journal of Computer and System Sciences, 62(2):367–375, 2001.\n\n[Rub18] Aviad Rubinstein. Hardness of approximate nearest neighbor search. In Proceedings of the 50th annual ACM SIGACT symposium on theory of computing (STOC), pages 1260–1268, 2018."}, {"review_id": "NKVKsmaJKx", "reviewer": "Reviewer_ZPux", "summary": "The paper is a purely theoretical paper discussing improving the computational time required to solve the attention computation problem which is known to be quadratic in time. The idea hinges upon the use of the boundness of the entries of the matrices involved in such a problem. \nThe authors discuss mainly two cases, where the entries are bounded by a polylogarithmic term in the number of tokens which then the attention computation problem can be solved in time almost linear in the number of tokens. Otherwise, the authors prove via the \"Strong Exponential Time Hypothesis\" from fine-grained complexity theory that the running time is no better than quadratic in the number of tokens.", "questions": "1) On page 14 of the appendix, did you mean to put n instead of D at the top of the summation symbol (see equation 2)?\n2) Also on the same page, did you mean to write $\\tilde{D}_{i, i}$ instead of $\\tilde{D}_i$?\n3) Shouldnt be the $B$ in the denominator of the equation stated at Lemma 3.4 be $B^2$?", "limitations": "The authors adequately addressed the limitations of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1) The paper is well-written and easy to follow.\n2) The paper proves the existence of almost linear time algorithms for the attention computation problem, as it depends on the boundness of the entries of the matrices. \n3) The paper also shows that in case the entires of the involved matrices in the attention computation problem are \"not bounded\", then the quadratic time associated with the attention computation problem can not be improved.", "weaknesses": "The only weakness is the lack of practical evaluation. Although the paper sets a theoretical justification for, e.g., the use of quantization or low-degree polynomial approximation of the attention computation problem, it would be great to observe the quality of the proposed algorithm on the quality of transformer-like models which practices the attention computation problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful review. We really appreciate the reviewer for pointing out these typos, we will fix them in the next version. (We remark that none of the typos impact the correctness of the paper.)"}, {"review_id": "yluZnDds48", "reviewer": "Reviewer_LkJK", "summary": "This paper shows that whether or not there exists a fast algorithm for approximately computing the attention matrix ($A \\in \\mathbb{R}^{n \\times d}$, where $d = O(\\log n)$ depends on the bound of the entry value $B$\n\n1. Particularly, the authors show that when $B = o(\\sqrt{\\log n)}$, their proposed algorithm can compute the attention matrix within $1/\\text{poly}(n)$ error in time $n^{1+o(1)}$.\n2. Assuming the Strong Exponential Time Hypothesis (SETH) and the fine-grained complexity conjecture is valid, they show that there is no $O(N^{2-q})$ time algorithm that guarantees $1/\\text{poly}(n)$ approximation error.", "questions": "+ Although the paper provides an explicit way of the matrix construction $U_1$ and $U_2$, the reviewer is wondering how complex is the matrix construction in practice.", "limitations": "+ It would be nice if there were simulation results that demonstrate the practicality for some benchmark datasets on some target tasks.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "+ The results are interesting and correct\n+ \u001fAn improved hardness result is a good contribution\n+ The observation that the existence of a sublinear algorithm depends on the entry value is interesting.", "weaknesses": "The proposed algorithm is a bit incremental, both in terms of the proposed method, and the techniques to establish these results.\n\nMinor typo:\np5: In the proof of Corollary 2.2, it seems like $\\max$ operation is missing in the formula $g = \\Theta(...)$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful review. \n\nMinor typo p5: Yes, it is a typo and the max operator is missing. Thank you for pointing it out. We will fix it in the updated version.\n\nQuestion: The construction of U_1 and U_2 in Lemma 3.2 is actually very simple and practical. At a high level: each column of U_1 corresponds to a monomial of the polynomial, each row corresponds to an input vector, and the corresponding entry of U_1 is equal to that monomial evaluated at that input. (U_2 is similar.) Thus, U_1 and U_2 can be constructed in essentially the same amount of time as it takes to just write down their entries. See Appendix A.1 in our supplementary materials for the full details."}]

Author rebuttal:
We want to thank the reviewers for their thoughtful and detailed reviews.

We want to briefly expand on the connection between our result and practical papers [ZBIW19] and [KVPF20] which multiple reviewers asked about. At a high level, these two papers showed that techniques for enforcing and manipulating bounded entries can quickly perform attention computations in practice. A main motivating question for our paper is whether there is a theoretical basis for this improvement. We show that there is a fast algorithm with provable guarantees in this setting, but also that bounded entries are theoretically _required_ for a fast algorithm. In other words, we show that bounded entries are not just a heuristic used by the prior work, but that they are both necessary and sufficient for fast algorithms with provable guarantees.

Meta-review: This paper presents an analysis for the phenomenon that attention computation is much more efficient when the input matrices have smaller entries. While there is a discrepancy in the reviews, we believe that this work makes some solid technical contributions (on the theoretical side) to an important practical question. We hope the reviews and feedbacks during the discussions could be helpful for the authors to prepare the final version of this paper.
Decision: Accept (poster)