Improved Algorithms for Kernel  Matrix-Vector Multiplication Under Sparsity Assumptions
Authors: Piotr Indyk
Michael Kapralov
Kshiteej Sheth
Tal Wagner
Venue: ICLR 2025 Poster (2025)
Keywords: Algorithms
Kernel Matrix
Kernel Density Estimation
Locality Sensitive Hashing
Fast Attention
Primary area: learning theory
TL;DR: We design fast algorithms for kernel matrix vector multiplication motivated by fast processing of attention matrices in LLMs.

Abstract: Motivated by the problem of fast processing of attention matrices, we study fast algorithms for computing matrix-vector products for asymmetric Gaussian Kernel matrices $K\in \mathbb{R}^{n\times n}$. 
$K$'s columns are indexed by a set of $n$ keys $k_1,k_2\ldots, k_n\in \mathbb{R}^d$, rows by a set of $n$ queries $q_1,q_2,\ldots,q_n\in \mathbb{R}^d $, and its $i,j$ entry is $K_{ij} = e^{-\|q_i-k_j\|_2^2/2\sigma^2}$ for some bandwidth parameter $\sigma>0$. Given a vector $x\in \mathbb{R}^n$ and error parameter $\epsilon>0$, our task is to output a $y\in \mathbb{R}^n$ such that $\|Kx-y\|_2\leq \epsilon \|x\|_2$ in time subquadratic in $n$ and linear in $d$. Our algorithms rely on the following modelling assumption about the matrices $K$: the sum of the entries of $K$ scales linearly in $n$, as opposed to worst case quadratic growth. We validate this assumption experimentally, for Gaussian kernel matrices encountered in various settings such as fast attention computation in LLMs. Under this assumption, we obtain the first subquadratic time algorithm for kernel matrix-vector multiplication for unrestricted vectors.

Published as a conference paper at ICLR 2025
IMPROVED ALGORITHMS FOR KERNEL MATRIX- VECTOR MULTIPLICATION UNDER SPARSITY ASSUMPTIONS
Piotr Indyk
Michael Kapralov
Kshiteej Sheth
MIT indykémit. edu
EPFL
EPFL
michael, waw.aaleee.paalleee.ll...c
kshiteej kshiteej.shetth@ppL] sheth@epfl.ch.
Tal Wagner
Tel Aviv University
t www.etaue..ta..cooom
ABSTRACT
Motivated by the problem of fast processing of attention matrices, we study fast algorithms for computing matrix-vector products for asymmetric Gaussian Kernel matrices K E TRaxn, K's columns are indexed by a set of n keys k1, .2.....e kn E RR rows by set of queries G1 122.-...9 In. E Rdd and its i, entry is Kij e 11ll=k,11/220 for some bandwidth parameter o > O. Given a vector x E R²" 0, task is to output y R such that / vll2 S c||=/||2 in time subquadratic in n and linear in d. Our algorithms rely the following modelling assumption about the matrices K: the sum of the entries of K scales linearly in n, as opposed to worst case quadratic growth. We validate this assumption experimentally, for Gaussian kernel matrices encountered in various settings such as fast attention computation in LLMs. We obtain the first subquadratic-tines algorithm that works under this assumption, for unrestricted vectors.
and error parameter
INTRODUCTION
Linear-algebraic operations kernel matrices play important role in machine learning. One of the most widely used operation computes a product of a Gaussian kernel matrix with another matrix or vector. Formally, let & Rd x IR > R+ be such that k(x, y) e je-vll//2e for some parameter 0. The kernel matrix is defined by two sets, keys { .2,..., k, } and queries {qi 92....... In where ky's and g are elements of Rdd The entries of K are defined as Kij k(g;, kj) for all i (n]. The computational task is defined follows: given Ky "s qi 's and product Kx, or its approximation. In typical applications, both d are large but > d.
R compute the
The kernel matrix-veet vector product has many applications machine learning and artificial intelligence. For example, if æ is the all-ones vector, this operation corresponds to Kernel Density Estimation, classic tool in non-parane non-parametric statistics, where the kernel function is used to extend the empirical distribution function over discrete set of points smoothly to the whole space. More recently. the problem emerged as a key computational subroutine in transformers (Vaswani et al., 2017). One of the key computational task in training and inference of transformers is to compute the product AV. where A = e Ag))) is the "attention matrix" and V consists of d column vectors 2,. A recent paper (Zandieh et al., 2023) gave a reduction that replaces attention matrices with Gaussian kernel matrices, so that the algorithms for Gaussian kernel matrices could be applied to attention matrices as well. fast kernel matrix vector product for Gaussian kernel matrices can then not only be used for fast attention computation but for other important computational tasks such as investigating the spectrum of attention matrices quickly by computing its eigenvalues using the kernel noisy power method presented in the work of Backurs et al. (2021). Thus motivation is study the kernel matrix vector product, rather than solely focus on fast attention computation which is the case the works of Zandieh et al. (2023); Han al. (2023) for example.

---

Published as a conference paper at ICLR 2025
A direct algorithm for kernel matrix-vector product takes time O(n²d). The quadratic dependence on n has been widely identified as a significant bottleneck in many applications, including transform- ers (Kitaev et al., 2020; Choromanski et al., 2021; Beltagy et al., 2020; Chen et al., 2021; Wang et al., 2020; Zaheer et al., 2020; Xiong et al., 2021; Zandieh et al., 2023; Han et al., 2023). Unfortunately, Backurs et al. (2017); 7); Keles al. (2023); Alman & Song (2023) gave evidence that algorithms that compute K2 or AV in time sub-quadratic in n are unlikely to exist for oigh-precision algorithms (i.e. algorithms that can achieve (/poly(n) error in polynomial time), in the worst case. In the low precision (i.e. algorithms that can achieve 1/ s/poly(logm) n.) error in polynomial time) high dimensional regime, the work of Backurs et al. (2021) gave a on?²) time approximate kernel matrix vector product algorithm, however it could only handle multiplying the matrix with non-negative vectors. This forms the baseline for our work.
Most of the algorithmic efforts have focused on designing approximation algorithms for the special cases of matrices which occur in practice. The contributions of these studies! are two-fold. First, they identify classes of matrices that accurately model the matrices occurring in practice. Second, they develop efficient algorithms for the identified classes of matrices.
1.1 OUR RESULTS
In this paper we present new model for Gaussian kernel matrices that observed in practice especially in the context of large language models, and propose improved approximate matrix-vector multiplication algorithms. Formally, for error parameter e 0, keys k ...k, kn, and queries QI .I....9 -n defining K, and a vector x, we want to output a vector in time o(n²). -olyy (d, 1/e) such that ||Kx - y < e||=r|12-
It has been observed in practice that on average over an input sequence of length n, each token in the sequence has high correlation with only few other tokens. This implies for self-attention, Gaussian kernel and other similarity matrices there are about n large entries. This motivates our modelling assumption about Gaussian kernel matrices K::
The ratio of the sum of all except the largest n entries of K (i.e. the sum of the tail of K)) and the sum of the largest entries of K (i.e. the sum of the head of K) is at most constant e > independent of n.
(A)
In Section 4 we validate this assumption for a collection of Gaussian kernel matrices K derived from attention matrices obtained by running BERT (Devlin et al., 2018) on sentences from Stanford Question Answering Pataaet((aajpukka) et al., 2016) (Section 4 contains formal details about obtaining Gaussian kernel matrices from self-attention matrices). For each attention head and layer in BERT, we compute the head-to-tail ratio as a function of matrix size. Our experiments shows that the maximum value of this ratio e is most 4.6, over all sentences, heads, layers and matrix size values. This confirms the validity of our assumption. We also perform this experiment, as well as additional experiments on the scaling behaviour of с with the context length on BERT and other language models such as RoBERTa (Liu, 2019) and GPT (Radford et al., 2018) in the Appendix A.I.
In Section 4 we also investigate a stronger assumption, where (informally) one postulates that there is a small uniform upper bound on the values of the entries in the tail of the matrix, which is orders of magnitude smaller than the values of the entries in the head of the matrix. 2 This is similar to the assumption made in Han et al. (2023), though in this paper we consider it in the context of Gaussian kernel matrices K, not attention matrices. Our experiments indicate that this assumption does not model matrices K well. Specifically, we show that the median ratio between the smallest entry of the head (i.e., the th largest entry of K) and the largest entry of the tail (i.e., the (n + 1)th largest entry of K) is very close to 1. In fact, even the median ratio between the n and (2m)th largest entry is about 20 in most cases. This demonstrates the usefulness of our assumption, which quantifies the tail according to the & norm, not the foo norm. Please refer to Section 4 for precise details.
Our algorithmic result is encapsulated by the following theorem.
Theorem 1.1. Under the assumption that K satisfies A, then in time (dn 1.89 1 the Algorithm 3 APPROXKMV outputs y H. such that it satisfies 1| ull2 S e||z||2 with probability 0.99
'See Related work for the overview.
NNote that if this gap is large enough, it implics our assumption.
Success probability can he achieved for any 8 > O, with an additional log(1/8) factor in the runtime.
2

---

Published as a conference paper at ICLR 2025
The complete algorithm and its proof is presented in Section 3. Crucially, the running time is o(n²), Prior to our work, subquadratic time algorithms in the nigh-dimenssonal regime (i.e. running time depends polynomially rather than exponentially on d) for kemel matrix- vector multiplicatton were not known for general vectors x, see Section 1.2. To summarize, our contributions are as follows:
We put forward a new modelling assumption for kernel matrices;
On the one hand, we show empirically that our modelling assumption holds for kernel matrices that arise in modern transformer based language models;
On the other hand, we show that our modelling assumption provably leads to subquadratic time algorithms for approximate matrix vector multiplication. As result, under our modeling assumption we obtain sub-quadratis time algorithm for high dimensional approximate kemel- matrix vector multiplication, that runs for general vectors.
1.2 RELATED WORK
We follow a line of work on hashing- based algorithms for kernel computations on -igeesinnnii points, pioneered by Charikar & Siminelakis (2017), and continued in Backurs et al. (2018); Siminelakis et al. (2019); Backurs et al. (2019); Charikar et al. (2020); Backurs et al. (2021); Karppa et al. (2022); Zandieh et al. (2023). Starting at the problem of kernel density estimation (KDE), Charikar& & Siminclakis (2017) considered the data structure setting, defined as follows: Let X be dataset of points in RR, and let μ E (0, 1) be precision parameter (for intuition, it is instructive to consider μ = 1/n where n = [X The goal is to preprocess X so as to enable efficiently reporting the KDE value [x] Ex.y k(x,, y) at any incoming query y. as long as the its true KDE is at least μ. By vanilla uniform sampling, KDE queries can be answered up relative error + e in time linear in 1/p, namely 2(d/epu). Charikar & Siminelakis (2017) showed that, by using locality sensitive hashing (LSHH) (Indyk & Motwani, 1998), it is possible to answer KDE queries in time (d/ (()) with p < 1, which is sublinear in 1/μ. For the Gaussian kernel, currently the best known value for p is = 0.173 + o(1), due to Charikar et al. (2020).
Charikar & Siminelakis (2017) also observed that their techniques can be used for fast algorithms for estimating the matrix product Kr of a kernel matrix K and vector x. In Backurs al. (2021) this was formalized into algorithm that, given an x kernel matrix K and R". outputs vector y that satisfies | vll S <>|K21 in time (en2+t/3**2) provided that x has only non-negative entries. For Gaussian kernel matrices, by plugging the aforementioned bound on p from Charikar et al. (2020), the dependence on n is n² (11 o(n²). To our knowledge, this is the only prior subquadratic time algorithm for kernel matrix-vector multiplication in the high-dimensional (i.e. when d is very large) regime.
The main limitation of Backurs et al. (2021) is the requirement that is non-negative. They used their kernel matrix-vector multiplication algorithm as a subroutine for estimating the top eigenvalue of K, which based on the classical Perron-Frobenius theorem, allowed them to only deal with non-negative vectors. However, in many applications, there is no way to enforce the non-negativity of x. Note that this limitation is inherent to their approach: the error in their approximation guarantee is el| x: lz. which in general can be zero (if x is in nullspace of K). Thus, in general it may require computing Ka exactly, which takes time 2(n2).5
To overcome this, we study the natural approximation guarantee | - u|| I <||æe|| instead of el| [ see Theorem 1.1. This notion of error independent of whether lies in the nullspace of K or not. This allows us to achieve subquadratic time algorithms without any restrictions, and in particular removes the non negativvity restriction on xr
Nonetheless, we note that our algorithm improves over Backurs et al. (2021) even for inputs restricted to their setting, i.c., where x is non-negative. This is true in two senses. First, for such inputs, their
algorithm's 's eтror e|| Kx 2|2 is always at least as large as our eтror, e <||z l2. This is because K, being
kernel matrix, has non-negative entries with an all-1s diagonal, hence [ krl 2 = |x (K T)x 1
"Our result of Theorem 1.1 has subquadratic runtime even for d (1)) as compared to previous work of (Alman & Song, 2023) which was only applicable for d = O(lo n).
For example, with kernel matrices, one can essentially realize zero matrix Ko. and also "hide" a single 1-entry in an otherwise zero matrix K1, see, e.g., Backurs al. (2017). Computing Kr: exactly entails distinguishing between Kg and K2 with high probability, which requires ((n) ) time.
3

---

Published as a conference paper at ICLR 2025
| + 2x - 1)x + II(K - 1 2 ||=|ß. Second, there are enror regimes where even for non non-negative x, their algorithm fails to run in subquadratic time, while ours does so. For example, consider the case when is the all ones vector denoted by I 1, Then the error incurred by the algorithm of Backurs et al. (2021) will be <l < and will run in time Onn ++/ /+-tt//240 /c²+2p) where p 0.173 mentioned previously. Consider the case when K contains one row of all ones and all other rows are O, then e e|| 1 = E n. Thus we would have to re-scale E by n to achieve our error guarantee of E|| 2 n=te. (n ()) multiplication. On the other hand our algorithm achieves this guarantee in o(n²) time.
n..5, Thus the runtime of Backurs al. (2017) will be at least ((n failing to achieve subquadratic time better than naive matrix-vector
We note that besides LSH, there are other approaches for fast kernel computations that can be used with the above line of work, like the fast Gauss transform (Greengard & Strain, 1991). While this also leads to kernel -vetox multiplication algorithms with running time subquadratic in n, the running time depends exponentially on the dimension d of the underlying points {k, } that define the kernel matrix, and is thus unsuitable for high-dimensional regimes, and particularly for deep learning models.
1.3 OVERVIEW OF OUR TECHNIQUES
We now give a high level overview of our algorithm, its details with proofs are presented in Section 3. Recall our goal is the following: given an error parameter >> O, keys kn ... kin and queries QI .....-00 Ur defining K, and vector 2, we want to output a vector y such that || ylla S e||2||2.
Pre Processing x: Firstly since our guarantee is free from the scaling of æ, we assume | = n. Now we pre-process x to explicitly calculate the contribution of extremely large entries of I to Kx, since || | = n we can't 't have too many extremely large entries in x. Next we round the extremely small values of z: to O, since the entries are extremely small and entries of K are bounded by 1 this incurs negligible error. This pre-processing of described formally Section 3.1, and it renders a's remaining values to be in bounded range.
Finding heavy keys: In the next phase for every query 9: for i E [n]. we will find all the keys k; for j E [ri] such that k(gs, k;) is large. We call such keys "heavy" for query q. Then we will calculate exactly the contribution of such heavy keys to (Kx); for every i E [m]. We will show this can be done in time o(n²) by first showing that assumption A on K implies we cannot have too many heavy keys per query on average, coupled with a fast locality sensitive hashing based recovery procedure to find all heavy keys per query. This is discussed with all details in Section 3.2.
Estimating the contribution of light keys: The final phase of our algorithm will be a random sampling based procedure to estimate the contribution of all the non-heavy, henceforth light, keys corresponding to query q to (Kr); for all E [n). We will uniformly sub-sample each light key with probability 1/n and calculate the (scaled) contribution of the surviving keys get a basic unbiased estimator for the contribution of all light keys. We will show that the variance of this estimator will depend on the sum of squares of the contribution of every light key to (Kr); This variance will also be the number of repetitions, up to poly(log n. 1/c) factors, we need to do of the basic estimator to reduce its variance by averaging to within our error bound. Our main innovation is to show that the number of repetitions for each row, which may potentially be different across rows, can approximated using a fast Gaussian kernel density estimation primitive. Please refer to Section 3.3 for full details.
2 PRELIMINARIES AND NOTATION
For any integer n > O we let [r]] to denote the interval {1, ,... We let 1, E R denote the all ones vector and we use 1 g be the indicator variable for any event E. For any matrix A E Rx for some integers m, n > O, we denote its j.j entry for any i E [m].j E [m] as Ar,. We let Al: i, j]] to be the sub matrix of A that contains first i rows first j columns for any E [m] and j [m]. For any vector & we use ||~|||,, ||>/|] to denote its f2, & norms respectively. For any matrix A we use ||Al|] denote the sum of all of its entries. We use Ó(-) suppress poly(log n. factors.
.n..
The first tool we will need in our algorithm are locality sensitive hash (LSH) functions which are used for solving high-dimensional approximate nearest neighbour search problems (Indyk & Motwani,
4

---

Published as a conference paper at ICLR 2025
1998; Andoni & Indyk, 2008). We first state the following claim about the LSH function of Andoni Indyk (2008) stated in a convenient form for us as Claim 19 in Charikar et al. (2020).
Lemma 2.1 (Claim 19 of Charikar al. (2020)). For any constant E [0, 1), there exists family of hash functions such that for Thear V2022 a n, the following holds for any T rffor 2 Thear.
1. F (p(pp) = h(qg)l 2 n for any |lp 4l1e S Tnear.
2. Ph (h(p) = h(4)] S n 2a for all 1lp min farr/Tenr))) log 17 m
gll2 = Tfar and c
We will also use recent algorithms for fast Gaussian kernel density estimation (KDE) (Charikar et al., 2020; Charikar & Siminelakis, 2017; Backurs et al., 2019). In this problem we are given dataset P Rdd containing n. points |P| n, the Gaussian kernel k<p; g) e 1 for some bandwidth parameter a > 0 and p,4 q E Rd The goal is to preprocess the dataset to create a data structure such that the query phase when given query q IR the data structure can approximate (Ep p k(p, /))/n to = 8 relative error in time o(n). We will use the following fast Gaussian KDE result of Charikar et al. (2020).
Theorem 2.2 (Theorem 2 of Charikar al. (2020)). Suppose we are given a set points P Rdd and parameters 3, O. For any point q E Red let p(q) = (Σ e- e- (1200 )/n. Then there exists data-structure with pre-processing time (((3-2nn/p) 01 log(1/6)). such that for any query q the data structure can output approximation to p(g) (m())2n) up + relative error in time ((33-2) pd/1..173) log(1/6) and success probability 1 - 6.
3 ALGORITHM
The goal of this section is to describe the main algorithm and prove Theorem 1.1. We will go about proving this using intermediate building blocks. We will work the following convenient re-phrasing of our Assumption A - the assumption says that if we denote K as our Gaussian kernel matrix then || |1 minus the sum of the largest n entries of K entries of K, thus |K]] is at most a constant times the sum of the largest entries of K. Since each entry in K is bounded by 1, the assumption directly implies that | O(n).
a constant times the sum of the largest
3.1 PRE PROCESSING
This section describes convenient pre- pre-processing of x, starting with the following notation. Definition 3.1. Let 7 E [0, 1] be a threshold. Define the following subsets of [n]] as follows,
H1 = [j [n] : 2 2 n²) H2 = (i E [m] 27 S 1 }, H H2 U H2, and T = [ni] H.
Let VH Vr E R be defined as follows, (yH); = Eyen. ki), and (yr); Ejer k(g ky)n, for all [n].
We now state the following lemma which says that уи + yr are good approximation of Kr and can be computed in o(n?) time. Its proof is provided in Appendix A. Lemma 3.2. time (dd n² 2²) we can output the set H and vector yn- Moreover |[r (yn
yr) 1 e||xe|12.
3.2 FINDING HEAVY KEYS
The next objective to approximate yт. The goal of this section is to give the algorithm that explicitly finds for all queries 9; for E [n], the set of all keys k, which have large contribution to Z,k P,k(g. k,). We call such keys "heavy" and we now formally define them.
E
Definition 3.3. Let a E [0, 1) be a threshold. Consider any E [n)]. For query define the set of "heavy" keys S, [i E [n] k(gi, ky) n n).
"The o(1) factor in the exponent in far collision probability is Ollo log n/ log 1/3 n)), and it is justified long O(log n).
5

---

Published as a conference paper at ICLR 2025
We now state the main lemma which says that find the set of heavy keys for all rows in on?²) time, its proof is in Appendix A. The pseudocode of the algorithm presented in Algorithm
Lemma 3.4. In time Odd n Algorithm / FINDHEAVY returns all the sets S; for E [m]. The algorithm succeeds with probability 0.99.
Algorithm 1: FINDHEAVY
1: Input: Keys k1, h2. ...... kn Queries q1, 92.............................. In threshold a > O.
2: Output: Sets St for all i [n] per Lemma 3.4.
3: Let 7 10n log n.
4: Let H be an Hash family as per Lemma 2.1.
5: Sample 7' i.i.d. /i....,hy hr ~ H. Hash entire dataset using these '' hash functions.
6: for E [n] do
7: Scan all the buckets (x;) for all t E [T] and return all points in
S {a E P k(z 24) 2 n =]
8: end for
3.3 ESTIMATING CONTRIBUTION OF LIGHT KEYS
After finding S,, what remains approximating EjeTis, k(as ky), for E [x] up to additive error e. This is the main goal of this section formalized in the lemma below, its full proof is in Appendix A.
Lemma 3.5. In time Ö(d (12+1~a + 1.7877/// (e)) Algorithm 2 APPROXLIGHT, when executed on S, for i [m] per Lemma 3.4 and set T, returns a vector E R" which satisfies 1 Ejer(s, ki (its,kg)a, kig) S for all i E [n] with probability 0.99.
Algorithm 2: APPROXLIGUTT
Input: Keys h kg...., .k,... Queries Q1, 92. On vector æ, parameters a, 7, > 0, set and sets S; for all E [n].
2: Output: A vector 2 IR" as per Lemma 3.5,
3: Let B,, (j T x [[1 +c)=- (1++6 e)" for [-41ogi+e (ogi))))))
4: For every m and j Bm let let
(1 (1+6) +c
5: For every B,, create a data structure as per Lemma 2.2 with data set {ky : Bm), error parameter -0.218 1 = e /ollnnoon (n) + ( I), 1/n²,, and kernel function k2 (²,
6: for [n]] do
7: Let t, be the data structure output for query q.. and let
8; =, Ejes, agk(gg ky)² 0.218
8: Sub-sample every key in T St with probability 1/n, and sum xzyk(1i, ky) for every surviving key ky.
9: Take average of //2 such repetitions, then median of 10 10log n such averages.
Set 2; be this median.
11: end for
Return
We now have all the parts state the proof of our main theorem, Theorem 1.1. The pseudocode of the complete algorithm is presented in Algorithm APPROXKMV.
Proof of Theorem I.I. We first use Lemma 3.2 estimate yи in time O(d n==-). We then let T [n] \ H. Next, we run Algorithm 1 FINDHEAVY to find sets S; for all i E [n]] Its cor- rectness is guaranteed by Lemma 3.4, and it runs in time Ö(d nl222). Finally rithm APPROXLIIHT on the set T and sets S; for all i E [n]. to obtain the vector = in time Ö(d - (n2+v-a + n (11784%/62) /c2). satisfies the guarantees as per Lemma 3.5. We then define the vector ØT E Rn as follows, (Dr); = z + Ejes, k(g., ky)ær, for all E [m]] and let y = Ør + yи- Thus we get that ||r y112 < |Kx yn yrll2 2 + I| yrll2 S 2e|[x]12, where we used the fact that (VT); (yr)al ) 22 Σier,s, k(gi, ky )ay & for all i [n]. We scale down by
run Algo-

---

Published as a conference paper at ICLR 2025
and set y 0.109, a 1/3 balance the exponents in the runtime, to obtain the overall runtime of Ödn' 1.d91.8999 /9/8//// /e²). A union bound over success probabilities gives the final success probability.
Algorithm 3: APPROXKMV
1: Input: Keys k1, k2,.. ...,, kn. Queries (1, 12..... Qn. vector v, parameter > 0.
2: Output: A vector y IR" such that |K yll2 S <||ar|12-
3: Let H C [ra] and ыи E R" be the output of Lemma 3.2 for y 0.109. Tet = |n] H.
4: Let St for all i [n] be the output of Algorithm 1 FINDHEAVY when executed for a 1/3.
5: Let 2 E R" be the output of Algorithm 2 APPROXLIGHT when executed for set T, sets S;
Vi E [n]. 7 = 0.109, a 1/3 and
6: Output the vector y E R" defined as yı + (UH) Eees, (ggrkk))a.
4 EMPIRICAL VALIDATION OF OUR MODEL
In this section we empirically evaluate our modelling assumption on the Gaussian matrices observed in the context of fast attention computation for transformer models. We start by introducing our main computation problem of interest: multiplying the dot product self-attention matrix by a vector, an operation that naturally arises in widely used Transformer models (Vaswani al., 2017). Consider a sequence of n tokens. For each token i there is key, query and value embedding denoted by k; q:: U/ E Rd respectively, for all i E [n]] We use Q K, V E Rx denote the query and key matrices whose jth rows are q. k2, Us respectively for all E [n.. Let A to denote the n x n
un-normalized attention matrix whose (i, entry is Ag)/a for all (i,j) [n] [n]. Thus A = exp((KT)) /d) where exp(.) entry wise exponentiation. Let D= diag(A denote the diagonal matrix containing the row sums of A on the corresponding diagonal entry. The main computational problem in self-attention is to compute -VV,, which naively takes Q(n2 d) time.
Consider the computational problem of computing the matrix-vector product Ax for arbitrary vector x R", When x = I., the all ones vector, then AI, will be the vector of row sums and thus can be used to compute the diagonal scaling matrix D embedding of each token U; we can compute Av, for all i E [n]] to compute AV. We will now use the following lemma to reduce this problem to instance of the problem we study Gaussian kernel matrix-vecto computation. Its proof is provided in Appendix A. We note that a similar reduction from attention matrices Gaussian kernel matrices was presented in Zandieh et al. (2023); the new reduction we give here is preferable, as it has better precision guarantees, and is also independent of the vector x being multiplied with the attention matrix (thus, our reduction need only be performed once per matrix, rather than once per
diag(AIn) Finally for the value
pair as in Zandich et al. (2023)).
Lemma 4.1. For any collection of vectors {ki)iz1> {q: H S Rd, there exists corresponding collection of vectors [k]] (q/))== S Rd+1 such that for any vector x E R".
te ky va = ellv:113. max 1I
Σ JE
zye
Σ JE(n]
Ize
2va
Vi [n].
This lemma and the discussion preceding it imply that we can use Gaussian kernel matrix vector multiplication algorithm to calculate Ax for any arbitrary x E R
We formalize the modelling assumption A and state it as follows,
For a of n keys and queries [k.) (o S Rdd consider the self-attention matrix A E Rax= defined as Ai,j = e k,)Va for all i,j [n]. Let (K) (z)) C Rd+1 be the set of keys and queries obtained after applying the reduction of Lemma 4.1, and K E Buxn be the Gaussian kernel matrix obtained from them defined as KiJ = e llo! 2/13/vva Then the ratio of ||]]] minus the sum of the top n entries of K and the sum of the top entries of K is at most a constant e independent of
(A)
To validate this assumption experimentally, we proceed as follows:
7

---

Published as a conference paper at ICLR 2025
take attention matrices computed in practice by Transformer model real data;
for each attention matrix and its associated keys and queries computed by the model, we apply the reduction of Lemma 4.1 to obtain Gaussian kernel matrix;
verify our assumption (A) for this Gaussian kernel matrix.
Evaluation methodology: We consider a pre-trained BERT base (uncased) model (Devlin et al., 2018), which is transformer based model pre-trained on a large corpus of English data. We use the Huggingface transformers library for our experiments (Wolf et al., 2019). This model has 12 layers with 12 self-attention heads per each layer. We then obtain the attention matrices from this model as follows- We consider all sentences obtained from responses of all questions in the public Stanford Question Answering Dataset (SQuAD) dataset (Rajpurkar et al., 2016), Our experiments are performed on the Google colaboratory platform's 's free tier version. For each sentence we use the tokenizer used in the BERT pre- pre-training to tokenize the sentence. Then feed this sequence of tokens into BERT and inspect all the self-attention activations across each layer. Our code is in the supplementary material. We also present additional experimental evaluation other models RoBERTa (Liu, 2019) and GPT (Radford et al., 2018) in the appendix Section A.1.
Fix sentence, suppose it has tokens after tokenization, and pass it through BERT. Then fix layer and an attention head in that layer. We obtain the key and query embeddings {ki }c} produced by this attention head. Then we use the reduction of Lemma 4.1 to produce the modified set of keys and queries { } that we use to construct a Gaussian kernel matrix denoted by K E RR=x" as described in A. To demonstrate Assumption A, we consider all principal sub matrices of K. More specifically, we consider K|: i, il for E [50, n]. This is natural for studying how our model scales with input sequence length as K[: i,: is the kernel matrix obtained from the prefix of the input sequence containing the first i tokens. We choose min prefix length of 50 so as to start observing asymptotic behavior. The maximum n goes up to is 512, the max context length of BERT.
Experiment (i). For prefix length E [50, ni]. we compute the sum of the top i largest entries in K|: i,; denoted by a, and we compute the sum of the remaining j2 entries in K|: i, which will be ||]] i, a]llı a, We then compute the max of [[] i,: a]]lı az)/a, over all e [n]] We then take the max of max [[: i, i] sentences we consider. We thus get accumulated max ratio over all sentences for each head and each layer. Figure! lists these accumulated max ratios per layer per attention head.
a,) ag over every sentence in the collection of
Max Ratio Heatmap
AGG
1.6
0.35
0.23
2.5
1.9
2.5
1.6
0.19
1.2
0.4
J.b
2.4
0.79
0.01/
2.8
2.8
0.5
1.8
Se-0n
0:35
3.0
0.99
0.6%
3.4
18:00
1.5
0.71
1.3
1.6
3.4
0.21
0,2K
1.6
0.88
9.0
0.45
0.3
3
0.73
2.2
0.97
0.23
3.5
0.11
0.64
0.24
0.48
1.2
0.33
0.33
Laye tarber
0.37
0.34
2.5
0.33
0.3
0.15
0.2
0.43
0.16
0.094
0.24
0.32
0.56
0.27
0.14
0.27
0.26
3.3
3,2
6.1
0.35
0.025
0.57
0.3
0.25
0.34
0.032
0.64
0.35
0.32
0.44
0.3
0.22
0.26
0.69
0.27
0.72
0:24
0.56
11
0.66
0.46
0.07
0.59
0.6
0.46
Heed Number
Figure 1: Statistics of max ratios.
Experiment (ii). Next, consider set of experiments to show that the large values in the reduced Gaussian kernel matrices after the removal of the largest n elements, are comparable to the values in the largest elements.
We consider the same collection of sentences from the entire SQuAD dataset as before. We fix a sentence, with number of tokens denoted by n after tokenization, and pass it through BERT. Then for each layer and each head we extract the key and query embeddings and construct the reduced Gaussian kernel matrix K using Lemma 4.1. Then we calculate the ratio of the nh and th largest.
8

---

Published as a conference paper at ICLR 2025
well of the n and (n 1)th largest entries of K, and take the median of these ratio across all all sentences. Thus we get two median ratios per head per layer. Figure shows a visualization of these median ratios the reduced Gaussian kernel matrices.
Moattnup
123
1.1.
1.51
42
140
2.7
1.1
1.1
1311
122
140
2.2
1.1
1.23
1.1
5:27
22
2.22
1K2
392
100
X
1.00
2:02
Laye
XK2
322
L.0
XC2
322
100
52
SE2
282
LAN
LSS
3
222
L22
232
242
2
A.2
LSD
L.L
LG
F2
15
11
221
L22
30:
100
2:2
LO0
L.2
1.0
137
140
1:27
1:22
1:30
1.7
1.01
Houd Humber
(a) Median ratio 1h and 2nzh largest.
(b) Median ratio of of rh
and (1 + 1) largest.
Figure Statistics of ratio of th largest with 2nth and (n + 1)th largest entries.
4.1 RESULTS
Experiment (i): From inspecting the numbers Figure across 12 layers and heads per layer, we observe that all of these are less than 4.6, and often significantly smaller. We interpret this as strong evidence that the constant e in Assumption (A) small, thus validating our model.
Experiment (ii): From Figure 3a observe that for most of the attention heads, the median ratio of th and th largest entries of the reduced Gaussian kernel matrices is about 20 or less" This implics that in most cases, the th largest and the 2nth largest entries have comparable value. Moreover from Figure 2b we observe that for almost all attention heads, the median ratio of n1th and (n + 1)th largest is about 1. The implication of this result is that we cannot rely the strong assumption, that after the removal of the largest n entries, there is small uniform upper bound on the values of the remaining entries the matrices we study. We interpret this as further motívation for assumption (A), which only assumes total sum of entries in the largest entries and the sum of the remaining entries after removing the largest comparable.
CONCLUSION
In this paper we study fast algorithms for approximate Gaussian kernel matrix vector multiplication motivated by the problem of fast processing of attention matrices encountered in modern language models.
Our results two fold, first do an empirical study of Gaussian kernel matrices derived from attention matrices the context of fast attention computation using pre-trained language models to arrive at modelling assumption that the sum of all but the largest n entries of the Gaussian kernel matrix is comparable to the sum of the largest n entries. This modelling assumption implies the sum of entries of the whole matrix scales linearly in the matrix dimension as opposed to worst case quadratic growth.
Our second contribution for these class of matrices that runs in time subquadratic in the matrix dimension. Our algorithm is only faster than previous algorithms but also can handle multiplying the matrix with vectors that can have negative entries, which was not possible with previous algorithms.
design provable approximate matrix vector multiplication algorithm
"The blank white entries correspond infinite entries due division by 0.

---

Published as a conference paper at ICLR 2025
A limitation of our work is that algorithms operate under structural assumption the input matrices naariles--amely,, of the linear growth of the sum of the entries in the matrix K. Although we provide empirical validation of this assumption, the set of matrices occurring in practice is very rich, and no assumption will model such matrices perfectly.
ACKNOWLEDGEMENTS
PI was supported by the NSF TRIPODS program (award DMS-2022448), Simons Investigator Award, GIST-MIT Research Collaboration grant, and Wistron Corporation. TW supported by Len Blavatnik and the Blavatnik Family foundation and by an Alon Scholarship of the Israeli Council for Higher Education. TW is also with Amazon; this work is not associated with Amazon.
REFERENCES
Josh Alman and Zhao Song. Fast attention requires bounded entries, Advances in Neural Information Processing Systems, 36, 2023.
Alexandr Andoni and Piotr Indyk. Near-optimal hashing algorithms for approximate nearest neighbor in high dimensions. Communications of the ACM, 51(1)111--12,, 2008
Arturs Backurs, Piotr Indyk, and Ludwig Schmidt. On the fine-gre tnee-graied complexity of empirical risk minimization: Kernel methods and neural networks. Advances in Neural Information Processing Systems, 30, 2017.
Arturs Backurs, Moses Charikar, Piotr Indyk, and Paris Siminelakis. Efficient density evaluation for smooth kernels. In 2018 IEEE 59th Anmual Symposium Foundations of Computer Science, pp. 615-626, 2018.
Arturs Backurs, Piotr Indyk, and Tal Wagner. Space and time efficient kernel density estimation in high dimensions. Advances in neural information processing systems, 32, 2019.
Arturs Backurs, Piotr Indyk, Cameron Musco, and Tal Wagner. Faster kernel matrix algebra via density estimation. In International Conference on Machine Learning, pp. 500-510, PMLR, 2021.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004. 20:4005044..5. 2020.
Moses Charikar and Paris Siminelakis. bashinng-baed-esttmaal for kernel density in high dimen- sions. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science, pp. 1032-1043.
IEFE, 2017.
Moses Charikar, Michael Kapralov, Navid Nouri, and Paris Siminelakis. Kernel density estimation through density constrained near neighbor search. In 2020 IEEE 61st Annual Symposium Foundations of Computer Science, pp. 172-183. IEEE, 2020.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems, 2021.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. International Conference Learning Representations, 2021.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs 04805.
2018. URL
http:://arxiie wwtp:///////a//e....eeeae
Leslie Greengard and John Strain. The fast Gauss transform. SIAM Journal on Scientific and Statistical Computing, 12(1):79-94) 7994,,994411 1991.
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near- linear time. The Twelfth International Conference Learning Representations, 2023.
10

---

Published as a conference paper at ICLR 2025
Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: towards removing the curse of dimensionality. In Proceedings of the thirtieth annual ACM symposium on Theory of computing. pp. 604-613, 1998.
Matti Karppa, Martin Aumüller, and Rasmus Pagh. Deann: Speeding up kernel-density estimation using approximate nearest neighbor search. In International Conference on Artificial Intelligence and Stafistics, pp. 3108-3137. PMLR, 2022.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self oeff-attention. In International Conference on Algorithmic Learning Theory, Pp. 597-619. PMLR, 2023.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations, 2020.
Yinhan Liu. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:190..11692 364, 2019.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. OpenAl Blog, 2018.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Jian Su, Kevin Duh, and Xavier Carreras (eds.), Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 2383-2392, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/v1/ D16- 1264. URL https: accantoo wow.hhooooooyo...oooo ogy org DD 6 1264.
Paris Siminelakis, Kexin Rong, Peter Bailis, Moses Charikar, and Philip Levis. Rehashing kernel evaluation in high dimensions. In International Conference on Machine Learning, pp. 5789-5798, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 04768, 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowic7, et al. Huggingface's transformers: State-of-the-ard natural language processing. arXiv preprint arXiv: 190.0377, 2019.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströnformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35-16, pp. 14138-14148, 2021.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Yang, et al. Big bird: Transformers for longer sequences. In Advances in Neural Information Processing Systems, volume 33, 2020.
Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. International Conference on Machine Leurning, pp. 40605-40623. PMLR, 2023.
A APPENDIX
A. ADDITIONAL EXPERIMENTS
We consider the experimental setup for two additional models, RoBERTa (Liu, 2019), which builds on BERT and modifies key hyperparameters, removing the next-sentence pretraining objective and training with much larger mini-batches and learning rates, and GPT-1 model by OpenAl (Radford
11

---

Published as a conference paper at ICLR 2025
et al., 2018). Again we use the Huggingface transformers library (Wolf et al., 2019) for loading the pre- pre-trained models and we consider their default configurations in the library both of these models have configuration of 12 layers with 12 self-attention heads per each layer and max context length 512. We consider all sentences obtained from responses of all questions in the public Stanford Question Answering Dataset (SQuAD) dataset (Rajpurkar et al., 2016).
For each sentence we use the corresponding Huggingface tokenizer to tokenize the sentence. Then we feed this sequence of tokens into the model and inspect all the self-attention activations across each layer.
A.1.1 SETUP
Fix a sentence, suppose it has n tokens after tokenization, and pass it through the model. Then fix layer and an attention head in that layer. We obtain the key and query embeddings {ks, qi produced by this attention head. Then we use the reduction of Lemma 4.1 to produce the modified set of keys and queries {K. } that we use to construct a Gaussian kernel matrix denoted by K E R as described in A. To demonstrate Assumption A, we consider all principal sub matrices of K. More specifically, we consider K|: i, i] for E [50, z). This is natural for studying how our model scales with input sequence length as [[: i,; is the kernel matrix obtained from the prefix of the imput sequence containing the first i tokens. We choose a min prefix length of 50 so as to start observing asymptotic behavior. The maximum goes up to is 512, the max context length of each model.
A.1.2 STATISTICS OF MAX RATIOS
For prefix length i [50, n]. we compute the sum of the top largest entries in K[: i,; d] denoted by a, and we compute the sum of the remaining ,2 entries in K[: i, i) which will be |[[] i, ||]I - a,- We then compute the max of [|] i, i|]l az)/a; aa over all i E [rx]. We then take the max of [ax (Il [|| [ i, :|]l az)/ a; over every sentence in the collection of sentences we consider. We thus get an accumulated max ratio over all sentences for each head and each layer. Figure lists these
accumulated max ratios per layer per attention head for both RoBERTa and GPT-1.
Mac Ratii Healmop Ca
Mac Ratti Heatmmap RoBIBITs
R
0.54
042
823
0.5
033
824
0.99
0.36
1
1
(a) Max ratio heatmap for GPT-1.
(b) Max ratio heatmap for RoBERTa.
Figure Statistics of max ratios.
From inspecting the numbers in Figure 3 across all 12 layers and 12 heads per layer, we observe that all of these are less than 5.15 for GPT and 5.39 for RoBERTa, and often significantly smaller. We interpret this as further evidence that the constant e in Assumption (A) is small, thus validating our model.
12

---

Published as a conference paper at ICLR 2025
A.1.3 SCALING OF THE CONSTANT WITH CONTEXT LENGTH
We perform additional experiments validate hypothesis that the constant in Assumption does not scale increasingly with the context length. For each of the considered models, BERT RoBERTa and GPT, we consider the following experiment.
Recall the setup of Section I.I.I. We consider context lengths starting from 50 to 512 in increments of 50. Then for each context length in this list, again let d; be the sum of entries in K[: i, i] and compute the max over (I K[: é, d]]l2 a;)/a; across all layers and heads, and then take the average and standard deviation of this over all sentences in the dataset. Thus for each context length i from 50 to 512 in increments of 50, we obtain a max ratio across all layers, heads and sentence prefixes of length in the dataset. The maximum is over all sentences that are of length least after tokenization. We plot these averages with standard deviations as the width of the error bars on the axis and the context length on the x axis Figure 4. We can interpret from the figures that the value of stays constant within range of the figures as strong evidence for our modeling assumption that e stays constant with the sequence length.
Scal rg behastar of BE17
Scaling behavir of far Butti'ss
Scaling bahavke ofe GPT
(a) Scaling of for BERT.
(b) Scaling c for RoBERTa.
(c) Scaling c for GPT.
Figure 4: Scaling of with context length.
A.1.4 SUBQUADRATIC SCALING OF RUNTIME
We perform an experiment that runs basic implementation of our algorithm based on the KDEFormer implementation of Zandich et al. (2023) for normalized attention approximation, and compares the wall clock times with exact attention computation. We select matrices Q, V < Ra from the Glo Ve word embeddings with batch size 8, dimension d 100 and set selection for different sequence lengths ranging from 4K to 16K in increments of 1K. Our goal is to approximate normalized self attention for a single vector e = Vl:, 1]. that is to compute D-1 Av for A exp(QKT /vd) each sequence length. At a high level the KDEFormer implementation of Zandieh et al. (2023) uses locality sensitive hashing approximate the contribition of heavy entries in the attention matrix to the attention computation. It then uses column norm sampling to find a those keys in the attention matrix to
Q. We consider this data
subset of keys, and only uses the columns corresponding approximate the residual component of the attention matrix. Our implementation builds on top of this to compute the empirical variance of the column norm sampling based estimator for each query used to approximate the residual component of the attention matrix. We then compute the average empirical variance across all queries, and take 1.5 times more samples in column norm sampling for the queries with empirical variance higher than the average. This reflects one of our main ideas to use adaptive sampling budgets for each row/query of the attention matrix.
We then compute the ratio of the wall clock time for our implementation and the exact algorithm for each sequence length. For each sequence length our implementation's 's parameters are such that the ratio of the error in approximating normalized attention and the & norm of is always within 0.1 + 0.05. This allows us to observe the runtime behavior across the difference sequence lengths under approximately fixed error. We report the ratio of runtimes of exact vs implementation as a function of sequence length in Fig. 5. The ratios of exact to approximate runtimes increases with sequence length, suggesting sub-quadratic scaling of our runtime.
13

---

Published as a conference paper at ICLR 2025
exact/our
3.0
2.6
rati of
2.4
4000
6000
8000
10000
12000
14000
16000
seq length
Figure Ratio of runtimes of exact vs our implementation as a function of sequence length.
A.2 FULL PROOFS
In the appendix we provide the full proofs of Lemmas 3.2,3.4, 3.5 and 4.1.
Proof of Lemma 3.2. Since we know that Σje(n) 2 a simple Markov bound implies that [Hi] < n 1--7, Corresponding to the entries Hi we explicitly calculate уи using its definition in Definition 3.1. To do this we need to explicitly calculate |Hi entries of K, which takes time n. | (dd) = O(d n² n²²)).
Next, since each entry in the matrix K has value most 1, have that (Kr yи yr); n n n 33 for all i E [n]. Thus |r yu yrll S n vn I <||x||2 since e = ((1) and ||el| vn.
Proof of Lemma 3.4. Fix any query 4s for i E [n]] and let μs = EE k(gs k,))/n, thus using Markov bound we get the following,
S S n² (mp4) = n 1+ us.
Consider T 10ne log n independent LSH hash functions hy...., hy H. per Lemma 2.1. Then for any key ky for j E S; we have the following,
P[3 E [T] s.t. he(gs) = (k,) 1 (1 (1-1/n") 10n" log
1/n 10
Taking union bound over all rows E [n] and at most n heavy points per row, we get that with probability least 1/n, S; can be recovered during query time by scanning the buckets that g: hash to for all i E [n]]
Now for any E [n]] let Lim (j E [n] k(g, ky) [2-
2 2- for {a log l)g(1/m:)
logn Ii.m- Then again by a Markov argument we know that [T,m < 2 for all i E [n]. Note that for any independent copy of the LSH hash function hy we have the following for all j L.m
and let L, = U 1/1g 1/14
P(h+(g:) = he (k2)) n
an-a(1-o11)) -a-l--n-------
wlew mm1 mm(1-(1))
Thus by linearity expectation we have that
Ell L h((g)) he(k)))l S L (2-m(1--(1)) 2n 2
for j. Thus again by linearity of expectation this implies that
Billje L, 3t e [T] s.t. he(g.) he(k)))]l < (on++0+ pa).
Thus we get that expectation the number of non-heavy points across all rows that we may have to scan due to collision is at most Σn= 1 +o(1) (u)) = Ö(n (t) since En Pi 17K1/n O(1). This also holds with probability at least 0.99 due Markov's inequality.
14

---

Published as a conference paper at ICLR 2025
This implies that in time T Ö(n we can hash all keys during pre- propprccsssing. Then for every row i. we can scan all the buckets that query qi hashes to across all repetitions and return the union of all keys 2y landing in the same bucket as qi satisfying k(q,, k2) 2 n ", As per our previous discussion we get that with probability 0.99, this scan will take us time
T Σ ((S)) + Ö(n (ot = Ō(n!
and we will recover S, for all i E [n]. For every row i, we will brute force calculate Ejes, 8 Fzk(gisk)) and this will take us overall time
Σ
S Sn
Σ 14 = Ö(n ta). 1=1
Next we present the proof of Lemma 3.5
Proof of Lemma 3.5. Let T, = T\ S, for every E [n]. Let Ky be the i element of K, and let K2 denote the 7th row K. For each row i, we will sub-sample every key in Le with probability 1 1/n (This can be done by sub-sampling every key with probability 1/n and only retaining those keys with index Lj). Thus define the following random variable X,, = n Kig with probability 1/n and otherwise, thus IE EEjE, X1y] = =EeE, By Kiy- Thus Var(X.,) S (n ne.zy K2;)² (a:))//n-- = n az K3-. Thus Var(E,er, Xa) I n (Ejer, k) Thus by Chebyshev's inequality Ejer, Xij EjEL, te with probability 0.9 for any fixed i if we take the average of 1On-(Ejez /²22 independent repetitions of Ejet, Xij If we take the median of 10 log(n) independent repetitions, then by Chernoff bound we get an estimator that is within ΣjEL, By Key with probability 1-1 /0n. Now by a union bound this holds for all rows with probability 0.9. The expected number of samples taken across all rows is
=00(
10 log n Σ n
Σ JET,
K3
-
=ō
n 1+7
ΣΣ
Σ
K2
(E(x)
iE|N
JEL
It can be seen that under the constraint that all Kij Sn L; and |K|| (m))
Kg K2nsnn Σ [Ki,
Kij 0
KE(N] JEL.
JEIN
Plugging this back into the expression on the expected number of samples across all rows and applying Markov's inequality, we get that with probability least 0.99 the total amount of samples taken is
Ön2 /e²).
What remains to estimate Ejet, 23 for each row i E [n]] to get the number of times we need to repeat the estimator for averaging to reduce the variance. We will do this using KDE data structure to estimate Σjer K3 and subtracting Ejes, Kł explicitly from the estimate for each i [m] We will do this as follows. Let B E [o, 1] be parameter. We will first do a convenient bucketing of entries in æ.
Rounding: First we will round the entries of x to the nearest powers of (1 + +)) for integers m in (oo11ee (n), log (n)). This covers all x E [n-A,n?], n?]. thus all j E T. Let Bn, = {j E T 2 E [1 + c) (1 + c) For every m E [-4log l (n)) and j E B, let
2) + e)m This implies the following for all E [n].
Σ
Σ JET
K2
2
<
2c Σ K2j03.
JET
JET
Estimation within each bucket: Fix an m [-4 +- (n), 7 log 4o (n)) Note that since ||=1| and for each B,, we have that 2² 2 (1 + e)m² have that Bmml S n/(1 (1 e)m Now for every
15

---

Published as a conference paper at ICLR 2025
Bm will create Gaussian KDE data structure with data set as {k, : E Bm), relative error parameter of n- failure probability = 1/n², and a KDE lower bound of (logee) B.nl This lower bound satisfies the following from the bound on the size of Bm
n (n)(1 + B) n² (og
Thus, the KDE data structure can be created and queried n times in time ((dn-(n2/c2200 200) = [ddn_.366+23 /.331.). This setting of the KDE lower bound implies that if for any row E [n]. the KDE value corresponding to this bucket less than this lower bound then its contribution is at most
Σ IjKij S
n Σ
K
jEBm
jeBm
2²
n1 (1 () Bml n lo)g((n) + e)mlBmll
logn
This implies that since there are at most Olog n) many buckets, ignoring the contribution of buckets with KDE smaller than the corresponding lower bound results in an additive error of in the end.
Thus without loss of generality we will assume that all buckets contributing to Σjer K2 for i [n] have contribution above the corresponding KDE lower bound. This implies in time Ödn' 1.346+28 / 0.346) we can output an estimate t, satisfying the following for all E [n].
Σ K sts S
Σ JET
2 K
+
Σ
35 K
JET
JET
We will use 4 Ejes, K3 +n as an estimate of Ejer, K3 This clearly an over estimate of ΣJEL, K2 from the guarantee on 4, and the over-estimation error will just lead to oversampling in the previous discussion. The additional number of samples we will take due to this oversampling due to the error
Ö((n/e) n
Σ Σ n, = Ö(n' 1847 JET
ΣEKH). TE(M)JET
Now we know that since Ki S 1 for all entries in K, we have that EiEin) Ejer K3 S Eijejn Kij On). Thus overall the additional number samples needed due oversampling caused by estimation error is 0 (e?) Thus combining this additional additive oversampling factor with the sample complexity bound of the equation 2, we get that the total sample complexity is
Ō + 1-2--8/62
The total time to estimate the sampling probabilities 1 0 8 (e2) we set 0.218. Plugging in these values, the overall runtime 2t n1.78+/e27) /e2).
1-346+29 0.334)) Balancing this with
We finally state the proof of Lemma 4.1.
maxje(m] || and let wy Append wy and 0 as (d + 1)th coordinates to ky and oj respectively to obtain k2, of Id+11 Then we can observe the following.
Proof of Lemma 4. I. Let
for all j [n].
10 10--111 13
ll. 1
mace celn
1R
(g..ky))
2vd
2v4
v/
Multiplying this with xy and summing up over all j E [n]. we finish the proof of the lemma.
16


Reviews:
[{"review_id": "CAVqT3LqKE", "reviewer": "Reviewer_oYAf", "summary": "A wide variety of applications in machine learning and computational statistics, including the Kernel Density Estimation problem and the Attention subroutine of transformers, reduce to the problem of kernel matrix-vector multiplication (KMV). This paper studies an important subproblem that features kernel matrices with 1-norm scaling linearly in n (as opposed to the worst case quadratic growth). The authors performed experiments that evidence the prevalence of such matrices in the context of large language models, which intuitively reflects the fact that each token in the input sequence has high correlation with few other tokens. For this restricted type of kernel matrices, a new algorithm in o(n^2)*poly(d, 1/eps) time is provided, improving on the previous best [Backurs et al. ICML2021] in terms of both efficiency and applicability (in certain parameter regimes).", "questions": "1. A few other subclasses of kernel matrices were proposed in recent years that correspond to easier KMV problems. One example is kernel matrices with stable rank constraints. How are those matrices compared to ones with weight distribution studied in this paper?\n2. Several variants of KDE problems with different levels of hardness are mentioned in this paper – KMV with vector entries being (1) all 1, (2) all positive or (3) arbitrary. For the relative-error setting, (2) is strictly easier than (3). Is this true in the additive-error setting as well?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper points out an important subclass of kernel matrices that assumes faster algorithms, while being general enough to encompass many real-life datasets. It in addition motivates and sets the stage for the interesting task of finding subquadratic algorithms for this subclass in other parameter regimes (e.g. low-dimensional) that get around the fine-grained lower bounds for general kernel matrices.\n2. This paper presents a separation between the additive-error model and the relative-error model. The authors give the first subquadratic algorithm for KMV that allows negative entries in the vector (wrt. the additive-error model), and provide formal argument on the inherent impossibility of achieving this in the relative-error setting (which is the focus of most previous literature including [Backurs et al.]).\n3.  For certain KMV instances to which both [Backurs et al.] and the new algorithm are applicable, the latter outperforms with higher efficiency.", "weaknesses": "1. The algorithm can be viewed as a reduction from the proposed KMV subproblem to the general KMV problem. The core techniques such as finding heavy keys using LSH and random sampling, are slight adaptations of the ones developed in [Charikar et al. FOCS20] (or earlier).\n2. Several fine-grained lower bounds for either the KMV problem or the Attention subroutine are known [Backurs et al. NeurIPS2017, Alman-Song NeurIPS2024, Alman-Guan CCC2024]. Showing how the new algorithm manages or fails to get around the known lower bounds might shed more light on the power and usefulness of the new subclass of kernel matrices.", "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": "NV4blctwB9", "reviewer": "Reviewer_Motp", "summary": "This paper designs an efficient algorithm for matrix-vector products of asymmetric Gaussian Kernel matrices $K \\in \\mathbb{R}^{n \\times n}$. This problem is motivated by a recent work (Zandieh et al., 2023) that replaces attention matrices with Gaussian kernel matrices. Therefore, the efficient algorithm for kernel matrix vector product of Gaussian kernel matrices can be used to help with the attention optimization problem. The quadratic complexity $O(n^2 d)$ in attention computation limits the efficiency of the large language model. Studying attention optimization is crucial and interesting.", "questions": "In Zandieh et al. (2023), their results (Theorem 3.5) hold for any arbitrary $Q, K, V \\in \\mathbb{R}^{n \\times d}$ to approximate the attention computation, and it does not seem that their Theorem 3.5 rely on the non-negative constraint. How does the result of this paper help with the attention approximation?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The most impressive strength of this paper is that, as mentioned by the authors, their algorithm is the first one that runs in sub-quadratic time for kernel matrix-vector multiplication for unrestricted vectors. It gives a more general solution than the previous works: Backurs et al. (2021) has to restrict the vector $x$ to be non-negative in order to obtain the running time $O(n^{1.173+ o (1)})$. \n\n2. Additionally, the empirical results from this paper are also very interesting and significant. It shows that the assumptions presented in this work do work well in practice in the setting of transformer-based large language models. It supports that the theoretical work in this paper is applicable to the transformers.", "weaknesses": "1. However, the experiment in this paper only includes a pre-trained BERT model, whereas the introduction claims that, “we show empirically that our modeling assumption holds for kernel matrices that arise in modern transformer-based language models.” Beyond the BERT model, I am uncertain whether this assumption would apply to other transformer-based large language models (LLMs) as well.\n\n2. While the experimental results support, to some extent, that the assumption made in this paper holds in practice, there are no experimental results verifying the improvement in running time. Since the main theoretical contribution is this improvement in running time, it would be preferable if the paper included experimental evidence to support it.\n\n3. Another work by Alman & Song (2023) employs the polynomial method to generate the low-rank matrices $L, R \\in \\mathbb{R}^{n \\times k}$, satisfying that the attention matrix $A \\in \\mathbb{R}^{n \\times n}$ is approximately equal to $LR^{\\top}$. This method only requires nearly linear time $O(n^{1 + o(1)})$ to approximate the attention computation $D^{-1} A V$. Although this method achieves better running time, it imposes a stricter assumption $d = O(\\log n)$ compared with $d = o(\\log^2 n)$ in Zandieh et al. (2023). Since this paper builds on Zandieh et al. (2023), I assume it follows the same assumption, $d = o(\\log^2 n)$ (if not, please point it out). It would be beneficial for the authors to include a more detailed comparison of this trade-off between running time and the constraint on the hidden dimension $d$. My concern here is that if existing work already provides a nearly linear time algorithm to approximate attention computation, why is there a need to develop kernel density estimation for attention computation in a less efficient running time?\n\n\nOne minor comment:\n\nAlman & Song (2023) was published in NeurIPS 2023, but in the paper, the authors cite it as Alman & Song (2024).", "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": "jU9TKef6Rx", "reviewer": "Reviewer_uVts", "summary": "The authors present a novel approach for accelerating Gaussian kernel matrices with general vectors, at a time complexity subquadratic in the number of tokens $n$. It is based on preprocessing the vector to explicitly calculate the contribution from its largest entries, finding the ‘heavy’ keys for each query using a hashing scheme, then randomly estimating the contribution from the remaining ‘light’ keys by uniform sampling. Under the assumption that the sum of entries in the kernel matrix grows linearly with matrix dimension, the method is $\\widetilde{\\mathcal{O}}(d n^{1.89}/\\epsilon^2)$ (where $\\epsilon$ is a parameter related to the quality of the approximation) with probability $p=0.99$. The authors are motivated by LLMs, so test their kernel assumption experimentally in this setting.", "questions": "1. What is the space complexity of your algorithm? I assume quadratic in $n$. This is not a reason for rejection, but running OOM is a real practical concern so it would be good to flag to readers whether or not you help here. \n2. Lemma 3.1 is applied to the attention matrix to obtain a related Gaussian matrix. Why not simply use $K_{S.M.} = \\text{diag}(\\exp(q_i^2/2)) K_{G} \\text{diag}(\\exp(k_j^2/2))$? Diagonal matrix multiplication is trivially linear. Your guarantees might need to be tweaked slightly. Would this be simpler, or are there good reasons to prefer the current formulation?\n3. Presumably one can trade off the time complexity and probability that the estimate is accurate, changing the exponent on $n$ and the value of $p$. Can the authors comment more on this? \n\nNotwithstanding my comments above, this is a nice paper. I would be happy to raise my score if the authors can test assumption A more carefully and give some wall-clock run times, or convince me that these things shouldn’t matter. I thank them again for their time and efforts.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Thanks for the nice paper. The work enjoys a strong practical motivation and brings technical novelty, especially improving upon Backurs (2021) and contributing to the growing literature on hashing-based algorithms for kernel computation on high-dimensional points. It is broadly well-written and validating your assumptions on BERT is a nice touch. I think the efficient transformer and kernel density estimation communities would enjoy this work.", "weaknesses": "1. *In places, contributions seem overstated*. The authors describe their algorithm as ‘the first algorithm for high dimensional approximate kernel-matrix vector multiplication, that runs in subquadratic time for general vectors’ (110), and similar elsewhere in the text. I agree that they make a nice contribution to fast linear algebra under assumptions on the structure of $K$ but this particular claim is overblown – e.g. taking a low rank decomposition to the kernel using random Fourier features achieves the same, though with a different set of assumptions and guarantees. I’d consider phrasing this a little more carefully to orient the contributions in the broader literature.\n2. *How practical is the algorithm?* It would be nice to actually see the algorithm in action with some wall clock times for real or synthetic data, to verify the subquadratic scaling with dimensionality. I suspect it will be slower than vanilla matrix vector multiplication for small $n$, but become faster at some sequence length. Roughly when does this occur? It’s fine if this is a big number, but I think it’s important to give readers a sense of how practical your scheme is. Subquadratic time complexity is not the same as a fast algorithm. \n3. *Do you really test assumption A?* My reading of your core assumption is that the ratio of the sum of all but the largest $n$ entries of K by the sum of the largest $n$ entries of K is at most a constant $c$, independent of the sequence length $n$. Great idea to test this on BERT, but to convince the reader I think you need to show how the maximum of this ratio changes as you vary sequence length, rather than just reporting its maximum value over all sequence lengths. I can’t see any evidence that it’s (approximately) independent of $n$. Another nice idea (which might be harder) would be proving whether your assumption holds under different data distributions (queries and keys that are uniform on a hypersphere, Gaussian etc.), to get a sense of whether we should be surprised that it empirically holds in LLMs or whether this is a general property of the Gaussian kernel.\n4. *Presentation*. This is a stylistic point. I would consider bringing Alg. 3 or a related schematic to Sec. 1.3. It feels a bit weird to wait until page 9 to see a clear presentation of your entire algorithm. I would also swap the order of Secs 3 and 4.", "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": "SRVv729zMN", "reviewer": "Reviewer_uqhw", "summary": "This work obtains a faster algorithm for evaluating matrix-vector products when the matrix is a kernel matrix which additionally satisfies a “heaviness” assumption. This assumption is newly introduced by the authors, and empirically validated for kernel matrices derived from attention matrices arising in transformers trained in practice. The techniques of this paper extend methods from the literature of kernel density estimation algorithms, largely based on hashing and sampling methods.", "questions": "- It may be worth emphasizing the new assumption in the title, since the title suggests an unconditionally faster algorithm for kernel matrix-vector multiplication, which is somewhat misleading.\n- How does the running time of the algorithm scale with the constant $c$ made in assumption A? Is the running time something like $\\exp(\\mathrm{poly}(c)) \\cdot dn^{1.89}/\\epsilon^2$? It would be good to mention this so readers can have an awareness of what happens if their matrix encountered in practice has values of $c$ that are a bit larger than what is assumed in this work (e.g. what if it’s 10? 50?)\n- Theorem 1.1: you state the theorem as a constant probability statement, but is there a way to boost the probability of this procedure with logarithmic dependence on the failure probability? Typically we need to compute many matrix-vector products, so this is an important point to consider. For example, training transformers would need a huge number of such products, and inference when serving the model could be even larger. It is also unclear how this subroutine would perform when used as part of an algorithm for power method (perhaps the authors might consider working out this corollary if they think it gives improvements?)\n- Line 151: should be $\\lVert …\\rVert_2^2$ rather than $\\lVert …\\rVert_2$\n- Do your results imply improved results for the special case of KDE? power method?\n- Maybe too much to ask for of the authors since this is largely a theoretical contribution, but it would have been more convincing if there were experiments which showed that this new algorithm could be used to speed up the KDEformer application of Zandieh et al. Again, that's ok though since this is a theoretical paper.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper works on a highly practical problem of speeding up kernel matrix vector products, which has applications to fast transformers, with improvements based on randomized algorithmic techniques such as hashing and sampling. The authors also find a nice structural assumption which helps in getting fast algorithms, while being a practical and realistic one, as the authors show.", "weaknesses": "While the result is highly practical, the authors unfortunately don’t take it all the way to evaluate the speed up implications of this result for transformers (although this is a lot to ask for).", "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)