The Fine-Grained Complexity of Gradient Computation for Training Large Language Models
Authors: Josh Alman
Zhao Song
Venue: NeurIPS 2024 poster (2024)
Keywords: Strong Exponential Time Hypothesis
Fine-grained Complexity
Polynomial methods
Gradient Complexity
Primary area: learning_theory
TL;DR: A theoretical study of gradient computation, from both algorithm and hardness perspective

Abstract: Large language models (LLMs) have made fundamental contributions over the last a few years. To train an LLM, one needs to alternatingly run `forward' computations and backward computations. The forward computation can be viewed as attention function evaluation, and the backward computation can be viewed as a gradient computation. In previous work by [Alman and Song, NeurIPS 2023], it was proved that the forward step can be performed in almost-linear time in certain parameter regimes, but that there is no truly sub-quadratic time algorithm in the remaining parameter regimes unless the popular hypothesis $\mathsf{SETH}$ is false. In this work, we show nearly identical results for the harder-seeming problem of computing the gradient of loss function of one layer attention network, and thus for the entire process of LLM training. This completely characterizes the fine-grained complexity of every step of LLM training.

The Fine-Grained Complexity of Gradient Computation for Training Large Language Models
Josh Alman Department of Computer Science Columbia University joshêcs. josh@cs.coluubia columbia. edu
Zhao Song
Simons Institute for the Theory of Computing University of California, Berkeley magic. linuxkeedgmailll com
Abstract
Large language models (LLMs) have made fundamental contributions over the last few years. To train LLM, one needs to alternatingly run 'forward' computations and 'backward' computations. The forward computation can be viewed as attention function evaluation, and the backward computation can be viewed computation. In previous work by [Alman and Song, NeurIPS 2023], it was proved that the forward step can be performed in almost-lineat time in certain parameter regimes, but that there is no truly sub-quadratic time algorithm in the remaining parameter regimes unless the popular hypothesis SETH is false. In this work, we show nearly identical results for the harder-seeming problem of computing the gradient of loss function of one layer attention network, and thus for the entire process of LLM training. This completely characterizes the fine-grained complexity every step of LLM training.
gradient
Introduction
Large language models (L.LMs) have emerged as popular technologies, driving breakthroughs across many applications in natural language processing, computer vision, translation, and many other areas [VSP+ 17, DCLTI8, LOG+19 19, YDY 19, BMR+20, JZLD21, ZRG+22, 22, CND+22, T3, TMS+23, Man23, T2F YCRI22, WTB+22, WSD+23. 23, WCZ+23, ZJL 2JL+23, 24H+24, 24a, XSL24, CLL 24, WMS IMMS+241 24]. The training of these models is a computationally in- tensive process, characterized by alternating between two primary operations: forward computation and backward computation. Forward computation, or function evaluation, involves the propagation of input data through the network to generate predictions. Conversely, backward computation, or gradient computation, is the process of calculating the gradient of the loss function with respect to the model's parameters, facilitating the optimization of these parameters during training.
The efficiency these computations directly impacts the feasibility and scalability of training LLMs, particularly as models grow in size and complexity. Recent work by [AS23, AS24c, AS24a] has carefully studied the forward computation step. They demonstrated a sharp computational boundary, showing that how quickly the forward steps can be performed depends critically how large the entries are of the matrices which define the model parameters. They showed near-lineas time algorithm when these entries are small, and also proved that when the entries are large, there is algorithm much faster than the trivial algorithm, assuming the Strong Exponential Time Hypothesis (SETH) [IP01] holds.
The Strong Exponential Time Hypothesis (SETH) was introduced by Impagliazzo and Paturi [IPO1] over 20 years ago. It is strengthening of the P NP conjecture, and asserts that our current best SAT algorithms are roughly optimal (for detailed statement, see Hypothesis 3.1 below). SETH is popular conjecture from fine-grained complexity theory which has been used to prove lower bounds for a wide variety of algorithmic problems. See, for instance, the survey [Wil18].
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In other words, in some parameter regimes, the algorithm of [AS23] performs the forward steps about as quickly as one could hope for, whereas other regimes, assuming SETH, it is impossible to design a nontrivially fast algorithm. However, this leaves open many important questions about I.LM training. In the case when forward computation can be done quickly, can the same be said for backward computation? If not, then the entire training process would still be slow. Relatedly, in parameter regimes where forward computation is known to be hard, is backward computation also hard? If not, perhaps heuristic tricks could be used, or other details of the model could be modified, to speed up the overall training. As we will see shortly, the backward step is defined in much more complicated way than the forward step, and it is not evident that algorithms or lower bounds for one extend to the other.
Our study aims to resolve these questions and determine the fine-grained complexity of the backward computation phase. Our main result (which we state more formally shortly) shows that the same computational threshold from forward computation also arises for the backward problem, and that the problems are easy (or hard) in the exact same parameter regimes. Thus, the forward algorithm of [AS23] can be combined with our novel backward algorithm to perform each training step for LLMs in almost linear time when the parameter matrix entries are small enough, whereas when the entries are not small enough, neither step can be performed quickly.
In addition to characterizing the fine- mne-grained complexity of LLM training, our result for gradient computation is novel for few reasons.
Previous work on computational lower bounds only focuses on forward computations, see [AS23, KWH23, AS24c, AS24a]. To our knowledge, ours is the first work to prove hardness of a backward computation step for training an LLM or similar model.
There has been previous work on algorithms for backward/gradient computation problems [BPSW21, SYZ21, DHS+ 22, 23, GQSW24, SZZ24]. That said, most of these works focus on backwards computation in other settings. The only previous work we're aware of that studies the optimization of attention layers (for LLMs) is [GSWY23], which uses Newton methods that rely on Hessian computations. However, Hessian computation is substantially more expensive than gradient computation; our algorithm and results apply directly the gradient computation and get around the Hessian "barrier", allowing for faster algorithms in some parameter regimes, and more powerful lower bounds in others.
Bounded entries. Our result proves that the size of the entries of the matrices defining the LLM play a substantial role in determining how quickly LLM training can be performed. Prior work on LLM implementations has observed similar phenomenon, that algorithmic techniques like quantization [ZBIW19, 19, HCL* + 24] and low-degree polynomial approximation [KV2F,0), which require bounded or low-precision entries, can substantially speed up LLM operations. See, for instance, the discussion of these phenomena in [ZBIW19, 19, Section 2] and [KVPF20, Section 3.2. Our work can be viewed as giving theoretical explanation for this phenomenon.
Polynomial approximation. Our new algorithmic approach, which uses a polynomial to approximate the softmax function, is also not unlike algorithms which have found success in practice [BGVM20, KWH23, ZBKR24]. For example, see detailed discussions in in [ZBKR24, Section 4.1]. Our new algorithm improves these approaches by using theoretically optimal polynomials for softmax, and combining them with a number of linear algebraic techniques, to give provable guarantees about their correctness and near linear running time.
Follow-up work of this paper. Recently, a number of works have considered different extensions of this paper. [LSSZ24] extends our analysis into tensor attention gradient computation and [LSS+24] 24] extends our results to multi-layer Transformers. On the other hand, (..LS+244) borrows our tech- niques and provides a fine-grained attention I/O complexity for attention backward. [LLS 24c] uses our techniques to provide a fast attention gradient approximation based on Fourier transform. [LLS + 24d] computes a sparse attention matrix based on our analysis well.
1.1 Problem Definition
Before formally stating our results, we begin by precisely defining the problems we study. We begin with the following problem of computing general Attention forward layer.
2

---

Definition 1.1 (e-th layer forward computation), Given weights Q. V E Rdxd, and letting Eg E Raxd denote the f-th layer input, then E E Raxd is defined recursively
E1+1 t D exp(EE E, /d)
where
D diaglexp(E,QQK) ET /d)1n).
exp denotes the exponential function which is applied entry-wise, i.e., exp(A)is exp(As.) for all matrices A.
diag() operation takes a vector as input and generates a diagonal matrix with the entries of that vector.
In denotes the length-n all ones vector.
In mathematical terms, optimization in the context of attention computation is described as (by renaming the QK" E iddd to be X Rdxd and V Rdxd to be Y Rdxxd):
Definition 1.2 (Attention optimization). Given four d size matrices A1, A2, A3 and E E Raxdd Suppose that d x size square matrix Y E TK also given. The attention optimization problem formulated as:
min L(X) 0.5| (xpAAXXAX)AXAAAAAA E
XERaxd
Here D(X)= R is
D(X) diag(exp( XAX 1d)
and -1?. denotes the squared Frobenius norm, i.e. ||A| =
Remark 1.3. In principle, the loss function above, and resulting gradients below, should depend on both X and Y. However, since the final matrix computed in the norm L depends only linearly on Y, is straightforward to incorporate it into either an algorithm or lower bound. Thus, in this work, we focus on the case where X is variable and is a fixed input to simplify some arguments.
We thus define the Approximate Attention Loss function Gradient Computation problem follows:
Definition 1.4 (Approximate Attention Loss Gradient Computation (AAttL d, e))). Given four n X d size matrices A1 E Raxd A2 E R=xd As E Raxd E E Raxd and square matrix Y Raxd which we think of as fixed matrices. Assume that || S B, ||A2ll S B for positive parameter B. Further assume that all the entries of these matrices can be represented as
Ollog -bibit rational numbers. Let L(X) be defined as Definition 1.2. Let dL(X) denote the gradient of loss function L(z).
dx
The goal is to oulput vector ÿ such that
dl(X) dX 21 Sc
Here for matrix A, ||A max,, |Agl.
1.2
Main Results
Our main results show that there is threshold in the computational complexity of AAttL d O(logn)) depending on the bound B. When B o( o(vlog we give new near-linear-time algorithm, and when = w( <log n)), we show that such an algorithm is impossible assuming SETH. This matches the results of [AS23], where nearly identical threshold B around Vlog was also observed. Our results therefore imply that the entire LLM training process has this computational threshold.
Theorem 1.5 (Main result, Lower bound, informal version of Theorem E.5), Assuming SETH, there no algorithm running time ( for any q > for the AAttLGC(n, O(logn), w( eolloog n)) (see Definition 1.4).
3

---

Theorem 1.6 (Main result, Upper bound, informal version of Theorem D.6), Assuming entries are bounded, there is a n [++((1) time algorithm to solve AAtt d = Oloo n), B = ol o(vlog n)) (see Definition 1.4) up to 1/ poly(n) accuracy.
Our new algorithm (Theorem 1.6) builds on low-rank approximation for the attention matrix from prior work [AA22, AS23]. Incorporating these approximation into the gradient computation is not straightforward: in the forward problem, one simply multiplies the attention matrix by an input value matrix, but in the backward problem, it is combined with other matrices in an intricate (non-linear) way. We ultimately use tools from tensor algebra to get a handle on the entry-wise products and high-rank sparse matrices which arise in the gradient computation but do not typically preserve the needed low-rank structure.
Our new lower bound (Theorem 1.5) comes from a careful reduction from a special case the forward problem (where hardness is known from prior work) to the backward problem. Reducing from computing a function to computing its gradient in general is quite challenging or impossible without control over how quickly the gradient may be growing or changing, and in general, the gradient of the forward (attention) computation can behave quite erratically (which is likely necessary for the expressive power of attention units). Nonetheless, in the special case of the inputs for which attention computation is known to be hard from prior work, we are able to reasonably control the growth of these gradients and successfully perform our reduction.
Roadmap. We discuss other related works in Section 2. In Section 3, we provide the basic notation, definitions, backgrounds, and facts which we will use. In Section 4, we provide the proof sketch of our algorithm and defer the details to the Appendix. In Section 5, we briefly conclude our paper. In Section 6, we discuss the limitations of our paper. In Section 7, we provide the broader impact statement.
2 Related Work
Fine-grained Complexity. Numerous algorithmic techniques have been used in theory and in practice for attention computations. The first algorithm with provable guarantees, by Zandieh, Han, Daliri, and Karbasi [ZHDK23], used locality sensitive hashing (LSH) techniques [CKNS20), while later work by Alman and Song [AS23] used polynomial approxmation methods [ACSS20, AA22]. We particularly focus here on the latter technique, which the only algorithm we're aware of which achieves near-linear running time.
Keles, Wijewardena, and Hedge [KWH23] established the first lower bound on attention computation under the assumption of SETH. Their findings demonstrated that when d = w(og nt), it is not possible to execute forward computations in subquadratic time. The later lower bound of [AS23] further incorporated the magnitudes of the input entries into the lower bound to tightly match the aforementioned algorithms. Both use the high-level technique of [BIS17] from kernel density estimation, and build on methods derived from fine-grained complexity associated with approximate nearest neighbor search [Rub18] and the polynomial method [AA22].
Fast Attention Computation. Optimizing the computation of attention mechanisms in pre-trained LLMs, given their extensive parameter sets, has been a focal point of recent research. Various studies have explored the application of locality sensitive hashing (LSH) techniques to approximate attention mechanisms. [KKL20] introduced two methods to enhance computational efficiency, including the use of LSH replace dot product attention and reversible residual layer to substitute the standard residual layer. [CIP++11 refined this approximation, noting that LSH's efficiency does not require constant parameter updates. [ZHDK23] proposed an innovative estimator based on Kernel Density Estimation (KDE) to speed up the softmax function and matrix multiplication computations. Some recent works [HJK+ + 23, KMZ23] have specifically used sketching techniques to avoid large entries in the attention matrix. [PMXA23] developed techniques utilizing a transformer within a transformer (TinT) model to simulate the transformer's forward and backward passes, significantly increasing parameter efficiency. [MGN M3] tackled the challenge of fine-tuning LLMs with high memory demands by improving the classical ZO-SCD optimizer, creating a memory- efficient gradient estimator that requires only forward pass. [BS724] provided insights into dynamic attention problems, they provide algorithm and hardness for the dynamic setting of attention problem. [GSY+23a] introduces quantum algorithm for attention computation, opening new avenues for

---

efficiency improvements. [GSYZ24] provides result for computing the attention matrix differentially privately. [DMS23] introduces a randomized and deterministic attention sparsification algorithms for reeepparameterized feature dimension. [DLMS23] provides zero-th order method to accelerate the computation of attention. [FA23, SLBK23, LLSS224 use weights sparsity to accelerate the attention computation, but cannot reduce the time complexity. [SMN+24] compress the input token length to accelerate attention inference. [CLS 24] uses Half-Space Reporting (HSR) techniques to accelerate attention computation. [SYZ24] studies proxy for softmax attention such matrix exponential and provides fast algorithms for these proxies.
Transformer Training. Transformer architectures (the backbone of LLMs) have been trained with alternating steps of forward and backward computations since their introduction [VSP+ 17, DCLT18, LOG+19 19, YDY+ 19, BMR 20, ZRG+22]. In Appendix B below, perform computations to verify that our stated problems are the same as the forward and backward steps from the literature. Note that there are many weights update methods, such as LoRA [HSW* 21, ZL23, HSK+224, prefix turning [LL21, LSSY24), and many so on. In this paper, we consider the standard training algorithm with gradient back-propagation. On the other hand [HYW+23, HI.SL.2, WHHIL HCL++24, HWI.244, HCW+24 24] introduce the modern Hopfield models as a proxy for possible fast attention computation in training (and inference), which have been used in various applications [XHH*24, WHL 224.. Similar analyses of computational feasibility have also been conducted for transformer-bic based diffusion models, such as Diffusion Transformers (DiTs) [HWL,*244 244,, Ano24).
Preliminary
In Section 3.1, we define some basic notation we will use. In Section A.3, we state important facts related to fast matrix multiplication. In Section 3.2, provide the formal definition of the Strong Exponential Time Hypothesis. In Section 3.3, we define several intermediate functions related to softmax and exponential which will arise in our al algorithms. In Section 3.4, we define the loss function. In Section 3.5, we provide standard tensor tricks which we will use. Section 3.6, we show how to reformulate the loss function for our purposes.
3.1 Notation
For two same length vector æ and y, we use (z, y) to denote the inner product between æ and y. i.c. (x. y) Σi= T:Ui. We use I y to denote vector that i-th entry 2,y6: Let 1, denote the length-n all ones vector. It not hard to see that (Ir y, 1,,) (x, y). For vector r, we use x to denote the transpose of x. For matrix M, we use M to denote the transpose of matrix M. For a vector x. we use exp(z) )) to denote the vector that i-th coordinate is exp(=)). For matrix M,
For any positive integer n. we define [n] := (1,2....,,1) nn.
denote the matrix that (i,j)-th coordinate exp(Mr,s). For function f. we use Ō()) to denote f poly f). Let no. ni, mo, mi be positive integers. Let X E R and Y E [R=xxmm, Xmm, We define the Kronecker product between matrices X and Y, denoted X @ Y E R x x as (X @ 1) 1)jo-1)aatisde 1000-1)) is equal to X jo Yn where jo E [mo]; io E [mo]:Ji E [m1]. E [m1].
exp(M
3.2 Backgrounds on Complexity
Over 20 years ago, Impagliazzo and Paturi [IPOI] introduced the Strong Exponential Time Hypothesis (SETH), an enhancement of the P of NP conjecture. It posits that the existing algorithms for solving SAT problems are essentially as efficient as possible:
Hypothesis 3.1 (Strong Exponential Time Hypothesis (SETH)). For any > there exists a positive
integer & 2 3 for which solving k-SAT problems with variables O(2(1-e)m) time impossible, including with the use of randomized algorithms.
SETH, a widely recognized conjecture, has been instrumental in establishing fine-grained lower bounds across broad spectrum of algorithmic challenges, as highlighted in the survey [Wil18).
5

---

3.3 Definitions related with Softmax
Now, we start by some definitions about X Rdxd which will be helpful. Let x denote the vectorization of X.
Definition 3.2. Let A1, A2 E R be two matrices. Suppose that A= = A1 8 A2 E R=²xd" We define Ajo E "Raxdd be n x de size sub-block from A. Note that there n such sub-blocks.
For every jo E [m], let us define function u(x) R.d
IR" to be:
u(0) = exp(Ajo x)
nx1
Definition 3.3. Suppose that there are two n. X size matrices A1, A2 Raxd, We define Ajo Raxd be Lean n x of size sub-block from A. (Recall that A = A1 2 A2 E xd.)
For every index jo E [n], we consider a function, (x(x) : RRdd - IR as:
alz);0 == (exp(A x. 1n
nxl
nxl
Definition 3.4. Suppose that a(I) E R is defined as in Definition 3.3. Recall u(x) E R is defined in Definition 3.2. For fixed jo [n], let us consider function fx) 12dd R [(x) a
u())jo
scalar
nxl
Let f(x) E Raxx denote the matrix where ju-th row is (f(x) T. Definition 3.5. For every E [dl, define h()io RR.
R as:
h(z) == A3 Yo
nxd dx1
Here let Y E Rddx denote the matrix representation of y E R. Let h(y) E R matrix where in
column is h(y)io
3.4 Loss Functions
In this section, we introduce some helpful definitions related to both E RRdd.
Definition 3.6. For every jo E [n]. we use [x) E R" to denote the normalized vector defined by Definition 3.4. For every io [d]. we let h(y) be defined Definition 3.5.
Consider every jo [n], every in E [d]. Let consider c( dd
dR + R as follows:
c(z) == (f(z) h(y) Ejo.io-
Here Ej is the (jo, it)-th coordinate/location of E E Raad for jo E [n], [d]. This is equivalent to
c(r) f(x) h(y)
E
axd
Definition 3.7. For every ju E [n], for every in E [d]. Let us define L(x) to be = 0.5c(x)² 30.
3.5
Tensor Trick
We state the well-known tensor-trick. It has been widely used in literature of linear algebra related to tensor computations [SWZ19, DSSWI8, DIS+ 19, SWY721, AS24c, GSX23, Zha22, RS722, GSY23b, DSY23, DGS23J.

---

Fact 3.8 (Tensor trick). For two matrices A1 and A2 E Baxd define A = A1 @ A2. Let X E R
Let E Rdd denote the vector representation of X. Then have vec( xec(AI,XAA.) Az.
Using the above tensor-trick, is easy to observe that
Fact 3.9. For two matrices A1 and Ae E Raxd, denote A A @ A2. Let gdxd Let Ajo E Raxdd a submatrix of A (by properly selecting n rows of A). Let x E Rd denote the vector representation of X. Then, have
(ro(exp(AIXA)) exp(A x)
(exp(AIXAA) )o:n)) CAp
Here Xx )in is the jo-th row o matrix exp(A X A
Proof. We can use the definition in fact and Fact 3.8, to prove it.
3.6
Reshape the Loss function via Tensor Trick
Lemma 3.10. Given the below requirements
Here are three matrices A1 E RRaxd A2 E Radd and A3 E Ruxd
Let
A 2 Ag E pnn xd? to be the Kronecker product of the two matrices A and Az.
For every jo E [n], define Aj E Raxd to be a x de sized block in the matrix
A Ru²xdd
E Raxd be matrix. Define E the (Jo, 40)-th coordinate/location of E Baad for every pair of jo E [n] and io E [d].
Here are two square matrices X gaxd let Y E Rddd
Let L(X) be defined Definition 1.2.
For every pair of [ni] to [d], recall that definition of L be found in in Definition 3.7.
Then, we have
L(X) Σ
Σ
L(x)jo.io
JoE[n] KoE[d]
Proof. We can show that
L(X) 0.5 D(X) exp(A,xXA)
As Y
E E
XXn
nXn
axd
dxd
nxd
0.5 (((exp(Ajo 1)) eep((Ajo I)), A3Y..io) Fojo.io)²
Jo=1 f0=1
0.5(((z))) M()):o) - }
Jo=1 10=1
Σ
L(x)
Jo=1 10=1
where the first step follows from definition, the second step follows from writing down the summation, the third step follows from definition of [xx (recall the Definition 3.4) and h(y)), (recall the Definition 3.5), and the last step follows from L(z) jo (see Definition 3.7).
D
7

---

f(0)
71
dc
A
Figure I: An example of oxx, y) " Adoag(ii (000)))))(()))fffffo)))
Proof Sketch for General Upper Bound
The most straightforward way compute the gradient would take dd) time in order to explicitly write down the matrix A. We first show how to obtain an intermediate algorithm, which runs in slightly improved time O nd²)) to compute the gradient. Our final algorithm will build this idea.
Lemma 4.1 (Warmup, attention gradient computation, informal version of Lemma C.8). the following conditions hold
Define four d size matrices E, A2, Ag and two d x d square matrices X, Y.Y to be input fixed matrices.
Let Rauxd and Y E Rdxdd denote matrix variables (we will compute gradient with respect to X
LLt dL(X) ax
Then the gradient g can be calculated in Onn²d + nd²) time.
The key idea behind Lemma 4.1 is to use algebraic manipulations to quickly compute the quantities defined in the previous section. We first compute f() in O(nd²) time, then show c(z) and ((x) can be computed in O d) time. Using these, we compute p(x) in (n²) time, then putting in all together, compute g in O(n nd?) time. (We refer the details to Section C.)
For notational simplicity, also write R to denote the vectorized version of X, and similarly y TR x1 for Y.
Next, we will show how to improve the running time of computing the gradient from quadratic time (2 n²) to almost linear time n I+()), We build on the approach Lemma 4.1 for computing the intermediate quantities f,c, q. and p. but speed up the time it takes implicitly, rather than explicitly, represent these quantities. We want to emphasize that, although our algorithm relies on careful manipulation of the input matrices, multiplication (which may otherwise be quite impractical).
main algorithmic result does not make of fast matrix
We now sketch the main algorithmic ideas. First, by linearity of derivative, we can show that
dL(x)
dl(x)jo.io
dr
dx
jo
Based on calculations we perform in Section B, Section C, and several linear algebra facts, we can show that
aL(x)
dx
c(x))).00
AT. (dia((((),,) h(y).
scalar
xm
For any fixed jo E [n), consider this quantity. Since this expression involves an n n matrix, the most straightforward way to calculate would take ((n²) time, and so summing over all jo [n]

---

would lead to a cubic-time algorithm. It is not too difficult to improve this: the n x n. matrix (see Figure for an illustration)
(f)g(f()))))
(z)
diagonal matrix a rank 1 matrix
is easily decomposed into a low-rank part f())) which has size n x n) and sparse part (diag((()))o) which also has size n X r),, which reduces the calculation of each part to only (()) time, and the total running time to ((n²) time.
However, we are aiming for a limesr-linear time algorithm, and it not possible to achieve this by treating the different Jo separately, since given jo must take Q2()) time to process. Instead, use tensor techniques related to low-rank approximations to simultanouesly compute all jo together and sum them in almost-lincal time.
To do that, we create several extra artificial or intermediate matrices g(x) E I (see Section C), p(z) E R=x= (see Section C). We will show the gradient can be finally constructed using a simple chaining technique (see Section D for more details), from f. c, q. Pi (handling diag (x))) similarly), P2 (handling Jo similarly), p (p Pt P22 to dL Intuitively, the chaining shows that de
low rank representation for f yields one for and these in turn yield one for q, and so on.
In particular, using ((x), we obtain that d((a)
dz be written
Jo diagonal matrix rank 1 matrix) column vector
(e)s))())sa))
[w)
4(z)30
which fact notably removes the summation step of in 1 to Using the notation of p(x), we finally yield that we need compute AT A2.) Thus long as p(z) has low-rank representation, will find that p(z) is the entry-wise product of two matrices with low-rank representations from prior work, which we can combine using column-wise Kronecker product to approximate p(x) itself.
then can solve the in +=1+ time (see Section D for more details). In particular,
Conclusion
Our results give complete fine-grained analysis of the rumning time needed to train LLMs. We show that there is a threshold depending on the parameter B, the magnitude of the parameter matrix entries. In settings where B is small, a near-linear-time algorithm for LLM training is possible by using our novel algorithm for backward computation. In settings where is large, not only does our algorithm not apply, but show it is impossible to design a nontrivially fast algorithm (barring breakthrough in satisfiability algorithms that would refute the popular SETH).
These insights can guide LLM designers to more efficient algorithms. When B can be made small, it would lead substantial savings in the computational resources needed for training and expression. When B must be large (perhaps achieve high expressiveness"), may as well use straigthforward algorithms and focus on other aspects algorithm speedup such as parallelization. The magnitude of B needed has been studied more recently (c.g., [AS24c]), and the need for fast training algorithms may further motivate this direction research.
lower bounds show that one
Limitations
Our main algorithm shows that the polynomial method can be used to quickly train LLMs with provable guarantees. While polynomial methods are frequently used for LLM operations in practice, they are typically simpler than the algorithms we present here. Implementing our algorithm in practical way would require substantial future engineering work which is beyond the scope of this paper. Our lower bound predicated on the Strong Exponential Time Hypothesis (SETH), a popular conjecture from fine-grained complexity theory. As with most results in complexity theory, results proved using a conjecture like this naturally come with associated limitations: the hard instances of SAT may not translate to the most important instances of LLM training, or the conjecture may

---

not even be true! That said, we wish to emphasize that SETH is the most popular conjecture in fine-grained complexity, used prove the optimality of many algorithms nearly every domain of computation, and decades of research in satisfiability algorithms have supported its veracity.
7 Broader Impact Statement
We give a new algorithm with provable guarantees for LLM training, which can help to guide future algorithm design in practice. This will help to develop the many positive broader impacts of LLMs. Since this is purely theoretical work, which addresses theoretical computational concerns for implementing known algorithms, we believe does not introduce any negative societal impact.
Acknowledgement
We would like to thank Yichuan Deng, Yingyu Liang, Han Liu, Jerry YC Hu, Zhenmei Shi, and Lichen Zhang for valuable discussions and anonymous reviewers for helpful comments. About how to upload this paper to arXiv (arxiv link https://arrat http:://axvvvo.///./ee /pdf http//ff22400/ 22402. 04497 [AS24b]). the readers watch video (this link bttps:///yutte bet..ee///etiti (MbHH in this youtube channel:
https:/ //tps////wwwww/// http:///ww..o.uuuu youtube. com/ @2zhaooongn03..
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 Dagssuhl-Leiinniz-Zenttee
für Informatik, 2022.
[ACSS20] Josh Alman, Timothy Chu, Aaron Schild, and Zhao Song. Algorithms and hardness for linear algebra geometric graphs. In 2020 IEEE 6lst Anmual Symposium Foundations of Computer Science (FOCS), pages 541-552. IEEE, 2020.
[ALS*23] Josh Alman, Jiehao Liang, Zhao Song, Ruizhe Zhang, and Danyang Zhuo Bypass ex- ponential time preprocessing: Fast neural network training via weight-data correlation preprocessing. In NeurlPS. arXiv preprint arXiv:2211 NVV:2211.11222 141227, 2023.
[Ano24] Anonymous. On statistical rates of conditional diffusion transformer: Approximation and estimation. Submitted The Thirteenth International Conference on Learning Representations, 2024. under review.
[AS23] Josh Alman and Zhao Song. Fast attention requires bounded entries. In NeurlPS, 2023.
[AS24a] Josh Alman and Zhao Song. Fast rope attention: Combining the polynomial method and fast fourier transform. manuscript, 2024.
[AS24b] Josh Alman and Zhao Song. The fine-grained complexity of gradient computation for training large language models. NeurlPS, 2024.
[AS24c] Josh Alman and Zhao Song. How to capture higher-order correlations? generalizing matrix softmax attention to kronecker computation. In ICLR, 2024.
[BCS97] Peter Bürgisser, Michael Clausen, and Mohammad A Shokrollahi. Algebraic complexity theory, volume 315. Springer Science & Business Media, 1997.
[BGVM20] Kunal Banerjee, Rishi Raj Gupta, Karthik Vyas, and Biswajit Mishra. Exploring alternatives to softmax function. arXiv preprint arXiv:2011.115:: 2020.
[BIS17] Arturs Backurs, Piotr Indyk, and Ludwig Schmidt. On the fine- ane-grained complexity of empirical risk minimization: Kernel methods and neural networks. Advances in Neural Information Processing Systems (NeurIPS), 30, 2017.
[Bā11]] Markus Bläser. Fast matrix multiplication. Theory of Computing, pages 1-60, 2013.
10

---

[BMR 20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[BPSW21] Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (over- parametrized) neural networks in near-linear time. 12th Innovations in Theoretical Computer Science Conference (ITCS), 2021.
[BS724] Jan van den Brand, Zhao Song, and Tianyi Zhou. Algorithm and hardness for dy- namic attention maintenance in large language models. In ICML. arXiv preprint arXiv:2304.02207, 2024.
[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.
[CLL++24] Bo Chen, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. Bypassing the exponential dependency: Looped transformers efficiently learn in-context by multi-step gradient descent. arXiv preprint arXiv:2410.11268 2024.
[CLP+21] Beidi Chen, Zichang Liu, Binghui Peng, Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Re. Ro.Mooo Christopher. A learnable Ish framework for efficient neural network training. International Conference Learning Representation, 2021.
[CLS+24] Bo Chen, Yingyu Liang, Zhizhou Sha, Zhenmei Shi, and Zhao Song. Hsr-enhanced sparse attention acceleration. arXiv preprint arXiv:2410. 10165, 2024.
[CND+22] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebas- tian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint wwXiv:2004.02311, 2022.
[DCLT18] Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1xXiv:8810004800 2018.
[DGS23] Yichuan Deng, Yeqi Gao, and Zhao Song. Solving tensor low cycle rank approximation. arXiv preprini 0rX45:2304.06594, 2023.
[DHS+22] Yichuan Deng, Hang Hu, Zhao Song, Omri Weinstein, and Danyang Zhuo Training overparametrized neural networks in sublinear time. arXiv preprint 045550 2022.
[DJS+19] Huaian Diao, Rajesh Jayaram, Zhao Song, Wen Sun, and David Woodruff. Optimal sketching for kronecker product regression and low rank approximation. Advances in neural information processing systems, 32, 2019.
[DLMS23] Yichuan Deng, Zhihang Li, Sridhar Mahadevan, and Zhao Song. Zero-th order algo- rithm for softmax attention optimization. arXiv preprint arXiv:2307.08352 2023.
[DMS23] Yichuan Deng, Sridhar Mahadevan, and Zhao Song. Randomized and deterministic attention sparsification algorithms for over-parameterized feature dimension. arXiv preprint arX1v:2304.04399, 2023.
[DSSW18] Huaian Diao, Zhao Song, Wen Sun, and David Woodruff. Sketching for kronecker product regression and p-splines. In International Conference on Artificial Intelligence and Statistics, pages 1299-1308. PMLR, 2018
[DSY23] Yichuan Deng, Zhao Song, and Junze Yin. Faster robust tensor power method for arbitrary order. arXiv preprint arXiv:2306.00 200000000 2023,

---

[FA23] Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned one-shot. In International Conference on Machine Learning, pages 10323- 10337. PMLR, 2023.
[GQSW24] Yeqi Gao, Lianke Qin, Zhao Song, and Yitan Wang. sublinear adversarial training algorithm. In ICLR. arXiv preprint arXiv:2208.05395. 2024.
[GSWY23] Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. A fast optimization view: Reformulating single layer attention in Ilm based on tensor and svm trick, and solving it in matrix multiplication time. arXiv preprint arXiv:2309.07418, 2023.
[GSX23] Yeqi Gao, Zhao Song, and Shenghao Xie. In-context learning for attention scheme: from single softmax regression to multiple softmax regression via tensor trick. arXiv preprint arXiv:2307.0244 02419, 2023.
[GSY+23a] Yeqi Gao, Zhao Song, Xin Yang, Ruizhe Zhang, and Yufa Zhou. Fast quantum algorithm for attention computation. arXiv preprint (xXXi:2307708045 2023.
[GSY23b] Yeqi Gao, Zhao Song, and Junze Yin. Gradientcoin: peer-to-peer decentralized large language models. arXiv preprint arXiv:2308. 10502, 2023.
[GSYZ24] Yeqi Gao, Zhao Song, Xin Yang, and Yufa Zhou. Differentially private attention computation. In Neurips Safe Generative Al Workshop 2024, 2024.
[HCL+24] Jerry Yao-Chich Hu, Pei-Hsuan Chang, Haozheng Luo, Hong- Yu Chen, Weijian Li, Wei-Po Wang, and Han Liu. Outier-efficient hopfield layers for large transformer- based models. In Forty-first International Conference Machine Learning (ICML), 2024.
[HCW+24] Jerry Yao-Chieh Hu, Bo- Yu Chen, Dennis Wu, Feng Ruan, and Han Liu. Nonparametric modern hopfield models. arXiv preprint arXiv:2404. 000:2240403900, 2024.
[3]K+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.0586) 2023.
[HLSL24] Jerry Yao-Chich Hu, Thomas Lin, Zhao Song, and Han Liu. On computational limits of modern hopfield models: A fine-grained complexity analysis. In Forty-first International Conference Machine Learning (ICML), 2024.
[HSK*24] 24] Jerry Yao-Chich Hu, Maojiang Su, En-Jui Kuo, Zhao Song, and Han Liu. Computa- tional limits of low-rank adaptation (lora) for transformer- based models. arXiv preprint arXiv:2406.00 0...0.3/36. 2024.
[HSW+21] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint (tti:://66.0966 2021.
[HWI.24a] Jerry Yao-Chieh Hu, Dennis Wu, and Han Liu. Provably optimal memory capacity for modern hopfield models: Transformer- compatilee dense associative memories spherical codes. In Thirty-eighth Conference Neural Information Processing Systems (NeurIPS), 2024.
[HWL+24b] Jerry Yao-Chich Hu, Weimin Wu, Zhuoru Li, Sophia Pi. Zhao Song, and Han Liu. On statistical rates and provably efficient criteria of latent diffusion transformers (dits). In Thirty-eighth Conference on Neural Information Processing Systems (NeurIPS), 2024.
[HYW+23] Jerry Yao-Chieh Hu, Donglin Yang, Dennis Wu, Chenwei Xu, Bo- Yu Chen, and Han Liu. On sparse modern hopfield model. In Thirry-sevnnth Conference on Neural Information Processing Systems (NeurIPS), 2023.
[TPO1] Russell Impagliazzo and Ramamohan Paturi. On the complexity of k-sat. Journal of Computer and System Sciences, 62(2):367-375, 2001.
12

---

[JZLD21] Yanrong Ji, Zhihan Zhou, Han Liu, and Ramana V Davuluri. Dnabert: pre-trained bidirectional encoder representations from transformers model for dna-language in genome. Bioinformatics 37(15):2112-212)) 2021.
[KKL20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient trans- former. arXiv preprint arXiv:2001.04451, 2020.
[KMZ23] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast trans- formers via sketches for polynomial kernels. arXiv preprini arXiv:2310.01655, 2023,
[KVPF20] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Trans- formers are mas: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-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.
[LL21] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2/01.00/99 2021.
[LLS*24a] Chenyang Li, Yingyu Liang, Zhenmei Shi, Zhao Song, and Tianyi Zhou. Fourier circuits in neural networks and transformers: A case study of modular arithmetic with multiple inputs. arXiv preprint arXiv:2402.09469, 2024.
[LLS+24b] Xiaoyu Li, Yingyu Liang. Zhenmei Shi, Zhao Song, and Yufa Zhou. Fine- line-griined attention i/o complexity: Comprehensive analysis for backward passes. arXiv preprint WxXiv:2410.093/ 2024.
[LLS 24c] Yingyu Liang, Heshan Liu, Zhenmei Shi, Zhao Song, Zhuoyan Xu, and Junze Yin. Conv-basis: A new paradigm for efficient attention inference and gradient computation in transformers. arXiv preprint arXiv:2405.05219 2024.
[LLS+24d] Yingyu Liang, Jiangxuan Long, Zhenmei Shi, Zhao Song, and Yufa Zhou Beyond linear approximations: novel pruning approach for attention matrix. arXiv preprint arXiv:2410.11261. 2024.
[LLSS24) Xiaoyu Li, Yingyu Liang. Zhenmei Shi, and Zhao Song. A tighter complexity analysis of sparsegpt. arXiv preprint arXiv:2408. 12151, 2024.
[LOG+19] Yinhan Liu, Myle Ott. Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 19X.v:1977.1192. 2019.
24) Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, and Yufa Zhou. Multi-layer transformers gradient be approximated in almost linear time. arXiv preprint arXiv:2408.1322 13233, 2024.
[LSSY24] Yingyu Liang, Zhenmei Shi, Zhao Song, and Chiwun Yang. Toward infinite-long prefix in transformer. arXiv preprint arXiv:2406.1400 14036, 2024.
[LSSZ24] Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Tensor attention train- ing: Provably efficient learning of higher-order transformers. arXiv preprint axXiv2440..66411, 2024.
[Man23] James Manyika. An overview of bard: an early experiment with generative ai. Technical report, Tech. rep., Technical report, Google AI, 2023.
IMGN+23] 23] Sadhika Malladi, Tianyu Gao, Eshaan Nichani. Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. arXiv preprint arXiv:230S. 0.:005/7733. 77333, 2023.
[PMXA23] Abhishek Panigrahi, Sadhika Malladi, Mengzhou Xia, and Sanjeev Arora. Trainable transformer in transformer. arXiv preprint 20X.::2307.0118, 2023.
13

---

[RSZ22] Aravind Reddy, Zhao Song, and Lichen Zhang. Dynamic tensor product regression. In NeurlPS, 2022.
[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.
[SLBK23] Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306. 1695, 2023.
[SMN+24] Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, and Shafiq Joty. Discover- ing the gems in early layers: Accelerating long-context Ilms with 1000x input token reduction. arXiv preprint urXiv:2409. 17422, 2024.
[SWYZ21] Zhao Song, David Woodruff, Zheng Yu, and Lichen Zhang. Fast sketching of polyno- mial kernels of polynomial degree. In International Conference on Machine Learning. pages 9812 9823. PMLR, 2021.
[SWZ19] Zhao Song, David P Woodruff, and Peilin Zhong. Relative error tensor low rank approximation. In SODA. arXiv preprint D.XXv:1704.0820 2019.
[SYZ21] Zhao Song, Shuo Yang, and Ruizhe Zhang. Does preprocessing help training over- parameterized neural networks? 35th Conference on Neural Information Processing Systems, 2021.
[SYZ24] Zhao Song. Junze Yin, and Lichen Zhang. Solving attention kernel regression problem via pre-conditioner. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 208-216. PMLR, 2024.
[SZZ24] Zhao Song. Lichen Zhang, and Ruizhe Zhang. Training multi-layer over-parametrized neural network in subquadratic time. In ITCS. arXiv preprint 2tiiii2112.076222 2024.
22] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.082399 2022.
[TLI+23] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aure- lien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397 2023.
[3MS+ Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yas- mine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhos- ale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09280 2023.
[VSP+17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[WCZ+23] Yilin Wang, Zeyuan Chen, Liangjun Zhong, Zheng Ding, Zhizhou Sha, and Zhuowen Tu. Dolfin: Diffusion layout transformers without autoencoder. arXiv preprint mxXi:23310.66005 2023.
[WHHL24] Dennis Wu, Jerry Yao-Chieh Hu, Teng- Yun Hsiao, and Han Liu. Uniform memory retrieval with larger capacity for modern hopfield models. In Forty-first International Conference on Machine Learning (ICML), 2024.
[WHL*24] Dennis Wu, Jerry Yao-Chich Hu, Weijian Li, Bo-Yu Chen, and Han Liu. STanhop: Sparse tandem hopfield model for memory-enhanced time series prediction. In The Twelfth International Conference Learning Representations (ICLR), 2024.
[Wil18] Virginia Vassilevska Williams. On some fine-grained questions in algorithms and complexity. In Proceedings of the international congress of mathematiiion: Rio de janeiro 2018, pages 3447-3487. World Scientific, 2018.
14

---

[WMS* 24] Jiayu Wang, Yifei Ming, Zhenmei Shi, Vibhav Vineet, Xin Wang, Yixuan Li, and Neel Joshi. Is a picture worth thousand words? delving into spatial reasoning for vision language models. arXiv preprint arXiv:2406. 14852, 2024.
[WSD+23] Zirui Wang, Zhizhou Sha, Zheng Ding, Yilin Wang, and Zhuowen Tu. Tokencompose: Grounding diffusion with token-level supervision. arXiv preprint arXiv:2312.03626, 2023.
[WTB*22] 22] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprini wrXiv:2206.07663. 2022.
[XHH*24] Chenwei Xu, Yu-Chao Huang, Jerry Yao-Chich Hu, Weijian Li, Ammar Gilani, Hsi- Sheng Goan, and Han Liu. Bishop: Bi-directional cellular learning for tabular data with generalized sparse modern hopfield model. In Forty-fire International Conference on Machine Learning (ICML). 2024.
[XSL24] Zhuoyan Xu, Zhenmei Shi, and Yingyu Liang. Do large language models have compositional ability? an investigation into limitations and scalability. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
[YCRI22] Ann Yuan, Andy Coenen, Emily Reif, and Daphne Ippolito. Worderaft: story writing with large language models. In 27th International Conference on Intelligent User Interfaces, pages 841 2412 2022.
[YDY+19) 19] Zhilin Yang. Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
[ZBIW19] Ofrr Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Edition (EMC2-NIPS), pages 36-39. IEEE, 2019.
Computing-Neute
[ZBKR24] Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. In ICLR, 2024.
[Zha22] Lichen Zhang. Speeding up optimizations via data structures: Faster search, sample and maintenance. Master's thesis, Carnegie Mellon University, 2022.
[ZHDK23] Amir Zandich, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In ICML, arXiv preprint 2XXv:2332.002511, 2023.
[3] Zhihan Zhou, Yanrong Ji, Weijian Li, Pratik Dutta, Ramana Davuluri, and Han Liu. Dnabert-2: Efficient foundation model and benchmark for multi-species genome. arXiv preprint arXiv:2306. 15006 2023.
[ZL23] Yuchen Zeng and Kangwook Lee. The expressive power of low-rank adaptation. arXiv preprint arXiv:2310. 10..:523.../5/... 2023.
[2RG+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 arXiv:2205.0100 2022.
ZWW**44 Zhihan Zhou, Winmin Wu, Harrison Ho, Jiayi Wang, Lizhen Shi, Ramana V Davuluri, Zhong Wang, and Han Liu. Dnabert-s: Learning species- aware dna embedding with genome foundation models. arXiv preprint arXiv: 08777, 2024.
15

---

Appendix
Roadmap.
In Section A, we provide basic notation and facts. In Section B, we provide details about gradient computations. In Section C, we explain the computation time for the gradient of attention loss. In Section D, we show how to further improve the gradient computation from quadratic time to almost linear time. In Section E, provide our main lower bound result.
A Preliminaries
In Section A.I, we define some basic notation. In Section A.2, we state several facts which we will use.
A.1
Notation
For any positive integer n, we define [r]] == 21,
For two length vector and y.
(ar, y) to denote the inner product between and y, i.e.,
all ones vector. It is hard to denote vector that i-th entry is r:9i. Let 1, denote the length-n
(x, 9) = Ei=1 T. We use
that (x y, 1n.) (I, y).
For vector we use u to denote the transpose of u. For matrix M, we use M to denote the transpose of matrix M.
For vector u. we use exp(u) to denote the vector that i-th coordinate is exp(u:). For a matrix A, we use exp(A) to denote the matrix that (i. coordinate is exp( (A.,)).
We define the Kronecker product between matrices X and Y, denoted X @ E 10 [ as (X & +h [m.].
Xmomi
is equal to jo 2i where E in E [mo] E [m,],
For each positive integers my M2, ma, use Taat (m,, m22 m3) denote the time of multiplying
X m2 matrix with another m2 X my matrix.
A.2 Basic Facts
Fact A.1. Let R, R.,
Then we have
(z) x diag(y)2.
(x,1) = (x y. 1n).
Fact A.2 (Folklore). Let U1, V1 E Ranki Let U2, V22 E Raxks, Then we have (UiViT) (U2V2 (U, (2)
Here, given U1 E Raxki and U2 E Ranka the Un Ø U2 E Raxkika the row-wise Kronecker product, i.e., (U1 2 U2) 1)k; = (U1),, Un,, for all i E [n], I E [ki] and l2 E [zl
A.3 Matrix Multiplication
We define matrix multiplication notation and state some well-know facts here. Definition A.3. Let nh, n22, n3, denote any three positive integers. We use Tmae n22, ns) to denote the time of multiplying an n2 mairix with another n2 rig-
The straightgforward algorithm following the definition of matrix multiplication gives that Tmaa (n1, T22, n.y) < Onningns) In fact, we will only use this straightforward bound in all our algorithms in this paper, and avoid needing any other (potentially impractical) matrix multiplica- tion algorithms. Nonetheless, we will emphasize the appearances of Tmat in our algorithms below, since then any fast algorithms systems for performing matrix multiplication could be used speed up these steps our approach.
It is well-knows that
16

---

Fact A.4 ([BC997, Blä13]), Let n1, R2, (TTaa (2) n3, n2)) (T (12, n1, Rz)) ma n3, n1)) O(Tmat (1z3 R1, (g])
Rz. denote any three positive integers. Tiae (11, n2: nz.)
m(Tmat (13 R2, ).
B More Details about Gradient Computation
In this section, we provide details and calculations to assist with gradient and derivative computations. We remark that, in this section, for convenience of computing a closed form for the gradient, we ignore the 1/d factor in function f. Since is only a rescaling factor, won't affect how we compute these matrices in general.
Lemma B.1 (The gradient computation for several different functions with respect to Iz). For every i E [d_], define Ajous R be the i-th column for Ajo, E Faxd u(z)., IR", The scalar function a E IR, column finction f( E IR", scalar function c(r) E IR and scalar function defined as Definitions 3.2, 3.3, 3.4, 3.6 and 3.7 respectively.
jo
Then, for each E [dd), we have
Part 1.
dr
64 dr;
Part 2. For each jo E (n].
d j (Ajo)i
dx.
Part 3. For each jo [m]
du(x) dx;
Ajo.s ou(
Part 4. For each jo [n].
do((x)) dx;
(Ajo.s, u()))o)
Part 5. For each jo [m].
df(x);
A of())) (Asoii f())s)) f(x),
dx;
Part 6. For each jo E [n]. for each in [d].
aly))) (f(v)io Agan of(x)s)) - (M(y)ior (o f()) (Ajoid f(x)
dæ;
Part 7. For each jo E [n]. for every in [d]
dc(x) = (Ajo.t of (2) h(y)oo) (f h(v)o)) (Ajois (x)
dæ;
Part 8. For each jo E [n). for each io [d]
dl(x)jooo dz,
((h(y) A )z())))) (f(x) Ajon)) (N(y)io f(x) c(1)
Proof. Proof of Part 1. We have
dr da.
E;
17

---

Proof of Part 2. We have
d jo
dr
A
dz;
dx:
xxd?
x1
Ajo
Cj
axd
d3 x1
Ajo
Proof of Part 3.
can show
du(x),o
dexp(Ajo x)
dx;
dr; = exp(Ajo x)
dA x dx. exp(Ajo x.) Ajo.i = (1) Ajo.i
where the 3rd step follows from Part 2, the last step follows from definition of u(x) jo-
Proof of Part 4.
For simplicity of writing proofs, we use (+) to denote (x)).
We can show
da(:) dx;
d(u(-)so In)
dæ;
(u(-);o Aju 1n)
=(u(-) Ajo.a)
where the 1st step follows from definition of a(-), the 2nd step follows from Part 3, the 3rd step follows from Fact A.1.
Proof of Part For simplicity of writing proofs, we use (-) to denote (I)). We can show that
df(-), dæ;
da(-) (a()))7((()))
dæ,
du(-),o
dox()) dx
a(. aa)) dr;
Ju(-)
For the first term, we have
du(())oo
al-) 50 dx;
= aa())=2l( u(()))o o
Ajo.
=(())0 Ajo.i
where the 1st step follows from Part 3, the 2nd step follows from definition of f(-).
For the second term, we have
dax(.) ja dr;
da(())o u(:)io )u()) - -a(:) dx
- -a(()) (u(-) Ajoaa) u(-)so
- --(()) (f())so: Ajo.a)
18

---

72
71
f())
f/
72-
digg L (()
Figure 2: An example of diag(/()))n)) - [x))
where the 1st step follows from basic calculus, the 2nd step follows from Part 4, the 3rd step follows from definition of f())jo-
Using all of the results above, holds that
df(-), s(:).o o Ajo -(()) (J(+)o: Ajo.))
dæ;
Proof of Part 6. It follows Part directly.
Proof of Part 7. For simplicity of writing proofs, we use (-) to denote (I)).
Following the definition of in Definition 3.6, holds that
c(-)io.fo = (f(()) M(U)) - F
(1)
Thus it holds that
a((f()io h(y)co) - Eso.50)
dæ;
dæ:
d(f{())o) h())io))
dr;
(f(())o o Aj h(y).n) - (f(-) h(y).o) (fo.))):AA..))
where the Ist step is because of Eq. (1), the 2nd step is from d da.
O. and the 3rd step is followed
by Part
Proof of Part& 8. For simplicity of writing proofs, we use (+) to denote (z). Following the definition of L(-) in Definition 3.7. it holds that
I(-)io.to 0.5c(.)3..
Thus, we have
dLL())oooo dæ
d(0.5c(-))
dæ;
dc(.)
c(-)
dr;
((f(-) Ajo.in k(y))o)) (f(-) A(y).)) (f(-) Ajoi)),
where the 1st step is followed by the Eq. (2), the 2nd step due to the chain rule, the last step followed by Part
19

---

C Time for Straightforward Computation
In Section C.I, show the calculation of (Similarly Section B, still ignore the 1/d factor here) and h. In Section C.2, we show the way we calculate c in straightforward way. In Section C.3 and Section C.A, we define two artificial functions p and g, and show how to compute them. In Section C.5, we provide the way to re-write the gradient in an elegant way. In Section C.6, we finally put these all together and find the running time of our algorithm.
C.1 Compute f and h
Lemma C.1 (Computing f and h). Suppose the following objects are given
Let (x) be defined as Definition 3.4
Let h(y) be defined as Definition 3.5
Then, we have
f(x) can be calculated in time of Tmaa (n, d, n)) + Tmae (n,, d, d)
. -(ly can be calculated in time of Tmat (11, d)
Proof. Note that
f(x) = D exx((AAAAA))
and
D = diag(exp( diag(exp(A1XA)) A 1
We firstly compute exp(ArxXA.) this takes time of Tmat (n, d, d) and Timet (n, d, n).
Then we can compute D, which takes O(?²) time.
Then compute D-1 exp(ppppp X (A)) this takes O(n²) time.
Thus, the overall time is
Tiat (n, d, d) + Timet (n, d, n) + ( O (n, d, d) + Tmat (n, d, n))
Note that h(v) = AgY which takes time of Tmat (n,, d, d).
Thus, the proof is completed.
C.2 Compute
Lemma C.2 (Computing c). Suppose the following objects are given
Bd
f(x) E Raxn is given
h(y) E Baxd is given,
Then compute c(x) R in Tanat (n, n.d) d) time.
Proof. Based on Definition of c(x) E RRad which is
c(x: = f(x)h(y) - E Computing (x))h(y) takes time of Tmae (n, d), and calculating ((r))(y) - E takes time of ((nd).
Thus, finally, overall time is
T Taat (n, d) (
20

---

C.3
Computation for q
We will define q, and then explain how calculate q. Definition C.3. Define c(x) E Ruxd in Definition 3.6. Define h(y) Bxxd in Definition 3.5. We define (2x) E R as
g(x) = c(xx)(())
nxd
dxn
Then we use qx=)} T to denote the Jo-th row of q(x) E R
Lemma C.4. it holds that
Suppose c(z) E Ruxd is given
Suppose h(y) Baxd is given
Then, we can compute ((x) in the time of ( (n, n, d))-
Proof. Recall that ((x:) = c<)))(y)) Thus it takes time of T max (n, d, n) O (n, n, d)).
C.4 Computation for p(z)
Let us firstly define and then we can show how to construct it. Definition C.5. For every index ju E [n], we define p(z)jo E RR" as
p(z) == (diag(f(*)))) f(x) f()) (()))
We define p(z) E Rax= in the sense that p the jo-th row of p(z). Lemma C.6. the below requirements holding that
Suppose f(x) E Raxx is given Suppose q(x) _R is given
Then, we can compute p(x) in O(n²) time.
Proof. Since (i)g(f(z)))) is diagonal matrix and [f is rank-one matrix, we know that p(x) E IR" can be computed in O(n), for each ju [n]. Thus we can construct matrix p(x) E Bexn in n x Onn) = O(n²) time in total.
C.5 Analyze the closed form of gradient
Lemma C.7. Define the functions f(x) E RR c(z) E Ruxd, h(y) Raxd, g(x) guxn and p(x) E Roxn in Definitions 3.4, 3.6, 3.5, C.3 and C.5 respectively, A,, A2 E Badd are two given matrices. We define A A₁ @ Az. Let L(x) be defined as Definition 1.2. Let L(x)io be defined as Definition 3.7. Then,
show that dL(x) vec( p()
Proof. From the Lemma statement, we have
AL(x,y)ion v)s..0 ((((x), Ajo h(w) (f(x) M()) (f(x) Ajor)) (3)
dæc
Note that by Fact A. holds that
(f(x), Aja h(y)a)) AJ.s
and
(f(I) v) (f(x) Apois) = AT. f1)) h(2)
21

---

Therefore, Eq. (3) becomes
dl(x))ooiio c( V) (A).. diag(f())), - Ao. jo (1)
dæ;
c(2, jo A.. (diag( -0))) )
(4)
where the 2nd step follows from simple algebra. Recall the way define jo (see Definition C.3).
( Σ (())
(5)
Recall that p(x) IR" is define Definition C.5,
p(x), (diag(f(z))a) (()))
(6)
holds that
dL(x)
dx
n
dL Su
dæ
=o=1 10=1
n
cl
Ag (diag((()))) fo) {x)
hlu))o
=o=1
scalar
dxxn
axn
nxl
Σ Al. (diag(f(z)))) - (1)))
Jo=1
n
Jo p(x) Σ
1
vec( A A(x)A2)
where the 1st step is because of Definition 1.2, the 2nd step is based on Eq. (4), the 3rd step is followed by Eq. (5), the 4th step is due Eq. (6), and the last step tensor-trick.
C.6 Putting together
Lemma C.8 (Attention gradient computation, formal version Lemma 4.1). If holds that
Define A,, A2, E Ra Define X, Y Rdxd to be several input fixed matrices.
. XeetXYY E Bddd denote matrix variables (we will compute gradient with respect to X
For easy of writing, we also use vector variables æ dd and y B i.e. vec(X)
Let dL(X) da E ddd (where L(X) is defined as Definition 1.2)
Then we can show that gradient g E & can be computed Tanat (n, d, n) + Tina (n, d, d) time.
Proof. Step 1. we compute f(z), hly). This takes O (n. n, d) + Taaat (n, d, d)) time due to Lemma C.I.
Step 2. we compute c(x). This takes time of ( (n, n, d) + Tmae (n, d, d)) due to Lemma C.2.
Step 3. compute g(œ). This take time of (n,, n, d)) due to Lemma C.4.
22

---

Step 4. we compute p(r). This take time of O(n²) due to Lemma C.6.
Step 5. using Lemma C.7, we know that gradient is equivalent to vec(A] p (2). Suppose A] E Rdxn p(x) R=** A2 Raxd are given, then it can be calculated in time of ((Tmat (n, n, d) Tmaa(n d, d)).
Thus, overall running for computing gradient
(TTas (n, d, d) + Tmat (n, d, n))
time.
D
Fast Running Time via Polynomial Method
Recall that in the previous section, for convenience of computing the derivative, we ignored the d factor in f. That factor d doesn't impact the running time of our algorithms since it just rescaling factor. To apply the tools from previous work [AS23], we will now reconsider the 1/d factor in f. In Section D.I. we will show how to efficiently and explicitly construct low rank representation for f. In Section D.2, show how to create low rank construction for c(x). In Section D.3, Section D.4 and Section D.5, we further give low rank presentations for ((x), Pr (x)) P2(x). In Section D.6, we prove final algorithmic result by putting everything together.
D.1 Low rank representation to
Using [AS23]'s polynomial method result, we are able to obtain the following low-rank representation result,
Lemma D.1 (Section of [AS23]). For any B = of (lVolog)) there exists a kı = o() such that: Let A1, E Raxd be two matrices and X E Rdxd be a square matrix. It holds that [|A] X| S B, 12 S B, then there are two matrices U1, V1 Baaks such that V f loo S c/ cn)oly(n). Here f(x) D- exp(A1 /d) and we define diag(exp(A, /d)ln)) Moreover, these matrices U1, V1 can be explicitly constructed +o( time.
D.2 Low rank representation to
Lemma D.2. Let = Ollg n). Assume that each number the n x d matrices and h(y) can be written using (log n) bits. Let n x d matrix c(x:) be defined Definition 3.6. Then, there two matrices U1, V12 E Baxks we have [|U V, h(y)) E - -xx) S e/ poly(n).
Proof. We can show that
| hly) -5-o(al c(x)llo = [|11VV h(y))(())))))) - E f(x)h(v) + Eloo
= =I(UuvT f(x)))(uv)lo
ce/poly(n)
where the first step follows from c(x) = ((x)((v) E.
D.3 Low rank representation to
Lemma D.3. Let K2 = (1), Define c(x) E Raxxd to be as Definition 3.6. Define h(y) E Raad be as in Definition 3.5. Assume that q(x) h(y)c(()T E Raxn There are two matrices U2, V/2 E guxka such that 1|U2V2 ((x)))l S e/ poly(n). The matrices U2, V2 can be explicitly constructed n1to(11) time.
Proof. We define ((x) to be the approximation of g(x).
From Lemma D.2, we know that V, h(y) E is good approximation c(x).
Then should pick in this way ((x) (( h(y) T.
23

---

Now, let us turn (() into some low-rank representation
(()) = h(v)
h ET
nxd dxn
n.kky k in Xn
nxd dxn
is obvious that we should can first compute h(y) VV which only takes n + time. Then since all the low rank matrices are known, then explicitly construct U2. V/2 E Baxk, where l2 max k} + d no(1)
For controlling the error, we can show
[((x)) [xx) (MTT)(UIVTTTO)) E)T -1())((2) 11o0
< d. [|b())lbo (U E c(x))
S e/ poly(()
Thus, we complete the proof.
D.4 Low rank representation to p1 (1
Lemma D.4. Let kı = no()), Let K2 = n1))) Assume that (i()) = f(x) q(x)) Assume
U1, V1 Raxki approximates the f(x) such that [|U,V, f(x) llo I c/ e/poly(n). Assume U2, V/2 E Raxks approximates the q(x) E Ruxx such that [12V22 q(x)]l sc/ poly(n). Then there matrices U3, V3 Raxks such that [|UgV3 -x)(x) S e/ poly(n). The matrices U3. Vs can be explicitly constructed in n +o(1) time.
Proof. We choose U3 U1 2 U2 and V/3 V1 @ V22 This can be computed in n 1+o(() time.
For easy of writing proofs, call f(x) U1 VT and ((x) U2V.T
Using Fact A.2, we know that
|Ug Vg (x)]lx 2 |[UaV3 - ((xx o u(())|>
-I(U1 o UV) @ V2) f(x) (()))
= (
-If(()) ((z)))))())))))))) f(x) q(
= II/(x) 7(x) - f(x) o q(z) [()) f()))
s +lff(z) q(z) f(x) q())) (1((z) o q(z) f(x) o ((2)
(o)ooy(()
where the 1st step follows from the way we define pi (I), the 2nd step follows from the way we define U3 and Vg, the 3rd step follows from Fact A.2, the 4th step follows from the way we define ((x and 7(x), the 5th step follows from simple algebra, the 6th step follows by triangle inequality, and the last step follows by that entries are bounded and [[((x)) -f(x) S e/ poly(n) (Lemma assumption) and []](x) q(=))loc S poly)(n) (Lemma assumption)
D.5 Low rank representation (2
Lemma D.5. Let K (11) Let K2 (1)) Let Ks n((1) Assume that P2(x) is where jo-th column p2(x)io = (2) for each jo E [n]. Assume U1, Vi E guxki approximates the f(z) such that ||U,, f(x) I e/poly(n). Assume U2, V2 e gaxk, approximates the q(I) E Bax" such that [U2VV22 =x(x)]loo s/ <oly()n)). Then there matrices U4, V4 E Banks such that [[AVV p2(x)]l e/ poly(1). The matrices Us, Va can be explicitly constructed in n. [+o(1) time.
24

---

Proof. We define local vector function r(x) E IR" where r(z),o is f(x))) Let F(x) denote the approximation of r(z).
Note that (o. is good approximation to f(r)
Note that (o. is good approximation to g2))
Let F(z))o = (0 7(1)))) = Vi) (U2V2) J0.-
For the computation side, we firstly compute V/ This takes n (+((1) time.
Next, have
F(x)
(U212)
= iU ViVV (U2)
1xk1
in x
12x1
Once the V2 are pre-computed, the above step only takes Okk time. Since there coordinates, so the overall time is still O(nkk = n +
Let f(x) V denote the approximation of fa:) Then we just use f(x) and F(x) approximate (2())) in the following sense, (z) [(x) diaag((())) Since f(x) has low rank representation, and diag(F(x:) is diagonal matrix, then it is obvious how to construct U4 and V4. Basically U4 = Ui
and 112 = diag(Fia (F
Now, we need to control the error, we have
[ lx = (122(x)) - (2
(f(=))) f(z)))r(z)))))
JoEinl
[[(r)) i(x) (()))o(()) + ) f(z)
max JoE(m]
{())) f()) acc 11((2)) (
max joE(m]
where the 2nd step follows follows from definition of pa (x) and (2 For the first term, we have
F(z)),) - F())),(=)))lll max [1 (F(x);, r(x) JoE[m] sc/ (ol
max JuEin]
For the second term, we have
max (f(())))(z)sn)
JoE[n| I max [lf()) fx)) |r(x),a
JoEin]
S (o)y((())
Using the three equations we obtained above, the proof is completed.
D.6 Fast Computation in Almost Linear Time
Theorem D.6 (Main result, formal version of Theorem 1.6). Assuming the entries of A1, A2, X, Ag, Y, E reprexented using Oglog n) bits, there is a n ++(1)) time algorithm solve AAttLGC(nn n, d = O(logn), B = of o()logn) (see Definition 1.4) to 1 //pol(()) accuracy. In particular, our algorithm outputs gradient vector R&8 such that d gl S 1/ poly(n).
n matrices p(x) (Definition C.5), P1 (x) (see Lemma D.5) and p2 (1) (Lemma D.4), it is straightforward that
Proof. Recall definition
(x) P1 (z) -()))
Using Lemma D.I, Lemma D.2, Lemma D.3, we know that assumptions Lemma D.4 and Lemma D.5 are holding, so that we can use Lemma D.4 and Lemma D.5 to obtain that
25

---

Pi (z) has approximate low rank representation U3, Vs. let (x) denote U3V3
P2(z) has approximate low rank representation U4, V4. let F2(2) denote U
All of the Lemmas D.1, D.2, D.3, D.4 and D.5 are taking n (o(+) time.
According to the proof for the Lemma C.7, have that
L(X)
vee (oc(AA]i p(()))
dx
Thus, we firstly compute A U333 A2,
We compute A] U3 E gdxka, this takes n ++0(( time
We compute Vig A2 E Raxd this takes n [+((() time Compute Ug) (V3) A2), this takes d (2) time
Second, we can compute A U A2,
We compute A U4 E Rdxks, this takes n (++0(1) time
We compute V2 A2 R& this takes n ++((1) time
Compute (A (VT A2). this takes d (1) time
So, overall running time still n ++o(1)
We have
dL(X)
~l1x
|| veec - vec(A] ((x) ((x)
dæ
||App(x)A2 A A2)1
(pI - -[A{((((()))_(()))) -22()))
-
S (pT(((())) ( -=) + ||4 (P2(x) - A2(e))
||AlllollA2lloo n² (x) (x))
- -
$ </ poly(n)
where the 4th step follows from triangle inequality, the last step follows from entries in A1, A2 bounded, and || (x) - Pi (x) c/ poly(n), [l)2(x) - P2(x))loc I </ poly(n).
Picking 1/ (n)y we have the proof completed.
E General Lower Bound
We will critically make use of the known hardness result for attention computation itself, which we state now.
Definition E.1 (Attention Computation). Given input matrices Q,K, V e Raxd and parameter e > 0, compute a matrix T e =x satisfying
TT- D----1 D AVIL SE
where A exp(QKT) and = diag(A1,))
Lemma E.2 (Lemma 4.7 in [AS23]). Assuming SETH, there no algorithm running in time 0(22~))fe for any constant > that solves Attention Computation (Definition E.1), even when the inputs satisfy the following constraints, for any parameter k. 2
d= O(log n))
26

---

Ve {0, 1
There is a value B S Olog" n (1+ R)) such that every entry of QK" is in the interval (0, B) and at least half the entries each row of QK equal to B,
moreover |01 |o I OlyVogg(1++)) and
<
Next, we show that the attention optimization problem behaves particularly well when given matrices constrained as in Lemma E.2:
Lemma E.3. Let be a fixed n. x n matrix whose entries are real numbers the interval [o, B). and such that in each row of A, at least half the entries are equal to B. Let V be any d matrix whose entries are all in 1.. For X E R, define the n n matrix My,- = exp(AA), where exp is applied entry-wise. Define the function TR + IR. by
(M)1)) f(x) == [|diag(Mx1,) ( 'MVVIN 1|?
for all À E R we have
IF((N) S O(Bn).
='(N))=((B'r)
Proof. Let C denote the matrix C dlag(M)1,,)-1N)) For i,j [n], we calculate that Mx[i.j] cMiii and
cAA(F
Cii
Ek=1 CAAlE,
For E [d]. let S₁ C [n] be the set of is in column of V, i.e., =j [n] vb Hence,
for E [n] and & [d). the entry (i, of the matrix diag( (a0ng(M/1. 1,a MAV is given by
diag(M, - MAVii CV
E ci.bV
Σ Ciijll
JES.
Σ Djes, eAA(ta)
12 k== eAAlC
where the first step follows from definition, the second step follows from simple algebra. We thus get explicit expression for fA):
f(A) = Et= EEES, 2
k= edAli.)))
(=1
JEES, Eies,, A(AKA))AAAA))
Σk1=1 Eky=1 A(ARAA))AAA
1=1
We define
n
a(A,t) = Σ Σ Σ MARRA)AARA
ES
27

---

and then we define
b(A, = Σ
k1=1A22
Combining the above three equations, we can obtain
f(A) Σ ((A.i)/b(A.0).
Since, for each row of A, at least half the entries equal B, and all the entries are in the interval BJ, we can bound
n 2 2BA S ((A,)) S (n)² 2BA
(7)
Furthermore, since the derivative of c((Ali..aa+AAAi.kk)))) with respect to A is (Ali, + Ali, K2)) )+ we can bound
db(2. 22B-b(A.
2- b-b -
(8)
dd
We may similarly bound
0 0s a(), S n2. 2BA
(9)
and
2 a() da(),)) < . i).
(10)
dA
We can thus bound the derivative of f (where here, all the ' notation means derivative with respect to A):
a d'(,, b(A,)) - a(A, ((,,
((()) Σ
(b(2,)))
(2, b(A,
(b(2,)))
d'(2,)
b(),))
2B 2
(n/2)2 eBA
Σ 8B
8B n.
where the 1st step follows from definition. the 2nd step follows from simple algebra, the 3rd step follows from cancelling b(A, i). the 4th step is using Eq. (7) (for b(), i)) and Eq. (10) (for a'(A, i)). the 5th step follows from simple algebra, and the last step follows from simple algebra.
Similarly, we can provide a lower bound f*()),
F())= = Σ
a' d'(,,)) b(A,)) i) - (0(,)))((AA)))
(b(2,)))2
1=1
a(\,i)- ((),i)
(b(),
28

---

(n² (B (2B b(1,
((n/2)² (2BA). (b(2.
8B
1=1
-- n2
where the 1st step follows from definition, the 2nd step follows form simple algebra, the 3rd step follows Eq. (8) (for (X, i)) and Eq. (9) (for a(>, i)), the 4th step follows from simple algebra, and the last step follows from simple algbera.
Finally, letting f(A, a(A, 8)/b(A,)), i), we have again by the quotient rule that f*( is equal
d' (2, f(2, - 2- 2.,,( ()
b(2.
which we similarly bound in magnitude by O(B²n).
We recall a simple approximation from calculus:
Lemma E.4. Let [0, - R be a twice-differentiabl function such that If) S for all [0, 1]. For any positive integer m, define the sum
f(/m))
Im
m
Then,
I'm (f(1) f()) s b/m.
Proof. If two Aa A1 E [0, 1] have AAg - A1l & 1/m, then from our bound on (*(A), we know that If'()t) - f'(o)) b/m. We can thus bound the difference
- - Sr ('Aada
by
m-1
f(//m) + (b/m) tm + b/m
f(1) - (1)- =
Σ
m
1=
and
((i/m) (b/m)) lm b/mm
((1) ((0) >
Thus, we complete the proof.
Finally, we are ready for our main result:
Theorem F.5 (Formal version of Theorem 1.5). Let K. : -+ - by any function with k(n) (1)
and s(n) o(log r)). Assuming SETH, there is no algorithm running in time 0(r²~6) for any
constant 5 > 0 for Approximate Attention Loss Gradient Computation (Definition 1.4), even in the
case where d Ollog n) and the input matrices satisfy || Aglll As < O( OVllogn
k(n)), B = 0, Y = I, X XI for some scalar X E (0,1), and O ((1/(log))
Proof. Suppose there were such an algorithm. We call it ((logn))) times refute Lemma E.2 (with parameter K s(n)). Let Q, K, V be the input matrices to Lemma E.2, and set A Q, A₂ K, Ag V, Y I, and X XI for parameter À E [0, 1). Suppose the function f [0. - R is in Lemma E.3 where is the matrix A AIAJ. that M2 is the matrix exp (A, X A It follows from
Lemma E.3 that
If (A) < O(n log n (k(n))²)..
29

---

We can compute f(0) in (n) time since then My is the all- matrix, and our goal is to output f(1).
Thus, by Lemma E.4, it suffices to compute (A) on (1)g = ( n) points up to O(1/(log n) error, and return their average. But, since have picked X Xl, we can calculate (2) from the gradient dL(X) (from Definition 1.4), which is approximated by our assumed algorithm.
ax
30

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All the formal statements made the introduction, and summarized in the abstract, are proved in the appendix. We prove our algorithmic result in Section D, and prove our lower bound in Section E.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiaatinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
31

---

Answer: [Yes]
Justification: All our results are completely formally stated and proved. The main assumption we in our lower bound, SETH, is defined in Section 3 and we state that lower bound Theorem 1.5. Theorem 1.6 is our main algorithmic result, which we prove in Section D. Our warmup Lemma 4.1 is proved Section C. Theorem 1.5 our main lower bound, which we prove in Section E. All other Lemmas and Theorems are proved immediately after their statements.
it in our
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [NA]
Justification:
This paper is a purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
32

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification:
This paper is a purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
public/guidn httee//iide//coeeuueeaeee Code@uumissiooooll oliicy for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurlPS code and data submission guidelines (https: wwwodeSumisssoonoo....
//ni http:ppullcpldpdeeepppppppp
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: This paper is a purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: This paper is a purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
33

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification: This paper is a purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttaar/////eupss//e cc/ ittp://iiripccccppliiii///////// t
Answer:
[Yes]
Justification:
We have the read the NeurlPS code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We have discussed that in Section 7.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
34

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper is a purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: This paper is purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwwode.hhoe...c. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
35

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper is purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moddl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper is a purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects. and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper is purely theoretical paper, and it doesn't 't include any experiments.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
36


Reviews:
[{"review_id": "JUj1VplhxG", "reviewer": "Reviewer_UJu3", "summary": "The authors study Approximate Attention Loss Gradient Computation (AALGC) (e.g. complexity of training LLMs with backprop):\n- they show that the magnitude of entries of certain matrices (layer inputs times QKV matrices) affect whether LLM backprop is \"almost-linear\" or not. This connects techniques like quantization to complexity.\n- if the max entry of the products is bounded by B and B is o(root(logn)) it's almost-linear time, and when it's more than that, it's impossible under some assumptions (SETH)\n- for the almost-linear case, they construct a novel algorithm for AALGC\n- SETH: the assumptions are the Strong Exponential Time Hypothesis: for any eps>0 exists k>=3 for which solving k-sat with n variables in O(2^{(1-eps)n}) is impossible (even with randomized algorithms)\n- part of the construction involves a polynomial approximation to softmax that works well for non-large entries.\n- the authors use some nice techniques (like the tensor trick) that not everyone between theory + ML is familiar with, so the techniques in the paper are useful in their own right\n- more generally, the style of the paper is nice in that the results are clear + it's very pedagogical", "questions": "- I can't tell if it's just notation or not, but at some points like Definition 1.4 it seems like \"n\" refers to (1) the sequence length (2) the number of bits in the B bound B = o(root(log(n)) , and also (3) d = O(log n) where d is the embedding dimension ? Could you elaborate on which of these are tied together?\n\n- You allude to the near-linear time algorithms possibly being tough and needing more algorithms/hardware innovations to make practical.  Could you elaborate on whether there are any particular things that come to mind?\n\n- AALGC seems to switch between an epsilon parameter and B parameter around Definition 1.4", "limitations": "Yes", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- the results, as mentioned in the introduction: extending results in complexity of forward computation to complexity of backprop\n- discussion of how complexity changes when certain assumptions are made (magnitude of entries) or approximations are made (polynomial softmax)\n- connection to practice: discussion of e.g. quantization practices in the context of theory", "weaknesses": "- the softmax approximation passes by a little too quickly in the paper, I had to scroll up and down throughout the paper several times to piece the definitions together. If you control+F (search) for \"softmax\" or \"polynomial\" in the paper, only a few results come up and it's mostly references. It would be really great to add a committed subsection that says \"This is the precise softmax approximation we use\" and later in the algorithm section point out where exactly it's used in a clear way.\n\n- could use a little more discussion on whether such bounds on B hold in practice (e.g. in some open source models)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review. We especially appreciate the suggestions in the ``weaknesses'' section, and we will make these suggested changes in the final version.\n\nTo answer your questions:\n\n$\\bullet$ You’re right, our main problem definition (Definition 1.4) defines these as separate parameters, and $n$ is always the sequence length. Our results discuss when fast algorithms are possible, and the answer depends on how these different parameters relate to each other, which is why we set some of the other parameters in terms of $n$. For instance, we show that if $B$ is bigger than $\\sqrt{\\log n}$ then no faster algorithm is possible, but if $B$ is smaller then we achieve an almost linear time algorithm.\n\n$\\bullet$ Actually, some research groups are currently working to implement these polynomial approximation approaches in practice. See, for example, the papers [MKZ23,HJK+23] we cite which discusses this in a lot more detail.\n\n$\\bullet$ In our paper, the parameter B gives a maximum norm for the input matrices, whereas epsilon is the allowed output error. E.g., in Theorem 1.6, “1/poly($n$) accuracy” means that the epsilon is any 1/poly($n$). We’ll try to clarify this more in the writing.\n\n\n[KMZ23] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. In ICML 2024.\n\n[HJK+23] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hajy2yGLrO", "reviewer": "Reviewer_mody", "summary": "This paper presents a comprehensive analysis of the fine-grained complexity involved in gradient computation for training large language models (LLMs). Building on previous work that characterized the complexity of forward computations in LLMs, this study extends the analysis to backward computations. The authors establish a threshold for gradient computation analogous to those found in forward computation. They develop a near-linear algorithm for scenarios where the bound is within the threshold and prove the impossibility of a subquadratic algorithm under the Strong Exponential Time Hypothesis (SETH) when the bound exceeds the threshold.", "questions": "1. How does the proposed algorithm perform when the bound exceeds the threshold?\n2. Based on your conclusions, is there a principle in practice for adjusting the strategy for forward and backward computations?\n3. What is the memory cost associated with the proposed algorithm?", "limitations": "As the authors discussed the limitations in Section 6.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The analysis provides theoretical evidence supporting recent successes in the efficient computation of transformers.\n2. The conclusions offer valuable insights for future algorithm design.\n3. The paper is well-organized and clearly presented.", "weaknesses": "While the theoretical contributions are strong, the paper lacks empirical validation of the conclusions and the proposed algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review. Indeed our focus is on the theoretical foundations of these LLM computational tasks, although we are optimistic that these results will help future empirical work as well. For example, the prior work on the theory of these problems which we extend here [AS23] has motivated a number of new fast algorithms for attention ([KMZ23,HJK+23,ZBKR24,YAH+24]).\n\nTo answer your questions:\n\n1. Indeed, when the bound exceeds the threshold, our algorithm may give incorrect answers, as it is not designed to handle larger inputs. However, our lower bound result shows that this phenomenon is necessary: it is not possible to design a fast algorithm beyond the threshold using _any_ algorithmic technique.\n\n2. One of the major insights behind our algorithm is that: we prove that the matrices in training/inference of the attention network can be expressed as the multiplication of two low-rank matrices. Such an idea has been widely used in fine-tuning throughout the study of LLMs (such as LoRA [HSW+21]). In a sense, our method gives a theoretical foundation for such a practical method. LoRA is mainly used in fine-tuning, but our method is for backward computation, indicating that this low-rank idea may be applied successfully to the training stage as well.\n\n3. The memory cost is also almost linear: the algorithm simply performs simple matrix-vector multiplications where the memory and time usage are both almost linear.\n\n\n[AS23] Josh Alman, and Zhao Song. Fast attention requires bounded entries. In NeurIPS 2023.\n\n[KMZ23] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. In ICML 2024.\n\n[HJK+23] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In ICLR 2024.\n\n[ZBKR24] Michael Zhang, Kush Bhatia, Hermann Kumbong and Christopher Re. The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry, In ICLR 2024.\n\n[YAH+24] Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, Liwei Wang. Do Efficient Transformers Really Save Computation? In ICML 2024.\n\n[HSW+21] Edward J Hu and Yelong Shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Lu Wang and Weizhu Chen. Lora: Low-rank adaptation of large language models.  In ICLR 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WhiH5HC0RQ", "reviewer": "Reviewer_C3gd", "summary": "This paper studies the complexity of gradient computation in training models with attention layers. The paper claims that their results support the same complexity as the forward computation for the backward/gradient computation. The results show the impact of the size of the entries of the attention metrics on LLM training.", "questions": "* Could the authors explain the key differences in technique between this paper and existing work that also explores the complexity of forward and backward/gradient computations?\n* Could we verify the SETH assumption for attention models?", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The complexity of gradient computation for attention layers is an interesting and important problem for LLM training. The approach to complexity analysis can provide useful insights to the community.\n* The complexity results are theoretically sound, supported by proof sketches.", "weaknesses": "* The clarity of the paper's writing could be improved. The paper could define the problem more clearly in the introduction and present the contributions clearly. It would be better to follow a flow that first introduces the computation problem, briefly describes the techniques and assumptions used to derive the complexity results, and then presents the complexity results. Currently, these aspects are intermixed.\n* The significance of the paper is unclear. The complexity of gradient computation aims to solve Approximate Attention Loss Gradient Computation, but the accuracy of this approximation compared to practical gradient computation is not studied. Additionally, the analysis requires the SETH assumption, the validity of which is not provided in the paper. Thus, it is unclear if the complexity presented in this paper accurately reflects practical cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review. We aimed to organise the paper by first introducing the problem we study (Definition 1.4 in section 1.1) then explaining our two main results and how they relate to each other (section 1.2), explaining the context of other work in this area (section 2), and then describing how we achieved these results (section 3 for standard tools we use, and then section 4 introduces our new idea). We chose this flow, of stating and explaining both the algorithmic and complexity results before getting into technical details, because we think that the contrast between the two results is one of the most important aspects of the paper that we want to get across: that having a bound on the entries is both necessary (from the complexity result) and sufficient (from the algorithm) for fast LLM training.\n\nTo answer your other questions:\n\n$\\bullet$ The existing work requires $n^2$ time to compute the gradient. Our new technique, of applying the polynomial/low-rank factorization idea to speedup the gradient computation, is novel, and has not been used before in gradient computation to our knowledge. Moreover, to the best of our knowledge, we’re the first to give an almost linear time ( $n^{1+o(1)}$ ) algorithm to compute the gradient with provable guarantees.\n\n$\\bullet$ SETH is the most popular conjecture in the area of fine-grained complexity, and forms the backbone of many results in that area, showing hardness of problems from ML, graph algorithms, string algorithms, similarity search, and more. The fact that it gives tight lower bounds for a wide variety of different algorithms problems, and that we don’t know how to beat SETH lower bounds for any of these problems, gives strong evidence for SETH. As we discuss in section 2 and section 3.2, SETH is actually a strengthening of the P != NP conjecture, and so verifying it would require, at the very least, first solving the P vs NP problem (which is far beyond the scope of this paper).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ziL2mHEIOs", "reviewer": "Reviewer_MncV", "summary": "This paper studies the fine-grained complexity of gradient computation for training attention networks. In particular, the paper shows that under SETH assumption, when the norm of certain matrices are bounded by $B$, the gradient computation lower bound is almost linear. an almost optimal algorithm is used in the theorem to provide the upper bound which almost matches the lower bound.", "questions": "First of all, I am not very familiar with the strong exponential time hypothesis (SETH) or the related literature. I evaluate the work on LLM and numerical analysis perspectives, which might not be adequate. In general I think this is an interesting work with potentials of guiding the practices, yet it seems not fully explored. I have the following questions for the authors\n\n1. I think the paper is somehow over-claimed for \"training LLMs\" since the authors only consider one layer of attention network. How is the gradient computation error behave if we use backpropogation with multi-layers? There should be a remark or corollary for it.\n\n2. Can the authors discuss the implication of the algorithm used for construction the upepr bound in Theorem 1.6 for real applications? For example, since when $B$ is small we have a theoretical efficient algorithm, shall we adapt certain strategy to keep the norm smaller than certain constant? Also shall we follow the theoretically-efficient algorithm when computing the gradient in practice?\n\n3. Continue the above point, is there any numerical evidence to show that having a small norm or the algorithm used in constructing the upper bound have advantage in practice? This may greatly strengthen the point made in this work.", "limitations": "The authors are pretty clear about their limitations.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is theoretically well-rounded. The work is among the first few to study the computational complexity lower bound for gradient of attention networks.", "weaknesses": "(Pleae reply to the Questions section directly) The work is pure theoretical and lack of experiment verifications; The work is dedicated for one layer attention networks and may arguably not for \"large\" language models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review. Indeed our focus is on the theoretical foundations of these LLM computational tasks, although we are optimistic that these results will help future empirical work as well. For example, the prior work on the theory of these problems which we extend here [AS23] has motivated a number of new fast algorithms for attention ([KMZ23,HJK+23,ZBKR24,YAH+24]). \n\nTo answer your questions:\n\n$\\bullet$  Our algorithms and lower bounds can both apply separately to each attention head in a  larger LLM. In other words, in training, you can separately apply the algorithm for each attention head and attention layer, still giving a fast algorithm for any size model. We will expand on this in the final version.\n\n$\\bullet$ Indeed, a main takeaway of our results is that one should focus on maintaining parameters with smaller norms throughout the training process in order to achieve faster algorithms. This mirrors the phenomenon observed in practice that quantization or other parameter approximations often lead to faster algorithms [AS23]. And indeed, our approach could be used for gradient computations in practice, although as we discuss in section 6, a practical implementation would likely need more algorithms engineering work.\n\n$\\bullet$ Indeed, prior work on LLM implementations has observed a similar phenomenon, that algorithmic techniques like quantization [ZBIW19] and low-degree polynomial approximation [KVPF20], which require bounded or low-precision entries, can substantially speed up LLM operations. We discuss this in section 1 (lines 66-72) of our paper, and the introductions of those papers discuss the phenomenon in more detail.\n\n\n\n\n\n\n[AS23] Josh Alman, and Zhao Song. Fast attention requires bounded entries. In NeurIPS 2023.\n\n[KMZ23] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast trans- formers via sketches for polynomial kernels. In ICML 2024\n\n[HJK+23] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In ICLR 2024.\n\n[KVPF20] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Trans- formers are rnns: Fast autoregressive transformers with linear attention. In ICML 2020.\n\n[ZBIW19] Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing-NeurIPS Edition (EMC2-NIPS), pages 36–39. IEEE, 2019.\n\n[ZBKR24] Michael Zhang, Kush Bhatia, Hermann Kumbong and Christopher Re. The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry. In ICLR 2024.\n\n[YAH+24] Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, Liwei Wang. Do Efficient Transformers Really Save Computation? In ICML 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper under review delves into the complexities of Approximate Attention Loss Gradient Computation within the framework of training large language models. It presents a significant contribution by connecting the magnitude of certain matrix entries to the computational complexity of backpropagation, thereby linking quantization techniques to training efficiency. The authors establish a nearly linear time complexity for gradient computation when the maximum entry of the product of layer inputs and QKV matrices is bounded by a function that grows slower than the square root of the input size. Conversely, they argue that under the Strong Exponential Time Hypothesis, achieving subquadratic time complexity is unlikely if this bound is exceeded. Reviewers' concerns are properly addressed. Overall, the paper provides a solid theoretical foundation and meets the acceptance criteria.
Decision: Accept (poster)