LevAttention: Time, Space and Streaming Efficient Algorithm for Heavy Attentions
Authors: Ravindran Kannan
Chiranjib Bhattacharyya
Praneeth Kacham
David Woodruff
Venue: ICLR 2025 Poster (2025)
Keywords: transformers
attention
randomized linear algebra
leverage scores
Lewis weights
Primary area: foundation or frontier models, including LLMs
TL;DR: We show how to find all the ``large attention scores" of a matrix $A$ w.r.t. a function $f$, i.e., entries of $A$ which are at least a positive value $\epsilon$, in time $n$ poly$(d/\epsilon)$ for a positive constant $\epsilon > 0$

Abstract: A central problem related to transformers can be stated as follows: given two $n \times d$ matrices $Q$ and $K$, and a non-negative function $f$, define the matrix $A$ as follows: (1) apply the function $f$ to each entry of the $n \times n$ matrix $Q K^T$, and then (2) normalize each of the row sums of $A$ to be equal to $1$. The matrix $A$ can be computed in $O(n^2 d)$ time assuming $f$ can be applied to a number in constant time, but the quadratic dependence on $n$ is prohibitive in applications where it corresponds to long context lengths. For a large class of functions $f$, we show how to find all the "large attention scores", i.e., entries of $A$ which are at least a positive value $\varepsilon$, in time with linear dependence on $n$ (i.e., $n \cdot \textrm{poly}(d/\varepsilon)$) for a positive parameter $\varepsilon > 0$. Our class of functions include all functions $f$ of the form $f(x) = |x|^p$, as explored recently in transformer models. Using recently developed tools from randomized numerical linear algebra, we prove that for any $K$, there is a "universal set" $U \subset [n]$ of size independent of $n$, such that for any $Q$ and any row $i$, the large attention scores $A_{i,j}$ in row $i$ of $A$ all have $j \in U$. We also find $U$ in $n \cdot \textrm{poly}(d/\varepsilon)$ time. Notably, we 
(1) make no assumptions on the data, (2) our workspace does not grow with $n$, and (3) our algorithms can be computed in streaming and parallel settings. We empirically show the benefits of our scheme for vision transformers, showing how to train new models that use our universal set while training as well, showing that our model is able to consistently select "important keys'" during training. We also provide theoretical motivation by formulating a planted model in which our efficient algorithms provably identify relevant keys for 
each query.

Published as a conference paper at ICLR 2025
LEVATTENTION: TIME, SPACE AND STREAMING EFFI- CIENT ALGORITHM FOR HEAVY ATTENTIONS
Ravindran Kannan"
Chiranjib Bhattacharya
Praneeth Kacham Google Research
Simons Institute, UC Berkeley kannan1 gma 1. com
Indian Institute of Science chi ru@ i isc.ea in
pkacham@googles com
David Woodruff
Carnegie Mellon University dwoodrufecs cs cmu. edu
ABSTRACT
A central problem related to transformers can be stated as follows: given two n x d matrices Q and K, and non-negative function f, define the matrix A follows: (1) apply the function f to each entry of the n x n matrix QK", and then (2) normalize each of the sums of A to be equal to 1. The matrix A can be computed in (n)²d) time assuming I can be applied to a number in constant time, but the quadratic dependence n is prohibitive in applications where it corresponds to long context lengths. For a large class of functions f. we show how to find all the 'large attention scores", i.c., entries of which are least a positive value E, in time with linear dependence on (i.e., poly(d/e)) for a positive parameter ε > O. Our class of functions include all functions f of the explored recently in transformer models. Using recently developed tools from randomized numerical linear algebra, we prove that for any K. there is "universal set" U [n] of size independent of such that for any Q and any row i, the large attention scores A., in row of A all have j E U. We also find U in poly(d/z) time. Notably, we (1) make no assumptions data, (2) our workspace does not grow with n, and (3) our algorithms can be com- puted in streaming and parallel settings. We call the attention mechanism that uses only the subset of keys in the universal set as Lev identify the universal set U is based on leverage scores. We empirically show the benefits of our scheme for vision transformers, showing how to train new models that use our universal set while training as well, showing that our model is able to consistently select "important keys" during training. We also provide theoret- ical motivation by formulating planted model which our efficient algorithms provably identify relevant keys for each query.
form (xx |x]
the
since our algorithm to
INTRODUCTION
A transformer architecture is one of the most popular architectures for building foundation models, with applications to natural language processing, computer vision, and many other modalities and their combinations. is well-known that exact computation of their attention layers naïvely requires quadratic (in the context length) time, which poses a huge problem for scalability. A large body of work has tried to improve the efficiency of computing attention layers under variety of assump- tions, including imposing sparsity constraints (Parmar et al. 2018 Child et al. 2019: Beltagy al.) Kitaev et al. 2020 TTy et al.) 2020). kernel methods (Bello et al. 2021 Choromanski et al.,
2020
2021
Peng et al. 2021 Zheng et al. 2022). low rank assumptions (Wang et al.
2020 Xiong et al,
202
Ma et al.
2021). and assumptions of bounded entries or conditions on column row norms
Alman& & Song
2023 Han et al. 2024).
"R. Kannan is the first author and the remaining authors are ordered alphahetically.
Part of this work done while D. Woodruff was Google Research and at the Simons Institute for the Theory of Computing. D. Woodruff also acknowledges Office of Naval Research award number N0001421 12647.

---

Published as a conference paper at ICLR 2025
In each attention layer, one receives as input an x d' input matrix X, which may be the embedding of the tokenization of the imput, or an input from previous layers in the transformer. Here n is the context length and d' is the embedding dimension of each token, which is typically much smaller than n. From X we multiply by three d' x d learned matrices WQ, W and WV and define the query matrix
X Wº, the key matrix K X WK and the value matrix V X . WV One then outputs the attention matrix, which is defined to be D f(QK")- V, where D and f(QK") are each n matrices defined as follows. For non-negative function f. we apply entry-wise to the n X n. matrix QKT to form f(QK*). Then we let D be diagonal matrix with f(Q,, K2)), where Q: the i-th row of Q and K is the j-th row of K, respectively. Let A= D f(QK"). Note that the entries of each row of A are non-negative and sum to 1, and each row of A V can be viewed as non-negative combination of the rows of V. where the coefficients of the combination sum to 1.
D1. = 1
Instantiating the above framework with f(z) c corresponds to taking a softmax of each row of QQK"vvd and then multiplying by V. In this case, by appropriately scaling the hard instance in Alman & Song (2023), one can show that computing attention with high precision requires n²=o(1) time under standard complexity-theoretie assumption. In an attempt to bypass this hardness, recent work has considered replacing with other functions, such as f(x) = x² for positive even integer Indeed, both the PolySketehFormer (Kacham et al., (2024) well as the work of /Saróós et al. (2023) consider using TenssoSSkech to speed up the computation of an approximate attention matrix for such functions f. Motivated by these works, we define the f-sensitivities of of an x d matrix K for i = 1,,,..... n as follows:
5((K,, v)) V4O Σj= f((Kg, 9))
of sup
When fxx = x², these are just the so-called leverage scores of the rows of matrix K, which are well-studied in randomized numerical linear algebra (see, e.g., Mahoney et al. (2011); Woodruff ctt al. (2014). For general p. these are known as the L,- f.sseniitivities which are also well-studied (see, e.g., Woodruff & Yasuda (2023); |Padmanabhan et al. (2023), ), and they can be bounded by the so-called E,-Lewis weights of the matrix K (Cohen & Peng 2015) (for background, see, e.g., Section 3.3 Clarkson et al. (2019). Many interesting properties of such scores are known. One such property is the following. Let ff sup Σi= of (K). When f(x)
|x| P for 1 S p 2,
one has / < d, and when (xx = |r| for p 2 one has v < dP/2, These bounds do not depend on the context length n.
Critical to work will be the observation in practice that the matrix A is often well-approximated by retaining only its large entries, i.c., preserving all entries above a certain threshold s > 0 and replacing the remaining entries with 0, or perhaps fitting low rank approximation to the remaining entries (uupta et al., 2021 Wang et al. 2022). The entries of A are called the attention scores or attention weights, and we will say a score is large if its value is at least c.
1.1 OUR CONTRIBUTIONS
We outline our contributions below.
Existence of a Universal Set. We prove that for a large class of functions f. for any key matrix K, there a small "universal set" U C [m]= (1.,,...,,,) , of size independent of n, such that for any query matrix Q and any i E 21, the large attention scores Ai, in row i of A all have j E U. One of our results, which combines some ideas from Sections 2 [aand 3 is the following:
Theorem 1.1. Let f be non-negative functtin funntiooon and let of =upk Σ of (K). There is subset U C [n] of size V /e so that for any query matrix Q and i (1,2......
........... n), if Ar, 2 E, then j E U.
We note that for f(x:) = x², the f-sesiitivities are just the leverage scores of the matrix K, and it is known that So = d (see, e.g., Mahoney et al. 1): Woodruff et al. (0144)). In this case, by Theorem 1.1 we have |U| d/e. More generally,
f(z) R we have 4/ < d for 1 1 p S and &f S do/2 for p > 2 (see, e.g., Section 3.3 of Clarkson al. 2019)). and so [U] S max(d, //e Additional bounds are known for other functions f. e.g., from the work of
2

---

Published as a conference paper at ICLR 2025
Musco et al. (2022). There, for example, if f(x) is the Huber function f(z) x² for |a| z r and f(z) = ø for larger |zl, then = ('d or if f(æ) is the Tukey function f(z) = x² for
[2| S T, and f(x
for larger |z)]. then vi O(d [(dlogg) n). Here 7 > 0 is any specified threshold.
We give an outline of the proof of Theorem I.I here. If A is large attention score, then A.j = (3)
E and so of (K) 2 e. If we define U to be the set of i for which of (K) 2 e, then i E U and we have | E S N (K), and so [U] S T (K)/e.
(QQ 2 It follows that sup, E, Kerkk
We stress that our set U does not depend on any particular query or query matrix Q, i.e., for any possible future query q, any large attention scores it participates in necessarily involve keys in U.
Fast Algorithm for Finding U. There are very efficient algorithms for computing the set U for many interesting functions f. For example, when I [ar) = x²² these are the large leverage scores and simple way of computing them is by computing a QR- -R-dpccompoition of the matrix K and finding all rows with squared norm at least E. This takes (ndd) time and has the advantage of comput- ing the leverage scores exactly, resulting in the smallest possible set U. There are also sketching techniques for more quickly finding the large leverage scores (Drineas et al.) 2012 Clarkson & Woodruff 2013) in time nnz K) + poly(d/e) up logarithmic factors, where nnz denotes the number of non-zero entries of K. Similarly, for f(x)= = |x]| one can use the nnz () + poly(d/s) time algorithm for finding f,-Lewis weights Coheee & Pengl (2015). and the fact that a scaling of
these weights bounds the f-sensitivities (see, e.g., Section 3.3 of Clarkson et al. (2019)).
Time and Memory-E Efffiient Algorithm for Finding Large Attentions Given U. We do not make any assumptions on the data and our theorem holds for any query matrix Q, and thus given any potentially new query q. one only needs to search in the set U for the set of large attention scores involving q. Thus, instead of naively spending O(nd) time to walk through each possible key to find the large attention scores, one only needs to spend ////)) time, assuming that f can be evaluated in constant time. Notice that our time per query does not grow with the context length n. Further, as we only store the keys in the set U, our workspace also does not grow with the context length.
We also show how, for [((x)) = x² for even integers p it is possible to spend at most n poly(d) time preprocessing K so that from U and K, given a query q. one can output the exact value of all large attention scores involving q in poly(d/s) time. Note that this is not trivial simply computing (q, K2) for each K; E U, since we also would like to compute the exact normalization factor in the attention matrix, and do not want to spend n time to do so. For p 2 this amounts to computing || but by computing the SVD of K in a preprocessing step, this quantity can be computed in only (d²) time. For even integers p 2, we can reduce to the case of p by using Khatri-Ran products on the rows of K, as well as Khatri-Rao products of the query with itself. One can further speed up these computations by approximating the heavy attention values using sketching.
Extension to Streaming Environments. It is also possible to compute the set U' in a streaming environment. We illustrate the idea for fxx = x², though using the tensor trick discussed above, it also applies to f() = x" for any even integer p. The simplest algorithm is a two- two-pass algorithm over K, where we maintain the d x d matrix KJ K by summing n outer products as we read each of the keys. Then, in a second pass, the i-th leverage score is K (K" K)- K,, and we simply let U be the set of keys for which the leverage score is at least E. The total memory is only O(d² (d)/)).
More interestingly, one can compute U in single pass over K by again maintaining KT K but also storing all keys whose online leverage score [ohen et al. (0166a) is at least E. It is known that the online leverage scores are at least the leverage scores and sum to O(d log x. xO4), where OL is the online condition number, which can be bounded by n 0(4) assuming the entries of K have Ollog n) bits of precision. In this case, at the end of the stream, one simply evaluates KJ (KTK) -) K, for each key K, that was stored because it had large online leverage score. The total memory is poly(d/e) and the algorithm is a single pass algorithm.
Extension to Distributed Environments. It is possible to find U in distributed environment, where the keys are distributed across multiple servers. In this case if server i holds an n, x matrix ...., .,, Kd] if there are s servers, then the i-th server can communicate (Ki)" K to the first server for each i, and the first server can add these up to compute K K and send KT K to each server. Then the i-th server can send set U of keys Ky that holds for which
K*. so that K- IK'; K2;...
3

---

Published as a conference paper at ICLR 2025
K 2 E. Then U = U,Uن. Although our discussion has focused on f(z) similar ideas exist for f(x) xp for any p. well as the functions studied in Musco et al. (2022).
x²?
Planted Model. Finally, we formulate a planted model of keys and queries, where each query is noisy linear combination of small subset of keys. We show that in this model, we can find all keys relevant to a given query in sublinear time. We give deterministic conditions under which this is possible, and describe natural stochastic setting where these conditions are realized.
Experiments. We perform experiments on pretrained ViT models to empirically understand the structure of the attention matrices that arise for typical inputs to model. Our experiments suggest that for many attention heads, small subset of universal keys along with a set of local tokens capture large fraction of attention weight for large number of tokens. This motivates the identification problem that we study in this paper.
We then evaluate the effectiveness of leverage score selection for the downstream task of image classification using the pretrained softmax model. In this experiment, at each attention head, we compute a subset of keys with the largest leverage scores and make the queries attend only to that subset of keys. We observe that while the accuracy of the model drops compared to the full attention model, the model still maintains non-trivial accuracy and that leverage score selection obtains better accuracy compared to e.g., row norm based selection, showing that leverage score based selection does compute a subset of important keys.
We also trained multiple VIT models from scratch using the leverage score based attention mecha- nism and observe that the model quality improves significantly compared to doing inference on the softmax pretrained models using the leverage score mechanism. Across all the models, we observe that the model quality achieves >90% accuracy of the full softmax attention while selecting the top 32 keys (out of 197 keys for I/16 and S/16 models and out of 785 keys for the L/8 model) using the leverage score mechanism at each attention head. On the negative side, we observe that models trained from scratch using squared row norm selection or even "random key selection", where we select a set of random key positions for each attention head and each query attends to only keys in those positions throughout the training/inference process, attains similar accuracies. However, unlike leverage score sampling, these latter methods do not have the "universality" property as dis- cussed in the introduction for relaxed definition of the attention mechanism. We leave as an open question how to make effective use of leverage score information to train better models.
Notation: For matrix B, we use nnz BB) to denote the number of non-zero entries of B. We let w be the exponent of matrix multiplication, so that multiplying two d x d matrices takes d time.
2 CONNECTION TO LEVERAGE SCORES
We define the Generalized Attention Problem (GAP): let and D 2 d be positive integers. Con- sider two known maps: v,, Φb Rdd - D. We assume that & and si are computable in time O(D). We are given matrices Q and K, where Q is n x d and K is n d. GAP is the problem of computing (v(Q)). () the n x n matrix A defined by Aij = >7 (1.)) ((())) where K, denotes the j-th row of K.
Remark 2.1. One special case is the SoftMax operation. In this case the value D is infinite, but Softmax can be approximated with finite D, see, e.g., Choromanski et (007)) Note that ò = v) in approximations Softmax. Another set of interesting special cases occurs when & and si are polynomial maps, as in Sarlós et al. (2023); Kacham et al. (2024).
GAP requires M(n²) time to write down the entries of A. We show how to improve this by finding only the large entries of A, i.c., those that are at least e.
Theorem 2.1. There is deterministic algorithm A that finds a subset of rows of K (i.e., U S (m]) satisfying: [U] $ D and for all Quxa and for all i E {1,2, .n. {j: Aij 2 e) C.
We show that it suffices to prove this for the special case D = d and 6, 1 are the identity maps. Lemma 2.2. Without loss of generality, we can assume in Theorem I that = and o, 1 are both the identity map (ie., d(x) = 1b(x)) = x).
4

---

Published as a conference paper at ICLR 2025
Proof. The proof follows from kernel-type construction. From K, define an n x D matrix K given by K5 (K()). For any x d matrix Q, define an x D matrix Q by Qj (Q2)). Further, define an n n matrix A by: Aoj Σi= R2), is easy to see that À A. Assuming that Theorem
R0),R
2.1 holds when D d and Ø, s are both the identity map, can find Ü with {j: Aug 2 =) - Ü. Since À = A, we get that {j: A, 2 E) C Ü proving the lemma.
L
Restricting the case when D= d and x(x)) =b))== =I, I. we formally define the set LA for large attentions scores: LA(K,
8, v) = (j: (K), u)² 2 Ei= (K,, v)²} for u E RRd Let U(K,)) = UVER (0) LA(K,,, E, v). The following is simple characterization of (K),)).
Theorem 2.3. Let (K,,e) be the set of rows of K with leverage score 2 E, then (U(K, s d/c.
Proof. Observe that j E (kk,)) e) if and only if there is a non-zero vector v for which (K,,)) Σ (Ke.ofz 2
which means that the the j-th f-sensitivity of (K) 2 e. Conversely, if of (K) E, then there exists a non-zero vector v for which e) (K.)z 2 E, and so j E U(K,s). It is a well-known
fact that for f(x = x² that of (K) is precisely the j-th leverage score of the matrix K, see, e.g., the solution to problem 2.2 in Woodruff (2021). Note also that j E U(K,e)
T(K j) >
E + LS(K, j) 2 E. As the sum of leverage scores equal to d (see, e.g., Foster (9533), we
have (U( -
Thus, to compute LA, simply need to compute the rows of K with leverage score at least Theorem 2.1 now follows from Lemma 2.2 and Theorem 2.33
Streaming Algorithms. The j-th leverage score of K is equal to KÝ (K"K)- (see, e.g., Mahoney et al. (20TT); Woodruff et al. (2014), and so a simple 2-pass streaming algorithm using of n outer products in a first pass over the keys O(d)) memory would be maintain K"K K K, exactly in the second pass over keys Kj. and store of K. and then to compute KJ "5[(""KO those Ky for which this quantity is at least This simple 2-pass streaming algorithm uses (dd) words of memory and can be implemented in O(nd + dd) time.
One can obtain a 1-pass algorithm compute the rows of K with leverage score at least e with only slightly more memory. The idea is to store the set S of rows K, with online leverage score Cohen et al. 60) at least & in the first pass, and to also store KTK. The j-th online leverage score is equal to K. (( Kj. where Ki-1 denotes the prefix of the first j - rows of K. The j-th online leverage score is at least the j-th leverage score since sensitivities cannot increase as more rows are added K, but more interestingly. Theorem 2.2 of Cohen al. (2016a) shows that the sum of online leverage scores is bounded by Odd xe TT) where n. Rº is the online condition number, see Lemma 3.3 of Woodruff Yasuda (2022). It is well-known for matrices with integer entries bounded by poly(n), which is an assumption often used to obtain meaningful memory bounds in data stream, that the online condition number is at most n ((d) (see, e.g., Lemma 4.1 of Clarkson & Woodruff (2009) (0099) which lower bounds the minimum non-zero singular value an n x such matrix by n O(d)), and thus the number of rows stored will be at most O(d log n). and so the memory is bounded by (d log n) words. One can also maintain KT K in the stream as before. At the end of the stream, for each row K, stored which had online leverage score at least e, one can compute its exact leverage score KJ(KT K)- at the end of the stream in order to exactly find LA. Overall this is a 1-pass algorithm with memory poly(dlogn) n) words.
Finding all Heavy Attentions for a Query. Given query for f(z)= x² we can compute all large attention score values it participates in exactly because we can first preprocess K in (ndd) time to compute its singular value decomposition (SVD) K = U'EVT. Then, given query we can compute (q. K2) 2 for each K; E U in UU| d = (dL/e) time, and we can also compute the -|ΣV'' g in Odd)) time since ΣV" is d x d matrix. Thus, after an initial preprocessing of n poly(d/E) time, we can compute all large attention score values involving a query exactly and in only (oly(d/=) time. Note that if one would like to instead approximate the large attention values up to a 1+ factor. instead of computing the SVD of K, one can use a random sketching matrix S so that |KK]] (1 + Kqll If one uses the Subsampled Randomized
normalization ||Ka
5

---

Published as a conference paper at ICLR 2025
Hadamard Transform for example, then this improves the O(nd²) time required to compute the SVD of K to only nd poly n/E) time and incurring 1/poly(n) failure probability, e.g., using the analysis of Cohen et al. (2016b). is also possible to CountSketch in O(nd + poly(d)) time with a 1 1/pol((a) failure probability Clarkson & Woodruff (2013).
This section for f(x) x², while the next section handles f(x) |a| for any
1.
ATTENTION SCORES USING f(x) = |
We next consider f-sensitivities for f(x) = (ar|?.
Theorem 3.1. Let K be an n x d matrix. Let p [1, 00) and let f(z) |z|P. There exists a set S of rows of K containing all f-sensitivities at least with the following properties:
1. For 1 1 |S| O(d/e).
2. For p > [S| Odd////))
Moreover, S can be computed in (nnz(K) K) poly(d/e))ooly(ogg n) time.
Proof. Fix p 2 and let f(x) = |x]". Let T, denote the lp LewLss weights of K (Cohen & Peng 2015 These can all be computed up constant factor in total time (nnz(K) d" )poly(log time Cohen & Peng 2015).
For p < 2, it is known that Ti upper bounds the i-th f-sensitivity of For p > 2, it is known that dB/2-1r; upper bounds the i-th f-sensitivity. The Le-Lewis weights sum to at most d. For an exposition of these statements, see e.g., Section 3.3 of Clarkson et (2019) and references therein.
Therefore, for S p S 2, the set S of all rows with Ti 2 satisfies the desired properties. For p > 2, the set S of all rows with d2/2-17; 2 e satisfies the desired properties. By approximating the 7; up factor of 2 and including all i whose approximate т; is at least e/2 for 1 S p S 2, while will include all rows i whose f-sensitivity is at least &. Further, if the approximate т; is at least e/2, then the actual т; is at least e/ e44, and so we have |S] = Od/e) for 1 S p S 2 while |.|| = 0 (2P22/e) /e) for >> 2.
including all whose approximate т; is at least e/(20/2-1) for p 2
We next show how to efficiently approximate the normalization term for all 1 < p < co, in order to approximate all heavy attentions for a given query g by using the set U. We will later show how to compute the normalization term exactly for p an even integer.
Theorem 3.2. Let Q be an x d matrix and K be an d matrix. For < < co and any K.) can be estimated efficiently. Specifically,
j E {I....,n) m. the normalization term Σi= KQ3, there exists sampling and rescaling matrix S with the following properties:
1. For 1 I 2, S has d polylog(d/s)/e² columns,
2. For p > 2, has dP/2 polylog// columns.
3. With failure probability poly(dj simultaneously for all x E Rd, |or = (1 +
S can be computed in (nnz(K) + | T Sil can be computed in poly(d/e) time.
n) time, and after this one-time computation,
Proof. The normalization term can be written as ||Q,KT Ip. We can obtain the matrix by Lewis weight sampling on KT. This is a well-known technique for constructing subspace embeddings that preserve b norms simultaneously for all Cohen & Peng (201)). Specifically, by Woodruff & Yasudal (2022). there exists an algorithm that computes a sampling and rescaling matrix with the stated number of columns such that for all R2, I (1 + =)]l=k* The time to compute S is ((nz( K) + voly(d/e)) n). Once S is computed, we can efficiently estimate

---

Published as a conference paper at ICLR 2025
the normalization term for any j as [|Q2K* SIlz. Since S has poly(d/e) columns, computing this term takes poly(d/e) time.
Finding all Heavy Attentions Exactly for a Query q. Although the above procedure allows for quickly estimating the normalization term up (1 + =)-factor, one may be interested in the exact value of the attention score if the model is sensitive to slight perturbations. One can also do this with slightly larger time complexity for even integers p > 2, as the idea for p = 2 for computing heavy attention scores exactly extends to f(r) = xx for even integers p using tensor trick. Indeed, we can form the n x d2/2 matrix K' by taking the Khatri-Rao product of each row of K with itself p/2 times. Then we compute the SVD K' U'EVT. Given a query q. we can again compute (g. K; for each K2 E U in [U]
O (dd/2+1 /E) time, and now we can also compute the normalization Ej(Kj. = o'12 = |EVT, d'1 where q' is the Khatri-Rao product of q with itself p/2 times. Thus, for constant even integers p, after an initial preprocessing of n poly(d/e) time, we can compute all large attentions for any particular q exactly in only poly(d/e) time.
Remark 3.1. The choice of in the definition of f-sensitivity can significantly impact the identifi- cation of large attention scores. Consider the following examples:
Case P 2: Suppose one entry in row of the attention matrix is n n2/p, while all other entries are (1). This entry is a large attention with constant for f(x) = |x|", but it is not a large attention score for f(x x² unless E is inverse polynomial in
Case p < 2: Suppose one entry in a row is E, another entry 1, and the remaining entries are close to O. The entry with value e is large attention score for f(x) x²² with value E², but it is a large attention score for f(x) = |x|| with value E. Thus, using fs-sensitivities allows for faster identification and allows for storing smaller set of large sensitivity rows of K.
for Remark 3.2. In Musco et al. (2022). a method for efficiently computing f-sensitivities with respect assumed to satisfy basic properties such as subadditivity and symmetry. Examples include the Huber and Tukey loss functions. Musco et al.) (2022) show that the sum of these sensitivities is (d ni), implying the existence of a subspace embedding S with poly(d) columns to facilitate fast computation of the normalization factto Fur- thermore, the universal set U of rows of K now has size O(d log log n/c) and both U and can be found in ((nz(()) K) + poly(d) )poly(log n) time. These generalized loss functions may offer more flexibility and different efficiency versus accuracy tradeoffs.
a generic function is presented. This function
A PLANTED MODEL
We formulate planted model of keys and queries, where, each query is a noisy liner combination of a small set of keys; we call this set of keys, the "relevant keys for the query". We show that our algorithm, based on finding large attention scores, given K and query satisfying model assump- tions, finds the relevant set of keys. For a plausible set of model parameters, the running time of our algorithm per query, amortized over many queries, is sublinear in n. Throughout, K will be an x d matrix. K, denotes the 1 x d vector (the j-th row of K). Each row is a key. is a 1 x d vector, and will denote generic query. We assume there subset S of keys such that for E S, the correlation of K, to other keys is upper bounded. We will assume later that query vectors are convex combinations of keys in S. We let 51, 62 be parameters satisfying 0 < 82 S ö1 S 1/4. We assume there is subset S of keys satisfying:
Vj,Les, i e, KKyK7) S 81 mii (inn)) ll. ||
(1)
Yj S,E 5, K < 82 min(1] K (l))
(2)
Remark 4.1. S may be thought of as a " "stand out" subset of keys, since the correlation of E S to other keys is upper bounded.
We first argue that elements of S have high self-attention scores.
Lemma 4.1. Let A = KKT, For all i E (Aa)² 2 (A (1+67151) 82n).
Proof.
(Ax²)² (KKK?)² =1|K|l
7

---

Published as a conference paper at ICLR 2025
([Aa))=1KK3? [KKK)) + Σ(KKT)
< || ([+0((ss 1) +zn),
FESY
BES
using ① and 2
Now, both of them together imply the lemma.
Note that || canceled out because of the normalization. This demonstrates the effectiveness of row normalization, which a non-linear operation.
4.1 ASSUMPTION ON QUERIES
There is an unknown subset 5(1) D of S and unknown weights myy(g) for E S(1) satisfying
Σ
wz(g)= S 1 and uy(a) 2 48, for all E S(q)
(3)
jES(g)
and unknown d-dimensional vector z(4) P with
51 - Kell for all e [n], with q Σ jES(4)
|z(g)k7]
WyKy + 2(g).
(4)
4
We will just abbreviate (z)4) by wy.
4.2 PROPERTY OF RELEVANT KEYS
Theorem 4.2. We have Vj E S(q) 2kg 2 4 and Vj & S(9). aKT = K
Proof. For j E S(q): qKT wi [r 11 EeEstary (e(KK))) + =( 2 |K,18(45,
51 - -2/4) = (11/4)8||| using
and [ For j E S \\ S(q): lgk?
EiES(O) (e(KK) =(g) - || (81/4) 56, || /12/4, /4, using 3. (4) and
For j & S: 1gK3] = EnEscH) we + K3 S |KK,11(62 + (61/4) S (5/ K2 2
since 82 I 61.
4.3
ALGORITHM TO LEARN RELEVANT KEYS
A
Let +88 p. It easy to see that
Ej Ab is at most the i-th leverage score of K; since the
i-th leverage score sup Σ Ky. which at least as large as the value in this expression
(g.K,)²
with the particular value y K₁ Let U {i: LS(K)) 2 p}. Since the sum of all leverage scores is at most d, we have |U| - d(1 + 5/ + 52 For a suitable setting of parameters, we have
d(1 + 09|5| + nn E o(n) and |U| E o(n). Let U = {i: (A,)))))))) 2 p).
(A.))
We can find U using our algorithm for finding heavy leverage scores. We then compute | l| for all This takes ((nd) time. To estimate the row lengths of A KKT, we use a Johnson- Lindenstrauss sketch to find the row lengths of K KKFB for random O(lin) Gaussian matrix B (see, e.g., Woodruff et al. (2014) for background on Johnson Lindenstrauss sketches). When a query 9 arrives, we check for each U if qKT > (11/4)][ This suffices since S U'.
4.4 A STOCHASTIC EXAMPLE
We now present an example K in which the assumptions I and hold. The example has K as a random matrix described below. There is a latent subspace V of dimension k S d/4. For the conceptual description, we assume the first & of the d coordinates of vector are in V and the coordinates are in V1. Note that this is only for ease of description; we do not
remaining
Sg))) is the set of keys relevant to query *=(g) be thought of noise.
8

---

Published as a conference paper at ICLR 2025
actually know V. Intuitively, for i S, K; is mostly in V with small "noise" component in V whereas, for & S. K, is mostly in V with a small component in V. The matrix drawn below describes the distributions from which each part is generated. The coordinates are ii.d. inside each of the four parts with variance differing between parts. We assume 61 2 (4/k) and 82 2 (4e1 /k). No condition 6o is required. By standard concentration bounds, ( and (2) hold with high probability, so that Lemma (4.1) and Theorem (4.2) hold and our algorithm applies.
KL 1/011/K K3 =100k)
K² N(O, 8(0,e1/(d-) - K
Knxd=
N(O0, 1/(d k))
Here Kd snn and K2 is Epn x (d (d-k), and K is (1 -n) ks and K is (1-5o)nx -d-
EXPERIMENTS
5.1 STRUCTURAL PROPERTIES OF TYPICAL ATTENTION MATRICES
We consider a pretrained ViT model Dosovitskiy et al. (2021) for image classification and consider the properties of the attention matrices across all of the attention heads in the model. The model we consider has six layers and each of the layers has six attention heads. The inputs to the model are tokenized into 197 tokens- 196 tokens representing the input image and an additional token whose final representation is used for classifying the image (see Figure Tp. Therefore each of the attention matrices that we consider is of size 197 x 197.
For i, given the set of queries {Q1,...,Q ,Qa) E Rd and keys (Ki,.... .......,,). E TRd then for i E [n] and j E [n]. Ar,
V.) For each i, we define (1),,...,, (nn such that E, (G) A1,4(1 2 ... Ai,((())
Top Heaviness. We find that at many of the attention heads, the attention weight of a token is distributed in a "top heavy" manner, meaning that for many tokens i, the sum of the top-32 attention weights, defined as Σe= A,,.,)) is quite large. We show the dis- tribution of top-32 attention weights for an image input in Figure 2 in the Appendix. For given attention head, we observe that the distribution of top-32 attention weights remains similar across mul- tiple inputs.
Example of an image partitioned into 196 patches using 14 x 14 grid. The green patches represent the neighbors of the red patch within a distance of 3.
Figure 1:
Locality. We measure the amount of attention mass that is solely captured by the "neighboring" tokens (see Figure []. We observe that only at a small fraction heads, the attention mass tokens is fully captured by the local tokens and that in general, significant portion of the attention mass is distributed among non-local tokens. show the histograms of the atten-
In Figure Binn in the Appendix, tion weights captured by local tokens at each of the attention heads in all the layers of the model for typical input.
Important Keys. For j E [n]. we define W, to be the non-local
attention weight captured by j as Ei((.)) are not neighhors Aij. We define important keys to be the set of 32 keys with the largest W, values. We find that the initial layers of the models, at many attention heads, most of the attention weight is captured by the local tokens and small set of "important keys". In Figure in the Appendix, we show the distribution of the attention weight captured by the set of "impprtant keys" along with the local tokens.
5.2 INFERENCE FROM SOFTMAX VIT MODELS VIA LEVERAGE SCORE SELECTION
We consider three ViT models from the work of Dosovitskiy et al. (021): (i) S/16, a small 22M parameter model that splits image into 196 patches each of size 16 x 16 pixels, (ii) L/16, large

---

Published as a conference paper at ICLR 2025
Model
Accuracy on validation set
S/16 (softmax)
76.47%
S/16 (LevAttention, top-32, pretrained w/ softmax)
13.3%
S/16 (&2 norm selection, top-32, pretrained w/ softmax)
3.3%
S/16 (LevAttention, top-32)
68.30%
S/16 (LevAttention, top 64)
72.48%
L/16 (softmaxx
78.83%
L/16 (LevAttention, top-32, pretrained w/ softmax)
48.58%
L/16 (l2 norm selection, top-32, pretrained w/ softmax)
8.9%
L/16 (LevAttention, top-32)
75.12%
L/16 (LevAttention, top-64)
77.27%
L/16 (LevAttention, top-128)
77.17%
L/8 (softmax)
79.47%
L/8 (LevAttention, top-32. pretrained w/ softmax)
0.8%
L/8 (LevAttention, top-32)
71.96%
L/8 (LevAttention, top-64)
74. 74.64%
L/8 (LevAttention, top-128)
76.69%
Table I: Accuracies models with various attention mechanisms.
305M parameter model that splits an image into 196 patches each of size 16 pixels, and (iil) L/8, large 305M parameter model that splits an image into 784 patches each of size 8 x 8 pixels. Note that all the models have one token appended whose representation after processing through the transformer is used for classifying images. We train the models the Imagenet- (Russakovsky
ct al. 2015) dataset using the same hyperparameters as in the original work (Dosovitskiy et al. 2021 We find that the S/16, L/16 and L/8 models achieve accuracies of 76.47%, 78.83% and 79.47% respectively on the validation split the Imagenet- 1k dataset.
We then estimate the performance of leverage score selection these datasets. At each attention head, we compute the keys with the 32 largest fg leverage scores and then make each query attend only this set of keys in the attention mechanism. While find that the accuracies drop signif- icantly compared to the full softmax attention, our results show that the models still achieve non- trivial accuracies. In particular, the L/16 model has an accuracy of 48.58% using this mechanism. we instead select top-32 keys based on the squared row norms, we observe that the accuracy of the L/16 model drops to 8.9%. This supports the idea that leverage scores are much better predictors of the importance of the keys compared to row norms.
5.3 TRAINING VIT MODELS
In the previous experiments, we used models trained using softmax attention and used Attente-- tion only at inference time. To see if the performance of the models improve they are aware of LevAttention,
the leverage score selection based attention mechanism to train the models as well. Using the same training setup as the softmax attention models, i.e., the same learning rate schedule, batch sizes, and optimizer, For the L/8 and L/16, we train for the initial 15% of the steps with full attention to obtain the warm start parameters and then train the remaining 85% of the steps using the leverage score selection based attention. We report the results in Table To test training the models with the awareness of "leverage score selection" truly is useful, we train models using (i) "row norm selection" where at each attention head, we pick 32 keys with the largest & norms, and (ii) "random selection", where at each attention head, we pick 32 random key positions and only make the queries attend those po- sitions throughout the training process. We observe that these models achieve similar performance to the models trained via "leverage score selection". This shows that the "selection aware" training procedure is currently unable to translate the usefulness of leverage scores, as identified in the previ- ous section, into obtaining better models than those achieved by "row norm selection" and "random selection" attention mechanisms. We note that only leverage scores have some provable guarantees, while the other methods do not. We leave the important question of obtaining better models trained using LevAttentio future research direction.
significant improvements in the validation accuracies.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Alman and Zhao Song. Fast attention requires bounded entries. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances Neural formation Processing Systems 36: Anmual Conference Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Irwan Bello, William Fedus, Xianzhi Du, Yann Dauphin, Ashish Srinivas, Tengyu Zhou, Zihang Barret, and Bryan Zoph. Lambdanetworks: Modeling long-range interactions without attention. In International Conference on Learning Representations, 2021.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint (rXx:20044005500 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprini
2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Belanger, Lucy Sainath, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 2021.
Kenneth L Clarkson and David P Woodruff. Numerical linear algebra in the streaming model. In Proceedings of the forty-first annual ACM symposium on Theory of computing, pp. 205-214, 2009.
Kenneth L Clarkson and David P Woodruff. Low rank approximation and regression in input sparsity time. In Proceedings of the forty-fithh annual ACM symposium on Theory of computing, pp. 81- 90, 2013.
Kenneth Clarkson, Ruosong Wang, and David P Woodruff. Dimensionality reduction for tukey regression. In Proceedings of the 36th International Conference on Machine Learning, pp. 1262- 1271. PMLR, 2019.
Michael B. Cohen and Richard Peng. Lp sampling by lewis weights. CoRR, abs/ 150.007166 C7156
2015. (RRppttp://axxit
Michael B Cohen, Cameron Musco, and Jakub Pachocki. Online row sampling. arXiv preprint 2016a.
Michael B. Cohen, Jelani Nelson, and David P. Woodruff. Optimal approximate matrix product in terms of stable rank, 2016b.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //openzeview.
Petros Drineas, Malik Magdon-Ismail Michael W Mahoney, and David P Woodruff. Fast approx- imation of matrix coherence and statistical leverage. Journal of Machine Learning Research, 13 (1):3475-3506, 2012.
G. Foster. On the stochastic matrices associated with certain queuing processes. Annals of Math- ematical Statistics, 24:355-360, 1953.
Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. A fast optimization view: Reformulating single layer attention in LLM based on tensor and SVM trick, and solving it in matrix multiplication time. CoRR, 0ttx/2309.07418 2023. doi: 10.48550ARXIIV23000 074188 URL httpps//do/ org/10. 4854445550000 art.v 2309.07418
Ankit Gupta, Guy Dar, Shaya Goodman. David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention. arXiv preprint arXiv: 106. 06899, 2021.

---

Published as a conference paper at ICLR 2025
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In The Twelfth International Con- ference Learning Representations. ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- view.net, 2024.
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketching polynomial kernels. Forty-first International Conference on Machine Learning, ICMI. 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001 200X1:2000100451,, 2020.
Xiaoyu Ma, Yiding Wen, Yuanmeng He, Hai Zhao Liu, Dianhai Liu, Jianhua Li, and Xuan Dong. Nested attention for long document language modeling. arXiv preprint arXiv:2105.11877 2021.
Michael W Mahoney et al. Randomized algorithms for matrices and data. Foundations and Trends@ in Machine Learning, 3(2):123 2244, 2011.
Cameron Musco, Christopher Musco, David P Woodruff, and Taisuke Yasuda. Active linear re- gression for bp norms and beyond. In 2022 IEEE 63rd Annual Symposium Foundations of Computer Science (FOCS), pp. 744-753. IEEE, 2022.
Swati Padmanabhan, David Woodruff, and Richard Zhang. Computing approximate lp sensitivities. In Advances in Neural Information Processing Systems, 2023.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, pp. 4055- 4064. PMLR, 2018.
Hao Peng. Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXwv:100.002114, 2021.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJC), 115(3):211 252, 2015. doi: 10.1007/s1 1263-015-0816-y.
Tamás Sarlós, Xingyou Song, David P. Woodruff, and Richard Zhang. Hardness of low rank approx- imation of entrywise transformed matrix products. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, 2023.
Zhao Song, Junze Yin, and Lichen Zhang. Solving attention kernel regression problem via pre- conditioner. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li (eds.), International Con- ference on Artificial Intelligence and Statistics, 2-4 May 2024, Palau de Congressos, Valencia, Spain, volume 238 of Proceedings Machine Learning Research, pp. 208 216. PMLR, 2024. URL httppss//proceeeiggg... mlr ///os//////////////////
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng, Sparse sinkhorn attention. In International Conference on Machine Learning, pp. 9438-9447. PMLR, 2020.
Pichao Wang, Xue Wang, Fan Wang, Ming Lin, Shuning Chang, Hao Li, and Rong Jin. Kvt: k- nn attention for boosting vision transformers. In European conference on computer vision, pp. 285-302. Springer, 2022.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arxtv:2006.047663 2020.
http::/:wwwww....t..u/////////e www. haddlaafs/aa 110pg/5855992 w
David P. Woodruff, 2021. URL http:///wwwwoo.cooooooooa//e////
/woodruf/
David P. Woodruff and Taisuke Yasuda.
Online Lewis weight sampling, 2022. URL https:
//arx/w.orr//bb//2007
12

---

Published as a conference paper at ICLR 2025
David P. Woodruff and Taisuke Yasuda. Sharper bounds for Ep sensitivity sampling. In Proceed- ings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 36317-36336. PMLR, 2023.
David P Woodruff et al. Sketching as a tool for numerical linear algebra. Foundations and Trends@ in Theoretical Computer Science, 10(1-2)11-17,, 2014.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström- -ased algorithm for approximating self-attention. arXiv preprint arXiv:2102.003902 2021.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. CoRR, abs/2007.14062, 2020. URL
14062
Amir Zandich, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transform- ers via kernel density estimation. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Bar- bara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Ma- chine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Pro- ceedings of Machine Learning Research, pp. 40605-40623. PMLR, 2023. URL https: //prccediings mlr. wataas/v022/zaaaacm wttp://///////eedeee2...... html
Zhe Zheng, Da-Cheng Juan, Yi Tay, Dara Bahri, and Donald Metzler. Randomized attention with linear complexity. In Advances in Neural Information Processing Systems, 2022.
13

---

Published as a conference paper at ICLR 2025
6 APPENDIX: EXPERIMENTAL RESULTS
In this section, we present empirical observations on the behavior of attention matrices a ViT model trained using softmax attention. We use the same setup as in (Dosovitskiy et al. 2021 to train a 6-layer transformer model with 6 attention heads in each of the layers on the Imagenet- Ik dataset. We consider typical imput to the model and in Figures 2 and 4 we present the properties of the attention weights matrix across different layers and attention heads in the model.
Logaoh
Laye
Heal
Layen SHANAS
Lager
ASAITA LEVICI Fe5o3
23 K4
LiWe (18622
LEVe33 Hend
Layer (HM0O4
Liyer: Fand
32
Ligee H467
Layeed Held
Liger Hoad
Herd)
Liverd Ferd
Figure 2: Histograms of top-32 attention weights. Each of the histograms plots the distribution of the sum top-32 attention weights for query tokens. We note that many attention heads, for many query tokens, the largest 32 attention weights constitute a significant fraction of the total attention weight.
14

---

Published as a conference paper at ICLR 2025
Layer HMBGE
Liyal1 19851
Layect H46111
Laye: read
Mypec HHMES
53
63 150 aborisn wesi
Leyert
Layar2 Headd
Leverd Peart
Loved Fead
1eyer2/Hertth
LipAC3 H8831
Lagar2
25 M. 24
26 63 C5
34
Hour??
053-
Figure 3: Histograms of attention weights captured by local tokens. Each of the histograms plots the distribution of attention weight captured by keys within Manhattan distance of 3.
15

---

Published as a conference paper at ICLR 2025
100 132
Ligerl
Lagatl HeA.
54 saree
Levert PAACE
Loverd Peall
Lover 1ter17
Levert Pest)
LevenS
Figure 4: Histograms of Attention weights captured by the important keys along with local tokens. As discussed in Section 5.1 important keys are defined as the 32 keys capturing the largest attention weight given attention head. The histogram shows that a large number of attention heads, the important keys together with local keys are able to capture significant fraction of the attention weight for large number of queries.
16

---

Published as a conference paper at ICLR 2025
7 APPENDIX: OVERVIEW OF RECENT RELATED WORK FOR REDUCING COMPUTATIONAL COMPLEXITY
Zaheer al. 2020 propose Big Bird, sparse attention mechanism that combines global and local attention to handle longer sequences in transformers. While Big Bird employs a fixed sparsity pat- tern, our work introduces a data-dependent sparsity approach by identifying universal set of keys based on leverage scores. These approaches could be seen as complementary: Big Bird provides a general framework for handling long sequences, while our method offers a more fine- ane-grained approach for identifying the most relevant information in the attention matrix. Combining these techniques could an interesting direction for future work.
Song et al. (2024) address the computational challenge of solving attention kernel regression prob- lems, where the matrix exponential of the Gram matrix is the kernel. They propose using pre- conditioner accelerate the solution of these regression problems. While their work focuses specific formulation of the attention mechanism as a regression task, our work provides a more gen- eral method of approximating the attention matrix by identifying universal set of important keys. Our approach could potentially complement their work by providing a reduced set of keys, which could further speed up their pre-conditioned solver.
Gao et al. 2023) propose a fast optimization approach for training single-laye attention networks in L.MMs. They reformulate the attention mechanism using tensor and SVM tricks achieve training time comparable to matrix multiplication. While they modify the training process itself to improve efficiency, our work focuses on identifying a core set of important keys, which can be used to speed up both inference and training. Both their approach and ours could potentially be used in conjunction with learned leverage scores during training. They could incorporate learned leverage scores into their tensor-based based optimization framework, while we could use learned leverage scores to construct our universal sets. Exploring the interplay between these techniques could be an interesting avenue for future research.
precursor HyperAttention that uses fast kernel density estimation algorithms to approximate the softmax attention. These algorithms efficiently estimate the normalization factor for each row of exp(QK")) as well as the sampling probabilities used in its approximate matrix multiplication. However, KDEFormer requires at least n 1-13 time and relies on assumptions such as bounded diameter datasets and small stable rank to achieve its strongest results. In contrast, method directly identifies the most important columns with theoretical guarantees and runs in linear time with respect to n, without requiring these assumptions. Furthermore, as dis- cussed in work, the universal set we identify be used input to HyperAttention, effectively reducing its runtime by focusing its computation on the most relevant tokens. This could potentially lead to both faster and more accurate attention mechanisms.
KDEFormer (Zandieh et al. 2023
17


Reviews:
[{"review_id": "QhZk1uyCEq", "reviewer": "Reviewer_MVJp", "summary": "This paper studies efficient algorithms for attention computation in transformer models. In particular for query $Q$ and key $K$ matrices in $\\mathbb{R}^{n\\times d}$, and function $f$ applied to each entry of $QK^T$ (for eg. $f=x^p$ for some $p>0$), this paper studies efficient algorithms to approximately compute $f(QK^T)$ followed by a row-normalization. The main contribution of the paper is to show that for any $\\epsilon>0$ there exist a subset $U\\subset[n]$ of keys, that is the rows of the key matrix $K$, such that for any query $q$, that is a row of query matrix $Q$, if the attention score of $q$ with $k_i$ after normalization is greater than $\\epsilon$ then $i\\in U$. The size of this set $U$ is $(\\sum_{i\\in [n]}\\sigma_i^f(K))/\\epsilon$ where $\\sigma_i^f$ is the $i^{th}$ $f$-sensitivity score of $K$. For eg. when $f=x^2$ these sensitivities are nothing but the $\\ell_2$ leverage scores of $K$. The set $U$ is independent of $Q$ and thus can be used to compute attention scores with any $Q$ in the future for a fixed $K$. Moreover using fast algorithms developed in the literature for computing $f$-sensitivities for a broad class of functions $f$, the set $U$ can be computed efficiently. For eg for any constant $p$, for $x^p$ the set $U$ can be computed in time $nnz(K)+poly(d/\\epsilon)$ using the input sparsity time algorithm of Cohen and Peng for computing $\\ell_p$ leverage scores.", "questions": "The main question I have is around the weakness that I brought up regarding the experimental section - that is how would the case be handled when $K$ has high rank and every key has high sensitivity ? Morever in the experimental section have the authors considered evaluation with other methods which are popularly used in attention approximation such as flash attention ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The main strength of the paper is to apply tools from randomized numerical linear algebra to naturally arrive to the conclusion that if the attention score of a query is more than $\\epsilon$ with any particular key, then simply by definition it implies that the $f$-sensitivity (or leverage score for the special case when $f=x^2$) will also be more than $\\epsilon$. This directly implies that capturing all the keys of $K$ with sensitvities higher than $\\epsilon$ suffices to prove their result. Thus I think the main strength of this paper is develop this connection in more detail, and prove results in various computational settings such as streaming and distributed settings regarding how to efficiently compute this set of universal keys, which I think is a good contribution to the literature.", "weaknesses": "One aspect I would want to get more clarity on is in the experimental section. The authors do arrive at the conclusion that the mass of practical attention matrices for each query can be captured by a small set of universal keys plus a set of local tokens for that specific query. However if this is the case, then algorithms for computing the set of local tokens should also be considered on top of computing the set of universal keys. For eg. it may be the case that in natural language applications, there are sentences in which for each token, it has a high correlation with a few local tokens in a small window around it. In this case, the leverage scores/sensitivities of each key may be large resulting in a large universal set, and the attention matrix has a high rank. However, the attention matrix can still be efficiently approximated with a sparse matrix as it essentially amounts to computing, for each token, the small set of local tokens with whom it has a high correlation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "E5QXodBYEH", "reviewer": "Reviewer_k6Cn", "summary": "This paper develops efficient algorithms to identify significant attention scores in transformer architectures without computing the full attention matrix. The key theoretical contribution is proving the existence of a small \"universal set\" of keys, independent of sequence length, that captures all large attention scores for any query. The authors provide efficient algorithms to find this set and compute attention scores in streaming and distributed settings. The work provides both theoretical guarantees and practical benefits for improving transformer efficiency.", "questions": "1. What are the trade-offs between exact and approximate computation of attention scores?\n\n2. Can the approach be parallelized effectively across multiple GPUs (Tensor Parallel)?\n\n3. What is the impact on training time and convergence?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Rigorous mathematical proofs for existence and size of universal sets\n2. Straightforward integration with existing transformer architectures\n3. Strong empirical results on vision transformers", "weaknesses": "1. No evaluation on (large) language models. Limited ablation studies on prediction quality\n\n2. Limited ablation studies on prediction quality", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "Ao2o36hgAx", "reviewer": "Reviewer_n1PT", "summary": "The attention mechanism is a bottleneck for transformers in terms of efficiency. Standard method of attention mechanism takes quadratic time with respect to the input, where the main issue is to calculate the n by n attention matrix A. This paper studies a class of attention mechanism where we replace f(x) = exp(x) in the vanilla attention with f(x) = |x|^p for some p, and show that regardless of the input data, there exists a set of columns of A that include all the large entries. Moreover, one can find this set efficiently. As a result, one can hope to approximate A by retaining only these large entries efficiently. The paper then generalizes the results to streaming and parallel settings.", "questions": "The result in this paper, as the authors stated, applies to any Q. It seems like this is because of the way f-sensitivity is defined (by taking the sup over all y). In practice Q might be structured (I think there is a shared-QK transformer where Q and K are identical), and I wonder if we can say anything about it.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Speeding up/approximating the attention mechanism is an interesting topic to study. \n2. It is natural to think that as n (sequence length) grows larger, there will only be a subset of keys that are similar, i.e. have larger inner product, with the queries. This paper gives a theoretical analysis on how large such a subset can be and how we can find it. \n3. This paper is clearly written and the mathematical contents are interesting.", "weaknesses": "1. The paper mentions that the standard attention mechanism cannot be well-approximated unless SETH is false (Alman & Song 2023), so I guess the point is that we are hoping for a good/efficient approximation for other attention variants. Although f(x) = |x|^p has been proposed and used in previous work (PolySketchFormer, TensorSketch), I don’t think they are the “gold standard” for attention, and therefore the motivation of studying this polynomial attention is not so clear to me. \n2. Even if we can find all the columns containing all the large entries, does that imply a reasonable approximation after we multiply the attention matrix with the value matrix V? In addition, if we consider multiple self-attention layers, even though we can find the important columns of all attention matrices, it is not clear to me whether we can still have any kind of guarantee on the final output.\n\nOverall I think the problem that this paper proposes is an interesting problem in numerical linear algebra, but I am not really convinced about its impact on transformer theory.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "E5b4M39Dg0", "reviewer": "Reviewer_tuuu", "summary": "In this paper, the authors study the attention computation problem: given the query, key, and value matrices $Q, K, V \\in \\mathbb{R}^{n \\times d}$, the goal is to output $D^{-1} A V$, where $A = \\exp(QK^\\top/d) \\in \\mathbb{R}^{n \\times n}$ and $D = \\mathrm{diag}(A {\\bf 1}_n) \\in \\mathbb{R}^{n \\times n}$ is a diagonal matrix. Exactly computing the product of $D^{-1}$, $A$, and $V$ requires $O(n^2 d)$ time, which can become very large when the length of the input token $n$ is large, so in this work, the authors present a novel theoretical framework and algorithm for efficiently identifying the large attention scores (the entries of the attention matrix $A$), while maintaining the linear time complexity with respect to $n$. Additionally, this paper finds the large attention scores for a large class of function $f$ in $A = f(QK^\\top/d) \\in \\mathbb{R}^{n \\times n}$, which includes all $f(x) = |x|^p$.", "questions": "1. $\\epsilon$ is used to determine the large attention scores. How can we decide the value of $\\epsilon$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. This paper gives a very solid theoretical foundation showing the existence and properties of universal sets and finds all of the large attention scores with size independent of sequence length. It applies the techniques from randomized numerical linear algebra to study the attention problem. \n\n2. The theoretical results do not rely on any restrictive assumptions and are applicable to a wide variety of functions, including all $f(x) = |x|^p$. The algorithm can be computed in streaming and parallel settings.\n\n3. The experimental results support the theoretical findings by considering a pre-trained ViT model. Moreover, the model quality has more than 90 percent of the accuracy of the full softmax attention when only selecting the top 32 keys.", "weaknesses": "1. This paper does not include an overview of the recent works, which might make it hard for readers outside this field to understand the importance of the contributions. Works like [ZGD+20] develop the sparse attention, reducing the computation complexity to $O(n)$, and other works like [SYZ23, GSWY23, ZHDK23] are more related to this work: applying the numerical linear algebra techniques, including tensor trick, preconditioner, and sketching to reduce the computational complexity.\n\n2. The experiment only focuses on the pre-trained ViT model. This paper does not study any other language models. Therefore, it is hard to see whether or not the framework developed in this paper can be generalized to more language models.\n\n[ZGD+20] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham et al. \"Big bird: Transformers for longer sequences.\" NeurIPS'20.\n\n[SYZ23] Zhao Song, Junze Yin, and Lichen Zhang. \"Solving attention kernel regression problem via pre-conditioner.\" AISTATS'24.\n\n[GSWY23] Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. \"A fast optimization view: Reformulating single layer attention in llm based on tensor and svm trick, and solving it in matrix multiplication time.\" Preprint'23.\n\n[ZHDK23] Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. \"Kdeformer: Accelerating transformers via kernel density estimation.\" ICML'23.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)