kNN Attention Demystified: A Theoretical Exploration for Scalable Transformers
Authors: Themistoklis Haris
Venue: ICLR 2025 Poster (2025)
Keywords: efficient transformers
self-attention mechanism
sublinear algorithms
sampling
k-nearest neighbors
Primary area: foundation or frontier models, including LLMs
TL;DR: This paper proposes a theoretical framework for kNN attention and develops novel algorithms for sub-quadratic attention gradient estimation in Transformers.

Abstract: Despite their power, Transformers face challenges with long sequences due to the quadratic complexity of self-attention. To address this limitation, methods like $k$-Nearest-Neighbor ($k$NN) attention have been introduced [Roy et al., 2017], enabling each token to attend to only its $k$ closest tokens. While $k$NN attention has shown empirical success in making Transformers more efficient, its exact approximation guarantees have not been theoretically analyzed. In this work, we establish a theoretical framework for $k$NN attention, reformulating self-attention as expectations over softmax distributions and leveraging lazy Gumbel sampling [Mussmann et al., 2017] with $k$NN indices for efficient approximation. Building on this framework, we also propose novel sub-quadratic algorithms that approximate self-attention gradients by leveraging efficient sampling techniques, such as Markov Chain-based estimation. Finally, we demonstrate the practical effectiveness of these algorithms through empirical experiments, showcasing their benefits in both training and inference.

Published as a conference paper at ICLR 2025
KNN ATTENTION DEMYSTIFIED: A THEORETICAL EXPLORATION FOR SCALABLE TRANSFORMERS
Themistoklis Haris Boston University tharis@bu.en
ABSTRACT
Despite their power, Transformers (Vaswani, 2017) face challenges with long se- quences due to the quadratic complexity of self-attention. To address this lim- itation, methods like k-Nearest-Neighbor (ANN) attention have been introduced (Roy et al., 2021), enabling each token to attend to only its & closest tokens. While ANN attention has shown empirical success in making Transformers more efficient, its exact approximation guarantees have not been theoretically analyzed. In this work, we establish theoretical framework for ANN attention, reformulat- ing self-attention expectations over softmax distributions and leveraging lazy Gumbel sampling (Mussmann et al., 2017) with ANN indices for efficient approx- imation. Building on this framework, we also propose novel sub-quadratic algo- rithms that approximate self-attention gradients by leveraging efficient sampling techniques, such as Markov Chain-based estimation. Finally, we demonstrate the practical effectiveness of these algorithms through empirical experiments, show- casing their benefits in both training and inference.
INTRODUCTION
Transformer models have become the dominant neural architecture across language, vision, and other domains (Vaswani, 2017; Dosovitskiy et al., 2020). However, scaling them to handle larger input sequences remains significant challenge (Tay et al., 2020), primarily due to the quadratic complexity of computing self-atrention. Overcoming this limitation is crucial for advancing neural networks. Extending context length would enable Transformers to tackle complex tasks like book summarization (Kryšcinskii et al., 2021) and time-series forecasting (Wen et al., 2022; Zeng et al., 2023; Zhou et al., 2021). Furthermore, improving attention efficiency would reduce the computa- tional burden of training, making these models more accessible. Bridging this "compute divide" is vital for democratizing AI (Ahmed & Wahed, 2020).
Efficient computation of self-attention has been focal point of research in recent years (Foumier et al., 2023). Flash Attention (Dao et al., 2022) and related work (Saha & Ye, 2024) optimize the exact calculation of attention by minimizing wasted computation during GPU 1/O operations. However, most approaches focus on approximating the attention function. Sparse Transformers improve efficiency by allowing each token attend only small subset of tokens (Meister al., 2021). These subsets are identified through deterministic methods (Child et al., 2019; Guo et al., 2019; Soldaini & Moschitti, 2020; Li et al., 2019; Qiu al., 2019; Beltagy al., 2020; Chen et al., 2021), randomized algorithms (Kitaev et al., 2020; Han et al.. 2023; Zandieh et al., 2023; Pagliardini al.. 2024), or adaptive techniques (Correia al., 2019). Additionally, self-attention is often approximated using low-rank matrices and kernel methods (Wang et al., 2020; Tay et al., 2021; Xiong et al., 2021; Katharopoulos et al., 2020; Choromanski et al., 2020). On the negative side, recent fine- grainnd complexity reductions indicate that achieving good approximation with sub-quadratic time is not feasible across all scenarios (Keles et al., 2023; Alman & Song, 2024a).
In this work, we focus on sparse attention methods where each token vector gi E Rd attends to the tokens ky E Rd with the largest inner products of' k, (Gupta et al., 2021; Wang et al., 2022), a kNN Attention. The Routing Transformer (Roy al., 2021)
paradigm refer
early example, using k-means clustering to ensure each query only attends to keys within the same cluster. Memorizing Transformers (Wu et al., 2022) later extended this approach by leveraging LNN search

---

Published as a conference paper at ICLR 2025
within stored memory, enabling models to memorize data during inference. More recently, Unlimiformer models (Bertsch et al., 2024) have improved efficiency by using a single kNN data structure (or index) across all attention heads and layers.
Previous works have empirically shown that ANN Attention not only improves computational ef- ficiency, but also enhances model architectures and capabilities. However, rigorous theoretical analysis of kNN Attention is still lacking. Key questions remain unresolved, including the precise approximation guarantees offers, the optimal value of k; and whether these methods can also help to yield efficient algorithms for approximating the backward pass.
For comprehensive outline of preliminary results and theory, please refer Appendix A.
Notation Let Q, K, E Bxnd be our query, key and value matrices. Let qi Qi E Rd be i-th row of Q written column vector. We will also denote the j-th column of Q by Q..- We define A QK" E "ax= to be the attention matrix, and O softmax(A) V E ig to be the output of the attention function. The softmax function applied row-wise to and is defined as a vector valued function R R":
exp(y:
(yi..... Σ (yp(
(1)
We also let [n]] == (1, and use the notation polylog(1) substitute of log" (n) for some arbitrary constant & Z that is independent of Finally, we use the O notation to hide polyloga- rithmic factors. We design randomized, Monte Carlo algorithms that can fail with probability d > O. We will often make use of the following boosting lemma:
Lemma 1.1 OMedian-ff-Meam Boosting, Chakrabarti (2020)). If an unbiased estimator of some statistic, then one can obtain an (E, (e.8))-mulillicaatie estimate of that statistic by suitably com-
Var Q In 3/2 independent samples of Q, where C is a universal constant.
bining K == O
1.1 OUR CONTRIBUTIONS AND RESULTS
Theoretical Framework for ANN Attention Our work provides a theoretical framework to explain both the efficiency and effectiveness of ANN Attention. Our framework reformulates self- attention as expectations over softmax distributions. These expectations are approximated by sam- pling from each distribution in sublinear time using Lazy Gumbel Noise Sampling. By connecting ÀNN, k-Maximum Inner Product Search (MIPS), and Gumbel noise sampling, we develop a new sub-quadratic self-attention approximation algorithm aligning with the ANN Attention paradigm with precise theoretical guarantees, given in full as Theorem 2.4.
Approximating the Backward Pass Our framework can be extended to solve the problem of approximating attention gradients. Even though backpropagation the main memory bottleneck for large models, few methods approximate attention gradients directly. The work of Alman & Song (2024a) most relevant, deriving inapproximability results for certain parameter regimes.
We present new approximation algorithms for self-attention gradients using &NN search. Our results are encapsulated in full by Theorem 3.1. key challenge is the need multiply by the transpose of stochastic matrix, which disrupts our expectation-basee reformulation. To address this, we use Markov- Chikn sampling technique, treating the attention matrix as a transition matrix and applying single-step iteration. We also employ other sampling techniques, such as CDF-based sampling, in novel ways to achieve sub-quadratic time complexity.
ANN ATTENTION AS AN APPROXIMATION ALGORITHM
We begin by viewing the self-attention output as a matrix expectations under various softmax distributions. This reformulation has been used in prior works (Kratsios, 2021; Singh & Buckley, 2023) but, to our knowledge, has not been used for analyzing sparse attention mechanisms. Let D, be the softmax distribution defined by Di(j) x exp(1] ky) over [n]. Then, notice that we can
2

---

Published as a conference paper at ICLR 2025
write:
Oig = Σ expla kr
Vi = Σ Di(r) V4J = E [Vry] mwD
(2)
e=1 exp(1 k.)
Thus, to approximate Oij. we have to estimate the expected value in Equation Let r be sampled according to D;- Then, the estimator Oij Vrj is unbiased, as EE-D, [Oig] = Oij- Assuming an upper bound on ||V and that the values Oi do not get arbitrarily small', we can bound the variance of Ozg and use boosting to obtain explicit multiplicative error guarantees:
Theorem 2.1. Suppose ||V I B = (log((n) and assume that for any i E [ni] we can sample from Dj in expected time O(T). Further, we assume that for some constant C, it holds that Ois 2 C for all j. Then, there exists an algorithm to oulput a matrix Ô E IRaxd such that:
O =Oi
(3)
for all (i. E [n] x [d] with probability at least 6, where E, are constants. The algorithm runs in O(nd T logg (ndd) logn) time in expectation.
Proof. Given that Ozy is an unbiased estimator of Osj, we can utilize Lemma 1.1 to get an (E, multiplicative estimator for Oij- To determine a sufficient number of samples of Ōij. we first bound the variance of estimator:
Var [.0
[ Σ
Di(r) - BBO..
(4)
Then, the number samples required is:
o ( 2.
2. log((15)) Var [
E [.s] 2) 0 0s (s-2 log((10)) log
(5)
due to our assumption ||V|| and the lower bound on Oij- To ensure that all nd elements of are approximated within the desired guarantees, we have to set 8' == 6/(nd) and union-bound over all nd elements of O. Since each sample requires (T)) time, arrive at the desired time complexity.
2.1 EFFICIENT SAMPLING FROM D; VIA LAZY GUMBEL SAMPLING
Theorem 2.1 previously assumed we could directly sample from the distribution D, in time O The Lazy Gumbel Sampling method proposed by Mussmann et al. (2017) provides way to sample from each D, in sublinear time, even with limited knowledge of D;- However, there an initial pre-processing step that takes a bit more than linear time across all the distributions. In this section we present this method and clarify value for T be used Theorem 2.1.
Fix some i E [n] and let Zij = q. ky_ In the Gumbel Max Trick (Lemma A.3), we form the random variables N.j = Zij + Gig where Gi ~ Gumbel(O, 1) for all j E [n] and sample arg max Nıj- This is equivalent to sampling (2017) observed that if we have the top Ao Zıj, values in a set S, and add Gumbel noise to just them, then for any j S, to be ultimately picked, its Gumbel noise G.y must be quite large. We can the concentration properties the Gumbel distribution to argue that in expectation we only need to sample elements not in S, Setting &o = /n allows to balance the two, resulting in sublinear time algorithm for sampling from D, An illustration of the idea can be seen in Figure 1, as it was presented in Mussmann et al. (2017). We can see that this method samples exactly from D..
[n]] from the softmax distribution over the Zij scores. Mussmann et al.
Theorem 2.2 (Correctness of Algorithm 1, (Mussmann et al., 2017)). After running Algorithm it holds that:
i arg maxig ky + Gij} JEIm|
(6)
where Giy ~ Gumbel(O, 1). In other words, 3 is sampled according to D,-
'If Ory is arbitrarily close 0 then can achieve low additive error by just outputting Ōi 0.
3

---

Published as a conference paper at ICLR 2025
Algorithm Lazy Gumbel Sampling from D;, for some i E [nt,, (Mussmann al., 2017)
1: Inputs: & E N. qi E R2, K E Raxd S == {the k keys j with the largest Zy == of k2).
2: Sample Gig Gumbel 1) for j E St
3: Let .M <- max Ziz + Giy) and Smin < min Ziy}.
JES:
JES;
4: Let B + M Mmn be the Gumbel cutoff.
5: Let m ~ Bin(n - k, - exp( ((x-xexp(-))) be the number of [n] \ S Gumbels greater than B. Sample points from [n]] \ S; and denote the set of sampled points T,
6: Sample Giz Gumbel (0, 1) conditionally greater than B for each j E T.
7:
return j c arg max (Zi + Giy) jES,UT,
Proof. The only way that we do not find the maximum is if one of the points in [m] \ (S; U T.) are the true maximum. However those points (by construction) have Gumbel noise at most B, so they cannot be the overall maximum.
In Appendix B, show that the expected num- ber m of large Gumbels is at most n/k. Our simplified proof uses the Gumbel distribution's 's Moment Generating Function, rather than the original exponential-based analysis.
Lazy Gumbel Sampling Illustration
100
8
Lemma 2.1 (MMussmann et al., 2017)). The
lowing holds:
Noley a7k,
80
Set
E [m] S
Due to Lemma 2.1, we see that we need to set A /n to optimize plexity. As a result, by combining the pseu- docode of Algorithm and Lemma 2.1, Muss- mann et al. (2017) arrive at the following:
70-
overall time com-
0
Index/EIn]
Figure Lazy Gumbel sampling
Theorem 2.3. Let ===v/ Vn. Suppose that we are able to retrieve the set S; in [nn, k) time.
Then, we can use Algorithm 1 to sample from D, in O( + f(n, vT)) time in expectation.
2.1.1 OBTAINING THE TOP i INNER PRODUCTS
Algorithm relies on obtaining the set S; of the top / inner products q k, for each i [n] in sub-quadratic [n, v time. Since the k, vectors are fixed, while the vectors act as queries, this setup is known as the k-Maximum Inner Product Search Problem (MIPS)
The A-MIPS problem can be reduced to the KNN problem using a transformation proposed by Neyshabur Srebro (2015). We add an extra dimension to normalize all key vectors. Specifically, the inner product of k- can be expressed as:
1
of ky
2 (114 + 1k 11:
ky
If the norms ||k are the same across all j. the problem reduces finding the k nearest neighbors to g. To enforce this, we define:
(K))=== kg /M 11k,
so that [1(*),lle M for all j [nll, where M is previously known upper bound. When querying with qs, we use:
(g.) [ooT
(10)
4

---

Published as a conference paper at ICLR 2025
This transformation preserves the inner products, allowing to solve the kNN problem for g. We can then use a KNN index to preprocess K and query it with each of to construct S, for all E [n]- We remain agnostic to the specifie &NN index one could use for this algorithm', but we assume that the construction runtime is slightly larger than linear and the query time slightly larger than k then Theorem 2.1 with T O(yn) give total runtime of zo O(d . log(1/8))
Algorithm ANN Attention
1: Inputs: Q, K, E Raxd error parameter >> 0,, confidence parameter 8 0, E N.
2: for j [n] do
Pre-Processing: Lines 2-4
3:
(RT)*=
M - 11y1 (d+1)x
4: +- Build ANN index from {Kg 1 E [n]}
5:
for i E [o]do
6:
(11)² + [n," [0] E Rd+1
7: Query H with of to get S; with [S;] = k.
8: for j E [d] do
9:
O.j f- Median-Of-Means with Algorithm as sampler <- (k, qs. K, S.).
return
Overall, our framework for &NN Attention be summarized by the following theorem:
Theorem 2.4. Let Q, K, V E Rnxd and e, 5 (0,1) be positive constants. Assume ||V| Ologg)). n). Then, Algorithm with /n outputs matrix Ô Rn such that:
Ö, Ozyl S =Oij
(11)
for all (i,j) [m] x [d] with probability at least and in expected sub-quadratic time and space.
Proof. Algorithm approximates each entry of the attention matrix by using Median-of -e-neans sampling. Since nn, the sampling from each softmax distribution is done using Algorithm according to the guarantees of Theorem 2.3. Theorem 2.1 thus gives us the desired runtime.
2.2 ANN ATTENTION WITHOUT BEDIAN-OF-MEANE
This section describes a simpler algorithm for computing the expected value needed for self- attention. The algorithm still uses A:N indices to find the top l: inner products per query, but usually outperforms Algorithm 2 in practice, due to its amenity for hardware-acceleraa vectorization, and thus our preferred implementation for experiments..
Building on Mussmann et al. (2017), we estimate E&-D, [Kks) using set S; by sampling & additional vectors outside S; (set T;) and upweighting them in the expectation sum, follows:
θ, Exes, oa k, Voj + E.er. ea Vej
(12)
E.Es, es pー E.ET. eook, k,
The quality of this estimator and the optimal choices for k and are derived as follows:
Theorem 2.5. Assumming ||V| S C, the estimator O.ş satisfies the following error guarantee with probability at least 8:
O. - [a)-Ou|se0 Se
(13)
if the following two conditions hold: 2 2 8n²=-2 log (4/5) and kl 2 2nE 2 (2/6). Setting O (n2/33- Vog(( /lg((1/1)) gives us an Ō Vogg// algorithm for estimating self-atiention within additive error OEE), assuming efficcent kNN implementation.
Foor a specific construction with precise thcoretical guarantees that uses Locality Sensitive Hashing (I.SH), please refer to Appendix C.
See Appendix G for PyTorch implementation of this algorithm.
5

---

Published as a conference paper at ICLR 2025
Proof.
The proof of the additive error guarantee be found in Mussmann et al. (2017).
3 ATTENTION GRADIENT ESTIMATION VIA MARKOV CHAIN SIMULATIONS
Next, we present randomized algorithms which can efficiently approximate the gradients of the self- attention function. First, we give exact formulas for the gradients in question. These can be obtained by applying the chain rule repeatedly, as shown in the proof of Lemma 3.1, in Appendix D.
Lemma 3.1 (Attention Gradients). Let Q, K Ruxd Let P == softmax (QKT) R be the normalized attention matrix. Let be differentiable scalar function of O and D = 06/00 E gaxd, Similarly define D D and DV. The following relationships hold:
DV pT.Do
(14)
D
Σ
Pik (D - Pt.)) Kaj
(15)
k=1
D
Σ
Pki (Dfa -D Pa.:)) Qki
(16)
k=1
where D.J := 00/0P; (D., Vg.).
Clearly, calculating DR, D and DV can be done O(d2²) time. We propose algorithms for estimating these gradients that run in sub-quadratic time. We first present Algorithms 3 and for the estimation of D whereas our algorithms for estimating DK and D' can be found in Appendices F and E. Ultimately, we prove the following theorem:
Theorem 3.1. Let & be a differentiable scalar loss function and 0s/00 E Rand Then, under certain assumptions the 11 norms of Q, K, V, DO, there exist sub-quadratic time algorithmms that output estimates De, DD DV E =Rn for which with probability at least 8 holds that: 1 0a/0Q1lso I Eq: 11 06/0K|| S ex and 115V - 06/ /a eev (17) where eq, ek, ev are explicit error parameters that can roughly be bounded by O(snn((lylllln polylog(n)).
Proof. We combine the guarantees of Theorems 3.2, F.1 and E.1, where the assumptions on the 11 norms of Q, K, V and Do can also be found.
3.1 ESTIMATING DV
Theorem 3.2. Given K,K,V V and DO, Algorithm calculates 06/8V.j D for all (i,5) e [n] x [d] within an additive approximation error of
= (D.,, 1n) + 2neMj; where Mj == min
Do
(18)
with probability at least 1 1. The time complexity is O(nde= 2 lo.
Proof. Suppose we want calculate the j-th column of DV:
D D.O
(19)
for j E [d]. Fix zj Do Raxl and suppose that æjj 0. We will soon explicitly relax this assumption. Then, yj == ボj #//||az /|55| is distribution over the universe [n]. Imagine random walk over [r] with transition matrix and initial distribution Øj. Then:
(20)
is the distribution after step in the process. Thus, we estimate क.j with Markov Chain simulations, by first picking an item E [n] from the distribution Tj. and then picking another item k E [n] with probability Pak. We make N independent length-1 random walks like this and let:
1, the s-th walk ends up in state e otherwise
((1*))
{o.
(21)

---

Published as a conference paper at ICLR 2025
We know that E ()) for all 8 E [N]. Thus, we can form boosted estimator (e) N E
s=1 x(,,), This estimator is unbiased due to linearity of expectation, so we can use the Hoeffding bound ensure that our empirical distribution close the true distribution as long as we take enough samples:
Pr (1)(()) - (T))) 2 E| S 2exp(-22ee
We set the probability of failure to 1/(dn²) so that we can union bound over all v E [x]] and all j E [d]. It follows that we require:
N O 2 In(nd))
(22)
Of course, we need to scale Tī back to recover D.,, We define:
DV | P.
(23)
Then we get that with probability least 1 1/n holds for all j E [d] that:
|D
D.y
| |1 11
oo S 8 ا
(24)
Relaxing the non-negativity assumption We now relax the non-negativity constraint on zj. ac- cepting some approximation error. Since normalizing zj with its LI norm fails zj has negative entries, we adopt a numerical stability technique to ensure we get valid probability distribution even when Ij contains negative entries. Let
Mj := min
(25)
VE E
(5)).
be the absolute value of the most negative entry of Zj.
2 then set Mj 0. Now, M₃ Mj 1" E Rax1, then n.j z. + Mj 2 0. Therefore, we can estimate
If
7, x バ : pT
(26)
using our Markov Chain method. Going back to our original goal of estimating J. we have:
Jj = pT zj = pT. (z) M2)
(27)
T PT M, J Mj PT.
(28)
where 1 is the all 1-s vector. So then we only need to additionally estimate pT 1", This can be done in the same fashion only once, as pre-processing step. Specifically, suppose that we estimate pT 1mt 3. We know from our prior analysis that using ( logn) random walks we get an estimate 3 such that:
g pT 1"d loo Sen
(29)
Putting it all together, our final estimator is then:
PJ = P Mj 2
(30)
Eventually, the total error for estimating ア.j becomes:
DD DX
10 M1 5 T + My Pr, 11"
(31)
< 10 / + M₁ ||TTT 1t - 511
-311
(32)
< 8 | + eM n
(33)
= E E 1n) + 2en.11
(34)
where the first inequality follows from the triangle inequality and the last equality follows from
Σ (1z)) + M₁1 Σ (1))) + Myl = (zj, 1n) + nM,
(35)
7

---

return E' M Published as a conference paper at ICLR 2025
Algorithm Estimating Pd x, with query access to P E RRxx stochastic matrix
1: procedure APPROXPOSPRODGP E Rux", I 2 0, > 0)
2: Let N 2 = and > + (z, 14)
O(n) time
3: Let 0 E R be our output.
4: for [N] do
5: Sample E [n] with probability x æ; using Σ as normalization factor.
6: Sample k E [n] with probability Pix.
7: 3k 3k ++
8:
return N -
9:
procedure ESTTMATEERODDCTT E RRxx SEERR Re 0,8 RR)
10: Let M <- minveen).e.500
O(n) time
11: Let x' + M 1"
O(n) time
Call ,PPROXPOSRRD(PPE"..)))
Ö(n ?) time
13:
Algorithm 4 Estimating DV
Input: K, D [@a error parameter
2: Let DV E Raxd be our output.
3: 3 *=APPROPPOPRRO(CPP 1",
Pre-Processing
4: for j E [d] do
5: DVG ESTTIIATTEPROOUUCCC,,. D.
6: return DV
Runtime analysis For each j E [d] we take N = O 2og n) samples. We can take one sample in (nd) time. In addition, we must pre-calculate the sums (æ,, 1") + nMj nM, + Σ DO for j [d], which takes ((nd) time.
Remark 3.1. Note that Algorithm 4 does not materialize the P matrix. Instead it accesses its elements by using Q and K in (d) time per element.
EXPERIMENTAL RESULTS
In this section we present our experimental results. Through them we can interpret our theoretical framework better and solidify our understanding of it.
4.1 FORWARD PASS APPROXIMATION QUALITY ON RANDOM INPUTS
We begin by evaluating the effectiveness of ANN Attention in approximating the attention function. We randomly sample matrices Q, K, V Rnd from uniform distribution over [-B, B nnd and assess the approximation quality on these inputs. Our focus is on the &NN Attention estimator we develop Theorem 2.5 with X 1, as used implementations like Bertsch et al. (2024) and Wu et al. (2022). We vary & to study how the absolute error decreases as l: increases and compare the efficiency to the naive O(n²) attention, expecting notable performance gains. This experiment implemented in Py PoTch,, rumning MacBook Air with M3 CPU and 8GB of RAM.
Efficiency of KNN Attention Our experiments confirm ANN Attention's superior speed, demon- strating sub-quadratic scaling. With batch size of and 10 attention heads, it handles self- attention for n 106, while the naive method runs out of memory beyond n 2 20000. Increasing further leads memory errors for n 50000, highlighting kNN Attention's 's memory efficiency. Detailed results are in Figure 2(a).
Role ls in the Approximation Error We investigate the impact & on the approximation error, predicting that error increases k decreases. The experiment confirms this, showing that for
8

---

Published as a conference paper at ICLR 2025
Efficiency Altention vs Brute Force
KNN Allention Error
Na ve (1
4000
K0M 12 V00
8
3000
1000
0.03
0.00
20000
40000
60000
80000
100000
logint
n/
Isput SequArTcR Length
(a) &NN Attention vs Brute Force. We are able increase the context length by a factor of without different values of B. As grows, the error becomes running out memory.
(b) Mcan error of KNN Attention function of k for
negligible. some cases, vn is too big a threshold.
Figure 2:
Experimental evaluation of the approximation quality and efficiency of KNN Attention on randomly generated input matrices Q, K, V.
A > 1/8 the error is minimal. Our theory suggests threshold closer to v which indicates that the optimal may vary by dataset. The results appear in Figure 2(b). Interestingly, the error is more pronounced for small values of both B and k, potentially due to the limited approximation power when & is small. For larger k, this difference becomes negligible.
4.2 BACKWARD PASS APPROXIMATION QUALITY
Next, we evaluate the quality of our algorithms for attention gradient estimation. We sample Q, K, V from normal distribution, as this strategy aligns with typical neural network weight initialization strategies, and approximate Di and using randomized techniques. Our goal is to assess the error introduced by the approximation and whether this error causes gradient descent to converge far from the minimum.
We set the sequence length N 100 and the embedding dimension d = 3. The learning rate is varied between 0.05 and 0.5, while the error parameter is fixed at E = 0.05 and the confidence parameter at 5 0.1. We experiment with both convex (Mean Square Error) and non-convex (Cross Entropy) loss functions to examine how approximate gradient descent behaves, using Py PyTorch's 's autograd to compute the exact attention gradients. As shown in Figure our approximation closely matches the expected results in the convex case but deviates from the optimal convergence in the non-convex case. more detailed investigation of the impact of gradient approximations in large language model (LLM) training is left for future work.
4.3 EXPERIMENTS ON LLMS
Finally, we experiment with incorporating KNN attention into LLMs to study its impact on training and inference. While previous work has explored KNN indices for efficient fine-tuning and training (Bertsch et al., 2024; Wu et al., 2022), our goal is to understand how Transformer LLMs respond attention function approximation, linking to our theory and providing practical guidelines. The architecture and training methods are adapted from nanoGPT (Karpathy, 2022), and our experiments are conducted on an NVIDIA 140 GPU with 48GB of memory.
Our first experiment trains a mini character-level Transformer on small Shakespeare dataset, re- placing attention with ANN attention. The sequence length is set to = 1024, the embedding dimension to d 768, learning rate to 0.005 and we perform T 5000 iterations. We compare training and validation perplexity to the exact method. Results in Figure 4 show that &NN Attention maintains small perplexity gap. However, overfitting occurs, the perplexity difference widens, possibly due to increasing maximum approximation error. Future work could explore the impact of larger l values on perplexity.

---

Published as a conference paper at ICLR 2025
Loss: Cross Entropy
gradieet
2.50
Na gr gradients
No vo gradienes
2.25
2.30
10-
1.75
0 0.0
1.50
1.25
0.6
1.00
0.4
0.15
0
100
200
300
400
500
00C
700
200
400
000
800
1000
1200
1400
Baration
Baration
Figure 3: Gradient Descent with Approximate Gradients against different loss functions Φ. Even with approximate gradients, gradient descent still makes adequate progress towards convergence.
We also experiment with fine-tuning large pre-trained L.M using &NN attention, the hopes that the approximation will not severly degrade the model's 's quality. We manage to fine-tune GPT-2 XI. on a Shakespeare dataset- a task typically infeasible on a single L.40 GPU due to memory constraints with quadratic attention. Examples from prompting this model can be found in Appendix H.
Perplexity of Asproximation Attention Throughout Training
Attention Approximation Error Throughout Training
1NN ix 51 Train
Mcan Error
4.0
Exact Vol
2.5
KNN 5)
2.0
1.5
500
2.5-
0.5
2.-
1000
2000
3000
4000
1000
1500
2000
2500
3000
3500
4000
4500
Trinning Iteratio
Training Iteration
Figure 4: The perplexity and approximation enror of kNN Attention throughout training
CONCLUSION
In this work, we developed a theoretical framework for leveraging &NN techniques design effi cient and effective Transformer architectures. We extended this framework by introducing Markov Chain-based methods to propose novel algorithms for efficient self-attention gradient computation. Empirical validation on both synthetic inputs and real-world datasets demonstrated that kNN ap- proximations closely match the original performance in LLM training, while significantly reducing computational costs during both training and inference.
Moreover, our research opens several avenues for future exploration. Key questions include the effectiveness of training with approximate gradients compared to exact ones, particularly when error distributions are tightly concentrated but unknown. Another open question is about explaining the practical observation that the optimal À value often significantly smaller than the predicted vn. Overall, an improved understanding of the theoretical guarantees of kNN Attention can help solidify its place among other sparse attention techniques that provably attain close linear time, while maintaining approximation quality.
In conclusion, by applying sublinear algorithm techniques, our work provides solid foundation for making Transformers more scalable and efficient while identifying critical areas for further research LLM approximation.
11.61B parameters, Radford al. (2019)
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
To aid in the reproduction of our experiments, we include our code in the repository at https:// gi thub. com/ sansui 123 /knn con_aateention For our theoretical contributions, full proofs of our claims and analyses of our algorithms can be found in the Appendix.
ACKNOWLEDGEMENTS
We thank Andrea Lincoln, Thien Nguyen, and Krzysztof Onak for valuable discussions at the early stages of this project. We are also grateful to Brian Kulis and Esty Kelman for their invaluable feedback on the manuscript and to Nikos Georgoudios for his assistance with the MGF proof of Lemma 2.1.
REFERENCES
Nur Ahmed and Muntasir Wahed. The de-democrativation of ai: Deep learning and the compute divide in artificial intelligence research. arXiv preprint arXiv:2010.1558/ 2020.
Josh Alman and Zhao Song. Fast attention requires bounded entries. Advances Neural Information Processing Systems, 36, 2024a.
Josh Alman and Zhao Song. The iin-grained complexity gradient computation for training large language models. arXiv preprint arXiv:2402.04497, 2024b.
Alexandr Andoni, Piotr Indyk, Huy L Nguyn, and Ilya Razenshteyn. Beyond locality sensitvve hashing. In Proceedings of the twenty-fifth annual ACM-SIAM symposium on Discrete algorithms, pp. 1018-1028. SIAM, 2014.
Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn, and Ludwig Schmidt. Practical and optimal Ish for angular distance. Advances in neurul information processing systems, 28, 2015.
Iz Beltagy, Matthe E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/55 2020.
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley, Unlimiformer: Long-range transformers with unlimited length input. Advances in Neural Information Processing Systems, 36, 2024.
Amit Chakrabarti. Data stream algorithms lecture notes, 2020.
Rajan Chattamvelli and Ramalingam Shanmugam. Gumbel distribution. In Continuous Distribu- fions in Engineering and the Applied Sciences-Part II, pp. 263-271. Springer, 2021.
Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Skyformer: Remodel self-attention with gaussian kernel and nystr\" om method. Advances Neural Information Processing Systems, 34:2122- 2135, 2021.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.1053 2019.
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. arXiv preprint MiXiv:2009.1479. 2020.
Gonçalo M Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. arXiv preprint arXiv:1909.0m 00015, 2019.
Tri Dao, Dan Fu, Stefano Ermon. Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 35:16344-16359 2022.
11

---

Published as a conference paper at ICLR 2025
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. / 1929, 2020.
Quentin Fournier, Gaétan Marceau Caron, and Daniel Aloise. A practical survey on faster and lighter transformers. ACM Computing Surveys, 55(14s):1-40, 2023.
Aristides Gionis, Piotr Indyk, Rajeev Motwani, et al. Similarity search in high dimensions via hashing, In Vidb, volume 99, pp. 518-52888 518-529, 1999.
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star- transformer. arXiv preprint arXiv:1902.0911 2019.
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention. arXiv preprint arXiv:2106.000 2021.
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David P Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. arXiv preprint arXiv:2310.05865 2023.
Iris AM Huijben, Wouter Kool, Max B Paulus, and Ruud JG Van Sloun. review of the gumbel- max trick and its extensions for discrete stochasticity in machine learning. IEEE transactions pattern analysis and muchine intelligence, 45(2):1353-1371, 2022.
Andrej Karpathy. NanoGPT. https: https:///ithhii
2022.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rns:: Fast autoregressive transformers with linear attention. In International conference chine learning, pp. 5156-5165. PMLR, 2020.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In International Conference on Algorithmic Learning Theory, pp- 597-619. PMLR, 2023.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0445), 2020.
Anastasis Kratsios. Universal regular conditional distributions. arXiv preprint arXiv:2105.07743,
2021.
Wojeiech Kryšcifíski, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. Booksum: A collection of datasets for long-form narrative summarization. arXiv preprint 10X.022005.08209, 2021.
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer time series forecasting. Advances in neural information processing systems, 32, 2019.
Clara Meister, Stefan Lazov, Isabelle Augenstein, and Ryan Cotterell. sparse attention more interpretable? arXiv preprint arXiv:2 10.e2/06.00000 000.07 2021.
Stephen Mussmann, Daniel Levy, and Stefano Ermon. Fast amortized inference and learning in log-linear models with randomly perturbed nearest neighbor search. In Gal Elidan, Kristian Ker- sting, and Alexander Ihler (eds.), Proceedings of the Htirty-Third Conference on Uncertainty in Artificial Intelligence, UAI 2017, Sydney, Australia, August 11-15, 2017. AUAI Press, 2017. URL pd.f.
http: atwaaaaaaaaa.//aaaa/aaaaaaa...aa www.or//aaaaaeeeeaaaaaee 2011111111 http//pooceeocec httpe/eddddng/ppppppppepeeee
Behnam Neyshabur and Nathan Srebro. On symmetric and asymmetric Ishs for inner product search. In International Conference on Machine Learning, pp. 1926-1934, PMLR, 2015.
Matteo Pagliardini, Daniele Paliotta, Martin Jaggi, and François Fleuret. Fast attention over long sequences with dynamic sparse flash attention. Advances in Neural Information Processing Sys- fems, 36, 2024.
12

---

Published as a conference paper at ICLR 2025
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. arXiv preprint arXiv:1911.02972 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8)99 2019.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-basee sparse attention with routing transformers. Transactions of the Association for Computational Linguis- tics, 9:53-68, 2021.
Barna Saha and Christopher Ye. The i/o complexity of attention, how optimal is flash attention? arXiv preprint arXiv:2402.07443 2024.
Anshumali Shrivastava and Ping Li. Asymmetric Ish (alsh) for sublinear time maximum inner prod- uct search (mips). Advances in neural information processing systems, 27, 2014.
Ryan Singh and Christopher L Buckley. Attention: Marginal probability is all you need? arXiv preprint 2xXiv:2304.04556. 2023.
Luca Soldaini and Alessandro Moschitti. The cascade transformer: an application for efficient answer sentence selection. arXiv preprint arXiv:2005.02534 2020.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:2011.0400 2020.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Re- thinking self-attention for transformer models. In International conference on machine learning, pp. 10183-10192. PMLR, 2021.
Ashish Vaswani. Attention is all you need. arXiv preprint 1rXiv: 2017.
Pichao Wang, Xue Wang, Fan Wang, Ming Lin, Shuning Chang, Hao Li, and Rong Jin. Kvt: k- 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 arXiv:2006.04768, 2020.
Qingsong Wen, Tian Zhou, Chaoli Zhang, Weiqi Chen, Ziqing Ma, Junchi Yan, and Liang Sun. Transformers in time series: A survey. arXiv preprint arXiv:2202.07/25, 2022.
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. arXiv preprint ar X2v::22030081133 2022.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 14138-14148, 2021.
Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In International Conference on Machine Learning, pp. 40605- 40623, PMLR, 2023.
Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting? In Proceedings of the AAAI conference artificial intelligence, volume 37, pp- 11121-11128, 2023.
Haoyi Zhou, Shanghang Zhang. Jieqi Peng, Shuai Zhang, Jianxin Li. Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting, In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 11106-11115, 2021.
13

---

Published as a conference paper at ICLR 2025
APPENDIX
In the following sections we deposit theoretical results, proof' and algorithms that are missing from the main paper, either due to space constraints or for the sake of clarity.
A PRELIMINARIES
A.1 SELF-ATTENTION AND APPROXIMATION
We start by defining self-attention and some of its variants.
Definition A.1 (Self-Attention). Let Q, K, V Raxd We can think of these matrices as col- lection of n d-dimensional query, key and value vectors respectively, We define the self-attention function follows:
-(Q,KKV)--''4
(36)
where A exp(QKT) E RR the attention matrix and D diag(A1"). D effectively implements taking a row-wise softmax of the entries of QKT. Many modern implementations of attention consider causal attention, in which we mask away the upper-triangular entries of A.
Remark A.1 (Normalization by vd). In most implementations we divide QK" by vd (Vaswani, 2017) because it reduces the variance of each element of A had Q,K, T been selected from uniform distribution. We will omit this technicality because it does not affect our algorithmic tech- niques. For the rest this puper we will assume that 1-1/2 has been pre- pre-normalized into K.
K,
Remark A.2 (Dropout). Many attention implementations also use dropout. Every entry of will be masked to 0 with probability p. where is set to small constant, like 0.1.
Definition A.2 ((E, 8) esstimmaooss) Let X be a statistic and X be an estimator we have for it. X is an (e, 8) additive estimator if with probability at least 1 1 it holds that
X
SE
Respectively, we call the estimator multiplicative if
X - X EX
We will often make use of the following boosting lemma from the theory of randomized algorithms:
Amplification Technique, (Chakrabarti, 2020)). If is an unbiased estimator of some stutistic, then one can obtain an (e, m)mlliiplica estimate of that statistic by suitably combining
Lemma A.1 (Median-Of-Nel
Varl
K ==
In
E E
independent samples of Q, where C universal constant.
A.2 GUMBEL NOISE
The Gumbel Distribution will be useful for sampling from sofumax distributions. We define it below:
Definition A.3 (Gumbel Distribution). The Gumbel distribution with mean f and parameter B has the following probability function::
1
Gumbel(p,, ()(x)
(37)
3
We will make use of the following properties of the Gumbel Distribution:
14

---

Published as a conference paper at ICLR 2025
Lemma A.2 (Gumbel Distribution Properties, (Chattamvelli & Shanmugam, 2021)). The following
are me.
The mean of a Gumbel distribution is ps Bry²
The moment generating function (MGF) of the Gumbel 3) distribution is:
M(t) ['(1 Bt)e
(38)
where the Gamma function.
We can easily sample a Gumbel random variable of mean and parameter 3 by using the uniform distribution:
= (=---l((nn(U)) where U unif([o, 1])
(39)
Let M1,... M, be (p. B) independent Gumbel random variables. Then,
M max M;
=E[n]
is a (n 3In n, B) 3) Gumbel random variable.
Next, we present the well-known Gumbel Max Trick, an alternative way to sample from softmax distribution:
Lemma A.3 (Gumbel-Man Trick, (Huijben et al., 2022)). Let T1, be real numbers and con- sider the softmax categorical distribution p where
pi exp(x:)
Σ exp(x:;)
Consider sampling Gumbel random variables GI..., Gn Gumbel(O, 1) and let
i E arg max + G;}
Then is distributed according to
Binomial Distribution As part our notational conventions, we also let Bin(n, p) be random variable distributed according to the Binomial Distribution with parameters and p.
A.3 LOCALITY SENSITIVE HASHING
In our theoretical exposition we will make extensive use of schemes for approximate nearest neigh- bor search. A very successful such suite of algorithms with a long history (Andoni et al., 2014; 2015) of provable theoretical guarantees is Locality Sensitive Hashing (LSH)):
Theorem A.1 (Existence of LSH, (Gionis al., 1999; Mussmann al., 2017)). Let V C U be set of size with a similarity measure Sim(..)) Consider a hash family H such that for scalars S S2 and Pi > P2
For any x, E V where Sim(r, y) 2 S₁. Praeh|h(as =h())= 2 P1
For any x. y E V where Sim(x, y) S Sz. PFREи [h(z) = h(v)] S
This is called an (S1. S2, P1, P1)P2) Locclity Sensitive Hash Family. Given such a family, one can construct a data structure which, given any query E U, does the following with high probability: there exists some point E with Sim(u. q) 2 St, returns a point o' V with Sim(v', 2 S2. If no point E V exists with Sim(v, 4) 2 82, it returns negative answer 1. Further, this can be done with Ö(nP) query time and (pp)+p) space where p logP1/ log P2
LSH also finds numerous applications in solving the Maximum Inner Product Search Problem (MIPS) as shown in Neyshabur & Srebro (2015), Shrivastava & Li (2014), and others.
by 0.577 is the Fuler- Mascheroni constant.
15

---

Published as a conference paper at ICLR 2025
A.4 RANDOM WALKS AND SOME CONCENTRATION BOUNDS
For our back-propagation algorithms we will make use of some elementary tools from the theory of Random Walks.
Definition A.4 (Random Walks), Consider state spuce V [n] and a weighted complete graph on V with weights w in [0, such that for all u E E
Σ VEV
Wwo
This graph represents random walk with transition matrix P E [0, 1]axn, where Pij Waj- P stochastic matrix because rows sum to a random walk, we start at some vertex and choose neighbor to jump to according to the probability distribution in P. The choice at each vertex conditioned on the previous transitions only depends on the vertex itself. This is known as the Markov Property.
A first elementary observation is that if we start with distribution p over V and we do single step in the random walk, we can obtain the resulting distribution by multiplying the original distribution with pr:
Lemma A.4 (Single Step Random Walk Transition). Consider distribution A(n) over V. Then, the quantity pr gives the distribution over V after one step of the random walk.
Proof. Let q be the distribution after one step. Let E V. We have by law total probability that:
q(v) = Σ p(u)= Wuv = (PT p).
wEV
Finally, we state a well-known concentration result about independent random variables:
Lemma A.5 (Hoeffding Bound). Let X XI,..., Xn be independent random variables where S X; S by almost surely. Let S. = X1 + + Xn- We have that:
212
[IS - E[Sn]] 2 2 2exp
Σ (b. - a;)
B PROOF OF LEMMA 2.1
We prove Lemma 2.1. Our proof deviates from the proof of Mussmann et al. (2017) in that it uses MGF-based argument, which we believe cleaner.
Lemma B.1 (Reminder). In the context of Algorithm I, we have that:
[m] < n
(40)
k
Proof. By Lemma A.2, we can generate Gumbel(0, 1) random variables as follows: Let U, be uniform in [0, 1]. Then:
Gig = 11=-In(nInU))
(41)
is distributed according to Gumbel(0, 1). We want Giy > B, which implies that:
IIn[ In(U,)) > B + U, exp( exx(exx(B))
(42)
(A(n) { E T x 0./1 } the probability simplex over [n].
16

---

Published as a conference paper at ICLR 2025
So the number of points for which the Gumbel noise exceeds B is distributed according to the Binomial distribution with parameters n - and exp( exp(-B) If we condition on M := max ky Gigl. h. have that: jES,
E mm M] (n k)(1 exp(- exp(-B)))
(43)
- n
(44)
where the last inequality follows by e 11=
exp((- - B)) < 1 exp(-B)) <1-(11exp(-B)) = exp(-B)
Now we can bound E[n exp(- - B)] by using the MGF of the Gumbel distribution (see Lemma A.2).
Let M' == maxjes, Giy. Recall by Lemma A.2 that M' a Gumbel random variable with μM
log k and Bar = 1. Let (t) = [[e be its moment generating function. We know that:
(t) r(1 - t) elog =I(1 t) k²
(45)
This allows to write:
B[n (app--Bip -B)] E[exp(-BB) B}]
(46)
n1 E[exp(Smmin - MM]
(47)
n Elexp(Smin max Lij + Git))l
(48)
jES:
S n Elexp(Smmn min Zi - M' JES,
(49)
n E[exp(-
(50)
(-i(-1)
(51)
n
(52)
k
where inequality 49 follows because
max JES:
{ Zij Giy) 2 min Zij} jES,
+ max Gij Smin + M' JES,
For quick proof of this statement, let j == arg minjes, Zij and j = arg maxjes, Gij. Also let
j == arg maxjes, (Zg + Giz). Then we have:
max Zij + Gij} JES,
Lig Gig.
(53)
2 2í + Gūj
(54)
Zg Gūj
(55)
Smin + M'
(56)
Finally 52 follows because [(2) 1. Now, via law of total expectation we finally get:
n n
m
EM
[m
M] < EM
(57)
k
k
C KNN ATTENTION VIA CONCENTRIC LSH
In the main paper, abstracted away the specific ANN method used to obtain the top-k key vectors kj for every query vector qt In this section we cover a method for solving the k-MIPS problem in sub-linear time per query that has sound theoretical guarantees. In the context of Algorithm this method could substitute the ANN index H.
This approach in question was proposed by Mussmann al. (2017) and it uses concentric LSH construction to get an approximation to this problem. First, let us define the approximate version of the k-MIPS problem, as is proposed in Mussmann et al. (2017):
17

---

Published as a conference paper at ICLR 2025
Definition C.1 (Approximate k-MIPS). We say that set Ss is approximate top-k inner product solution if |Sa] = and there exists constant e such that:
max of ky min q < JES,
(58)
yeS,
we are able to generate an approximate solution S; instead of an exact one, we have to lower our threshold B M Swin in Algorithm 1, which in turns implies that Elmm] I vn e.. This remains sublinear in n because с is a constant.
The solution to the approximate version of the problem constructed using LSH. Specifically, we build sequence of (polyloggn,, d)) LSH data structures, each with concentric approximation radii, and hash all the key vectors k, R into them. For each query vector qi. we hash across all these data structures and identify the first pair of consecutive LSH structures, D, and Di+1. where D++1 contains more than /n points in the buckets corresponding to di while D, contains fewer than vn points. For further details, see Mussmann et al. (2017). The following theorem ultimately holds:
Theorem C.1 (Mussmann et al. (2017)). Let0 < < 1 be a constant. There exists an algorithm for solving the approximate version of k-MIPS on any single query with probability least 1 1/n2 by using an explicit concentric LSH construction. The algorithm takes O(dn l polylog(14 d)) pre-pre processing time/space, and O(yn n² polylog(n, d)) time/space per query.
D
D
Figure 5: An illustration of the concentric LSH construction of Mussmann et (2017) the D341 band we find at least / points and in the Dj band we find fewer than vn points.
Remark C.1 (The role of p). The choice of p < allows us to compute S; in sublinear time for each i E [n]. The value of p is determined by the radii gaps in the concentric construction. Our algorithm for computing S, is sublinear in n because p 1. Depending on the particular input dataset, we could have p < 1/2, which which case f(n, VT. = Ö(yn) in the context of Theorem
C.I A COMPLETE ALGORITHM BASED ON SOLVING l MIPS
We now have complete algorithm to estimate self-attention with provable guarantees that is based on the solving k-MIPS problem for every query vector q. If we combine the boosted estimator approach of Theorem 2.1 with the Lazy Gumbel Sampling Technique of Algorithm 1 and the k- MIPS LSH technique of Theorem C. we arrive at the following theorem. We give pseudocode for the resulting algorithm, Algorithm 5:
18

---

Published as a conference paper at ICLR 2025
Theorem C.2. Let 0 and >> 0 be small positive constants. There exists an algorithm that can estimate Self-Artention in the same way as Theorem 2.1 and fail with probability at most + 1/n. The algorithm's 's time complexity is shown in the table below, where p (0, 1) is a fixed constant.
Pre-Processing
Main Computation
Complexity
PP
O (n (mxx(221 d log(1/6))
Algorithm Approximating Self-Attention using concentric LSH k-MIPS solver
1: Inputs: Q,K, V Ra error parameter > 0, confidence parameter 8 >>
2: -rreate Concentric LSH data structures for solving k-MIPS, as in Mussmann al. (2017)
3: Let Ô RRax be our output.
4: for i E [n] do
5:
S; 4- Query H for the vn indices j E [n]] with the approximate largest values of q kj
6: for j E [d] do
7:
Oij Median-Of-Meant with Algorithm sampler + (vTL, qú S.).)
8:
return Ô
Proof. Let k = vn. Suppose we construct the concentric LSH data structure according to Theorem C.1. This takes (dn" 1+)) time. Let condition on the event that for all queries the data structure provides a correct approximate answer S to the k-MIPS problem. This happens with probability at least 1 - 1/n by union bound over all queries. Now we use our sets S; in Algorithm to sample from D;. Since retrieving St takes f(n, k) = O(vn + nP- polylog(n)) time and space, Theorem 2.3 dictates that sampling from D; also takes Ö (ax12222)) time and space. Thus, in the context Substituting back gives us the desired runtime
of Theorem 2.1 have that T and failure probability guarantees.
D PROOF OF LEMMA 3.1: DERIVATION OF THE ATTENNITNN GRADIENTS
In this section, we will show the proof of Lemma 3.1. Suppose we have differentiable scalar function & that represents the loss when training our neural network after computing the output O: 86 for all E [m].] E [d] and stored it in an matrix
f ((O). Suppose that we have calculated 80,, D Raxd Now we will calculate the remaining derivatives by using the chain rule. A similar calculation is also done in the Appendix of Dao et al. (2022).
As Calculating JV, All these calculations just use the chain rule. One can simply draw a tree of dependencies and use it to perform the derivation. & depends on Ozg and Oig depends on all Vyj,
aΦ
dø 00r;
Σ Do 00.j
a
oV,
a0,g
oV,
Now, we calculate that:
00
Σ
ProVis = Pri
0V
oV, k=1
so that gives:
0Φ
n
n
Σ D P. Σ D
(59)
aV,
Thus, we can write the result succinctly:
DV = . DO
(60)
19

---

Published as a conference paper at ICLR 2025
Calculating
as BQ: To do this, we will first calculate
de a
and 0e "S,"
where S QK?
First, each Og depends all Pak, the chain rule gives:
06
dΦ
d0tk
d 80
Σ
Σ
D
aP3
i0
OP
OPs
We can calculate that:
80
VjN
OPij
and so:
OΦ
Djj iP.j
Σ
Die Víx (Dig Vis)
(61)
k=1
for all [n].j [n].
Now recall that Pij = exp(S:))
P. depends on all Si for k = 1...., n. Thus:
0Φ
0b OPa
D OP.
aS
0Pk
OS:j
öSij
OP
0Pp D
= =D
Σ
as;j
aSig
k=L,
We now calculate seperately the two cases by using the quotient rule:
k + j:
0 exp(Sat) - exp(Sat
exp(Sit)
OS.,
a5
2
Σ exp(Sir)
Σ exxp(S..)
= -
=
exp(Sig) Σ exp(Sir) exp(Sis) exp(Sis)
0P;
exp(S.g)
0S.
aS,j Σ exp(S.r)
Σ exp((Str)
Pij - P2
Now we can put it all together:
Os
n Σ
D.j- OS.j
Djt OS,1
as,y
k=
= DD (Pij - P.))
Σ D Psp
k= IKki
= -Dg Pij - Σ D PacPi
k=1
Pij (D!) (D Pa:))
Now finally, for i E [n], e [d]. Qú, influences Sip. for all A E [n], so:
0Φ =
do asa
(62)
0Q
USw aQ
k=1
=
Σ
Pae ( sk (D..,, Pt.)) Kaj
(63)
k=1
20

---

Published as a conference paper at ICLR 2025
de Calculating aKg
We know that Kij influences Ski for & [n], so:
dΦ
do
OSki
(64)
ikij
aSa akij
km
Pai (D - - Qud
(65)
E ESTIMATING D
In this section we give an efficient algorithm for estimating DR_ This algorithm is based on our i-N-Attention framework. Recall that we found that:
DO
Piee (Dl - (D Pa.)) Kkj
can write this expression as an expectation with respect to the distribution D;:
iQ ENND.
[ EA
[KRj Eg [
(66)
= E.DD.
D Kajl
ERnD,
[Kaz] E. [
(67)
Ea
Ea
This allows us to use any of our softmax expectation estimators. We choose the estimator for the purposes of clean analysis. We just have to do it three times and ensure that the terms we take expectations over are efficiently computable. Indeed, because
D (D) VE.),
(68)
compute all three of those expectations in sublinear time! Let E1, E2, E3 be the estimates we produce. Then, almost identically to the error analysis we did for the forward pass, we get an (e, -addiiive estimate for E,, where i (1,
Pr
Ei
Eil
VA
(69)
On
Pr
F22
2
V
(70)
06
Pr
Eg Egl
m
IA
(71)
And so, putting these three together and using the union bound we get that with probability at least
- ô it holds that:
F
Fog
F.g
F1
F.
E
<
F E
Pg Fg
F2
F3
(72)
E
E
Eg E' Es
E3
E2
Li
E3
Es
(73)
E E
E2
Es
Es
Es
E2 E2
(74)
S Eg +
(75)
S e F2 + +6) E
(76)
= = - <(E2 + Es)
(77)
In order to bound the variance of our estimators, we need to assume some bounds on the inputs,
analogously to || < B = O(log n) from Theorem 2.1. First, we assume that ||]]
& Bx
Inn Theorem 2.1 we multiplicative approximation. To get additive approximation guarantee we need ( 2 1og(1/8) Var(ō.,]) samples, where Var[O,,] B² Oppoyyogg(n)).
21

---

Published as a conference paper at ICLR 2025
O(polylog(z)). Second, have that || 1DD"||os < dB dB.1IDDl|> O(polylog(n)) if d Ollg n). This also gives that ||D 2 $ Bx || = O(polylog(n)))
These assumptions reasonable within the context of the hardness results proved for the atten- tion mechanism and the computation of its gradients" (Alman Song, 2024a;b). Given these assumptions,
also bound the emor more compactly. Starting from Equation 77, get: eq & (e)) + E(Bx + dB ||D lloo) = O(E polylog(n)). As a result, we arrive at the following theorem:
Theorem E.1. Assume that ||K|| Oppolylognn)). There exists sub-quadratic algorithm that takes as imput Q, K, D E Raxd and outputs matrix DQ E guxd such that:
Op d Ollogn) and 1 8c
1
S O (olyloooy)g)
(78)
This algorithm is shown as Algorithm 6.
Proof. The proof and analysis of Algorithm 6 detailed the preceding paragraph.
Algorithm Estimating D
procedure ESTIMATE- E,Q D,,,,, D.
F (DD f. Ve.) KxilE-1 E Raxi
b F will not be materialized.
Ei Median-Of-Means with Lazy Gumbel Sampling Q, F,S,
return E
procedure ESTIMATR- (2, K,S,, 6,
Eg t Median-Of-Means with Lazy Gumbel Sampling + Q, K, K.j, S,,5,
return Ez.
procedure ESTIMATE- K,V (,,,i,))
F F (DD. Vo, E Raxl
b F will not be materialized.
E3 Median-Of-Means with Lazy Gumbel Sampling + Q, K, F, St,,,
return Eg
Input: D ERxd Q K, V. Rexd parameters E,
Let BR E Raxd be our output.
for i [n] do
S + /n values E [n]] of the largest of Kn via LSH or KNN.
for j E [d] do
Ej + ASTIMATE-Ei(Q K, K,V, D (S,i,,,,,,))
E2 + ESTIMATE- E2Q, i,,,,,,,,
E3 ESTIMATE-E2 V, S,,,,,
D0 E, E
return D
F
ESTIMATING D
Finally, we turn to estimating D Our earlier calculations show that
06
Σ
Pai (Df
-DE Pk.)) Qxi
k=1
Another motivation for assumming upper bound on the norm of DO is to avoid the phenomenon of exploding gradients in training neural networks.
22

---

Published as a conference paper at ICLR 2025
We can break up this sum into two terms:
OΦ
Σ PRI
Pks (D.:: Pac)) Qui
(79)
Bij
We will estimate both terms separately:
F.1 ESTIMATING Arj
For i E [ni] and E [d]. we have:
Av Σ PaD
Σ
Paa (D Va.))
(80)
P Y (i)
(81)
where ylf kj Qkj (Dl.) Vi.). So we can write:
Ai = (PT)
(82)
nxl
We will use our familiar Markov Chain estimation method from Algorithm to calculate this quan- tity. However, in this case we only care about estimating the i-th entry in the vector (PT) (() which do by performing Ollog n --2) simulations. Ultimately, by following the same analysis as in Algorithm 3, we are able to estimate Ai with probability at least 1 1 and error:
Ag - Ang S SE(Y : 1n) + 2
(83)
=
Qui (De. Vin) + 2
(84)
where
M = min (0)
kE(M kj
yt
Remark F.1. Because we would need to calculate all n² values of M we will instead use a single
upper bound M 2 M, for all (i.j) [m] x [d] for this algorithm. We assume that we know a large enough M in advance and that M (polylognn)).
In the next paragraphs, will tackle some implementation issues that arise in this approach. We did not see these issues when estimating DV, and because they make the algorithm a lot more complicated, left them for last.
Pre-calculating the normalizing factors We need to pre-calculate the normalizing sums ND (x.)) 1") + nM for all (i,j) E [m] x [d]. Naively, takes O(n2d) time to calculate all those sums. However, with some preprocessing we can take the time down to (ndd). First, observe that we have:
NOO nM (Y. 1n) -nM + Σ Qki (DR.. Ve.) nM + (V Eau
Qkj D2.)
(85)
k=1
We can thus first pre-compute the d vectors time. Then, for each E [n]] and j E [d], we can produce 17(0) in O(d) time by using Equation 85, bringing the total time complexity O(nd)).
= Qkj DO. E ddd for each j E [d] in (ndd)
23

---

Published as a conference paper at ICLR 2025
Sampling according to y: + M 1 efficiently Unfortunately, because we are now estimating Ay individually for all (i.j) [n] [d]. we cannot spend ((n) time generate each sample. We need to generate samples in sublinear time with some pre-processing. This seems intuitively difficult first because we have ((nd) distributions over [m] and each distribution requires (2(n) time to sample one sample. However, we can take advantage of the structure between the distributions in order to reduce the pre-processing time. First, consider the following method of sampling from a distribution (p), ........P
1. Compute the cumulative sums s; Ei= Ps- We know that $1 = P1 and Sn 1.
2. Pick some x ~ Unif(O, 1) uniformly at random from (0, 1).
3. Find the interval [ps. Pe41] for i E [1. 1] in which falls in. That is, find the smallest i for which x S 8z. We can do this in (log n) time using binary search.
4. Output
CDF Sampling Illustration
CDF
Fx
3.6
0.4
0.2
Chosen Index
0.0
20
40
80
100
[c]
Figure 6: An illustration of the CDF sampling method: We form the CDF and then sample an index by choosing E (0, 1) and using binary search to find the corresponding bucket.
It is easy see that this method outputs value i with probability Pt. If we applied this method naively we would still take O(n2d) time because we'd have to calculate all the cumulative sums. However, the inner product structure again comes our rescue:
Y
(V... Qkj D2.)
(86)
So, we can create d cumulative sum tables Σ, for j E [d]. each of which stores n cumulative-sum dd vectors as follows:
(E))e= = Σ Q.j Do. E ER, Ve E [m]
(87)
This requires (ndd) time and space to construct. Now, in order sample with probability propor- tional to (()) ki + M given that we know (7(0) we sample Fig ~ Unif(0, 1) and perform binary search to find the interval Zij belongs to. At that point, we can calculate the O(log n) necessary cumulative sums in ((d) time each by using our pre-processing:
M)
RM + (Vi.. (Ex)e)
(88)
kj
This allows us to sample in Od n) time after O(nd?) pre-proces,nng. Our algorithm in total is included as part of Algorithm 7.
24

---

Published as a conference paper at ICLR 2025
Sampling with respect to D; Again, cannot afford to sample from the softmax naively with O(n) time. Thankfully, we know of a sublinear method that can allow us to sample from the softmax, with slightly super-linear pre-pnn processin time: the Lazy-Gumin Sampling method. We will omit the pre-processing details the algorithm pseudocode.
Algorithm 7 Estimating D Part I: Computing A
1: Input: K,K, D0 RRax error parameter
2: for j E [d] do
> Pre-Processing
3: Compute Ej =
- Qkj Dh. Rd
4: Compute the cumulative sums
Quj-
E2 for all l [n].
5: Compute @ p" 1at using Markov Chain simulations.
6:
Initialize a KNN index H.
7:
procedure COMPUTE-ACC K, K, DO E,E,,,,,,,,HH,. Let N t 2 2og -2
8:
9:
A - [0]exd is the output.
10:
for E [n]] do
11:
Query H to get set S;
12:
for j E [d] do
13:
N10 (V Eÿ) nM
O(d) time.
14:
for [N] do
Sample E [n] with probability x Y + M via binary search, Σ, and N7(0
15:
16:
Sample f E [n]] with probability Pae via Lazy Gumbel Sampling, given S
if then
18:
Aij - A.j +1
A.j + (Aig .2.)) M 3
20:
return A
F.2
ESTIMATING Bij
For (i,j) E [n]] x [d]. we first have:
Bij
Pai (DF., Pk.)) Qui
(89)
k=1
X
(90)
where XAj = (DP.. Pa.)) Quj. Notice that Xkj takes ((nd) time to naively compute, so we will first approximate it with Nzy. Observe that:
Xaj Qkj (DP. Pk.))
(91)
Qui Σ Dbs
Pko
(92)
Qui EonDy (Da
(93)
EanDe 1Qkj
Dio
(94)
XEj
(95)
Let us approximate BE-D, Quj (DP) s RNg using the Lazy Gumbel Sampling and Median-
method. This allows to get for all () E [n] x [d] with probability at least
that:
XAJ
XNs Se
(96)
25

---

Published as a conference paper at ICLR 2025
Remark F.2. To have an o(n) bound for the variance, have assume (again) that || xc (ppolylog(n)). This follows from the assumption that ||Q||e = O(polylog(n)) and || = (n)) The latter follows from | O(polylog(n)). So the assumptions here the same as in Theorem E.I.
Now we can define:
B
PhiXAs
(97)
k=1
We can bound the error of this upproximation using the triangle inequality:
Big Big
Σ
Pro - Xus)
(98)
k=1
Pki XRs Xns
(99)
k=1
n
Phi
(100)
k=1
E (P... 17)
(101)
Now the problem is calculating B. Note that we can write:
P= X
(102)
Finally, this takes back to the calculation of DV, We can use the exact same Markov Chain method and get final approximation B that with probability at least 1 - it holds that:
Big Bis S Se(R 1") + 2
(103)
where
M = min RAz
ke(n)
Then the overall error be bounded follows:
B.j Bij
< B.j Bij + B:j Bij
(104)
Se(P.,i )n) + 8(8 1 - 2en.
(105)
$ <P(,,, 12) + e(X.,,,17 1n) =n + 2en
(106)
= =e(P., + X.j, 1") + e2n - 2
(107)
To wrap up implementation details, we can calculate the required normalization sums as follows:
(8 17) =
Qky Dr
(108)
1=1
k=1
We can do this in 3dd33/2) time if we precompute in advance
D. = -Df Pk.))
(109)
using Lazy Gumbel Sampling for all & E [n). Further, each element X, can be computed in Ö(vn) time well in similar fashion. Finally, x can also be caleulated in such time. Our algorithm is given below as Algorithm By combining algorithms and 8 we arrive at the following theorem for Algorithm 9:
26

---

Published as a conference paper at ICLR 2025
Algorithm Estimating D
Part Computing
1: S Use an LSH or ANN index to calculate S; for all E [n].
2: <- ESTTMATERODDCTTOSTIVVVVE
3: procedure COMPUTR-Xx;y K, DO e,
4: F {Quj (D&. Vo,))> 1- Raxi
F will not be materialized.
5: KNj t Median-Of-Meant with Lazy Gumbel Sampling +- Q, K, F,S,,eã
6: return Xxj-
7: procedure COMPUTE-B(Q K, KV,,D.,,
8: Output B - [0]m
9: for j E [d] do
D (()) times
10:
B.j f E-ESSTIMATPPRDDUCTPPP. e,
8,
11: return B
Theorem F.1. There exists algorithm that approximates DK inputs Q, V, DO under our standard assumptions such that the estimate D satisfies:
||D.5 D.K
S + X. 1 + 22n + 2
+ E
Qkj ((D"),.. Vis) + 2enM
k=1
where:
= := min y( kj and M(A) = min X
M 2
(110)
kE(n]
kEinj N
Yi <0
under our previous definitions for all j E [d]. The algorithm runs in sub-quadratic time and space and succeeds with probability > 1 8.
Proof. We have that D. Aij - Big. We define D. Aus - -Bis and we established in the f)
preceding discussion that:
A Aul 8 n
Qki ((D)), Vo.) + 22nM
(111)
BBg Bil e(P., + , 1n) 22 2
(112)
Thus by the triangle inequality we get the desired error guarantee.
Algorithm 9 Estimating 8 Putting it all together
B < COMPUTE-B(Q. K, V, s) Á COMPUTE-ACC K, V, D E, return Ā A-7, B.
G VECTORIZED IMPLEMENTATION OF THE FORWARD PASS
We present the vectorized implementation of &NN Attention that we used in our experiments. This based on Theorem 2.5.
import
torch
This function calculates the attention ion mechani sm in the forward rd pass. Inputs:
27

---

Published as a conference paper at ICLR 2025
Q: tensor of shape (o, n, containing the quezy vectors.
K: tensor of shape (o, n, containing the key vectors.
V: tensor of shape (b, h, n, d) containing the value vectors.
- k: The number of vectors to attend to via KNN metheds.
10 Note that b the batch size, is the number of heads,
11 n is the sequence length, and d is the dimension of the vectors.
12
13 Outputs:
tensor of shape h, d) containing the output vectors. 15 def htt_fffowwrrbbtthhoe (2, V, k): N, s.shape
16
17
output torch. zeros H, dtypestorch. Float32)
for b in range(B))
20
for h range : A Get the top k indices of Qlb, h] 8 x(b, h] "I and the scores. scores, s topk (Qlb, h, i,; K k[b,h, i,i,, k, masking=Truc) scores scores (1 math. (D))
21
22
25
# Calculate the denominator.
26
M wwwch..aaxsccorreeertaaaaa dim-1) [0]
27
exp_scores ttorcheeppp (scores M. u.usqueezze (1))
denom torch. sum sexp_scores] dim=1) ansqueeze
29
30
# Calculate the numerater.
31
A Vbh[S] shape - (N, D,
numerator torch. bmm (exp_scores. u (1), v[b,h, numerator. squeeze
34
35
output [b, h] = numerator / denom
36
return output
Listing 1: kNN Attention Vectorized
H SAMPLES FROM FINE- TUNED GPT2-XL
Below we give some samples from our fine-tuned GPT2-XL. using fast approximate &NN attention:
My favorite thing in the world is: Just Cause And couldn't it stop playing it. It's mess. I'm not sure how to describe it. But it's just so fun. When you play it, you know instantly that this is a very different kind of game, one that takes place in an open world of sorts, and you have bunch of different characters. It's epic, and it's open world, and you can play as bounty hunter, or you can play as a doctor, or a cowboy, or you can play as soldier. You have a buddy system, and you can go another character's 's story mission if you want. This is a great game, and I' excited to see what the developers do with it.
Prompt:
Prompt: like being in nature because it's so natural, and because it' always changing. Nature is friend and a companion, and I I so thankful every time I'm able go into the outdoors. The second thing that love about being in nature is the comfort and energy that you can feel! I feel so happy and healthy on the trail when I'm interacting with nature. don't 't have to think about what I m wearing, or what I I'm doing, and can just be in nature in my element. also feel great being out nature because I I'm not spending all day at my desk, working in a office. I get to meet new people along the trail and can spend time talking to them.
28


Reviews:
[{"review_id": "tUy9uYKua1", "reviewer": "Reviewer_SjhM", "summary": "This paper theoretically analyzes kNN attention, which improves on the quadratic complexity (with respect to context length, $n$) of traditional full attention. The authors do this by viewing the self-attention matrix as an expectation over multiple softmax distributions, and then use the Gumbel-Max trick, along with the concentration properties of the Gumbel distribution, to efficiently approximate the original self-attention. This lazy Gumbel sampling, combined with k-MIPS, results in a total runtime of $O(n^{1.5})$. Additionally, the work approximates the attention gradients using length-1 random walk sequences, reducing the naive complexity from $O(n^2d)$ to $O(nd^2)$, providing high-probability approximation bounds.", "questions": "Please see weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper focuses on the important and relevant topic of the quadratic complexity of attention, given the ever-increasing context length in LLMs, and provides a rigorous theoretical analysis behind the empirical performance of kNN attention.\nThe paper is generally well-written, easy to read, and the ideas are clearly organized and discussed throughout. I enjoyed reading it, and liked how the authors first decompose the attention matrix as expectations over softmax, then use median-of-means boosting to achieve high-probability approximation bounds and runtime complexity bounds. The use of the Gumbel-Max trick and the concentration properties of the Gumbel distribution is also cute, ultimately leading to the gain over quadratic attention.\n\nThe use of 1-step random walks over the transition matrix P to approximate matrix-vector products (giving attention gradients in this case), although known, is also pretty nice. Overall, I appreciate how different well-known ideas are effectively combined.", "weaknesses": "I don’t have many weaknesses to point out, but I do have some questions/points I’d like to raise/discuss:\n\n\n—The experiments in Figure 3b) show that for $k>n^{⅛}$, kNN attention performs quite well, and there doesn’t seem to be much need for $n^{1/2}$ (as suggested in your theory). I understand you’ve mentioned this as potential future work in the conclusion, but why do you think this is the case? As far as I understand, the choice of $k=\\sqrt{n}$ in your theory arises because you want to balance the samples outside $S_i$​, which could, in expectation, ruin the top score after the addition of Gumbel noise, and the accuracy of lazy Gumbel sampling. This factor of $n/k$ also appears in the discussion in (Routing Transformers, Roy et al.), where the complexity is $O(nkd+n^2d/k)$, and $\\sqrt{n}$ is the optimal choice. What do you think explains this significant gap?\n\n\n—(Related to the above) For kNN attention without median-of-means (Sec 2.3), you randomly sample outside the top $k$ similarity products and upweigh them to capture the tail of interactions, and this is the algorithm used in the experiments. Median-of-means doesn’t consider the tail at all. Do you think capturing the tail behavior is critical to $k≪\\sqrt{n}$​ performing well?\n\n—Regarding the experiments in Section 4.2: The true benefit in the backward pass should only show up with large $n$. I understand that training larger models is difficult, but it would be interesting to see what happens with a moderate $n∼1000$ when training with cross-entropy.\n\n—What is $n$ for the final set of experiments (perplexity and approximation error)? Also, for the mean error of the kNN ($k$ vs. $n$) experiment, what is the range of $n$? I couldn’t find these details in the appendix.\n\n\n—Minor point: There is no figure number for the figure at the top of page 9, which I believe should be Figure 3. Please fix this and the rest.", "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": "2swTobLMrv", "reviewer": "Reviewer_WAZ2", "summary": "The submission presents a theoretical framework for *k*-nearest neighbors (kNN) attention, an alternative to standard self-attention where each token attends to only the *k* \"closest\" tokens instead of all tokens. Additionally, the authors propose approximation methods for self-attention gradients using kNN search.\n\nMore precisely, the submission is organized as follows:\n\n- In part 1, the theoretical results are outlined in brief.\n- Part 2 focuses on kNN attention, introduced as an approximation algorithm. In Section 2.1, self-attention is reformulated as an expectation, and Theorem 4 presents the primary approximation results, comparing the outputs of true self-attention and kNN attention, both of dimension $n \\times d$.\n  - Section 2.2 discusses efficient sampling from the softmax distribution (with each empirical distribution corresponding to a row of the attention matrix) via Gumbel sampling.\n  - Section 2.3 introduces an alternative method for computing kNN attention outputs, designed to be more compatible with modern hardware.\n- Part 3 describes randomized algorithms to approximate self-attention gradients (with derivations in the appendix). These estimations leverage random walk simulations, and Theorem 10 provides a theoretical runtime analysis.\n- Part 4 shows experimental results on kNN attention’s efficiency in terms of space and time costs (Figure (a)) and shows approximation error as a function of k (Figure (b)). Figure 3 compares learning curves for standard gradients obtained using backpropagation with those of the proposed approximation in part 3. Figure 4 evaluates perplexity and approximation error on real-world data using nanoGPT.", "questions": "Please see the questions in the previous part. \n\nIn addition: \n\n- Regarding the conditions in Theorem 8: are these practically achievable? For instance, can $n$, $d$, $T$, and $k$ be expected to yield satisfying approximation errors?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The problem studied in the paper is well defined in the abstract. \n- Understanding the approximation abilities of sparse attention models, of which kNN attention are a special case, is a significant research question. \n- Similarly, proposing alternatives to computationally costly gradient computations in Transformers could have a strong impact in the community.\n- The code is provided which enables to reproduce the experimental results.", "weaknesses": "I do not believe this paper to be ready to be published at ICLR for the following reasons:\n\nOverall clarity:\n\nThe paper is generally quite difficult to follow. Importantly, it is difficult to understand what is new in the paper. A clear statement of the contributions would be helpful. \n\nI am adding specific suggestions / questions for improvement below:\n\n- It should be clearly mentioned that Part 1 sketches the results that are later detailed in Parts 2 and 3. \n- Whenever a theorem is stated, a *detailed* proof should be attached to it, even if it is in the appendix. I could not find a proof for each theorem (for instance, theorem 10). Please let me know if I am wrong. \n- Lemma 1 is hard to understand as is. What is a multiplicative estimate ? If this is detailed in App A, then the remark in l. 77 should be added before the lemma. \n- Line 60 is unclear. For a reader encountering this sentence for the first time, the phrase \"how to extend the method to approximate the backward pass\" is confusing. Please clarify why this is important and what connection it has to kNN attention.\n- Mathematically, the paper lacks rigor. For instance, in line 104, is this an assumption on the differentiability of $\\phi$, or is it just notation? Also, where are the precise assumptions on the norms of $Q$, $K$, and $V$ stated? These should appear immediately following the theorem.\n- The source of the probability $1 - \\delta$ in all the theorems (e.g., line 92) is unclear. Please specify the source of randomness; is it due to sampling over the softmax-induced distribution?\n- In line 129, what does $T$ represent?\n- In line 120, the term $k_k$ in equation (3) is confusing.\n- Similarly, equation (5) is unclear because it takes the expectation with respect to $k$, yet there is an index $k$ in the sum. This could be clarified.\n\nTheoretical Contributions:\n\n- In general, I find the results challenging to interpret. How do these results compare to previous work? What are typical values?\n- Presenting attention as an expectation of the value matrix is not new. For instance, see *Attention: Marginal Probability is All You Need* (https://arxiv.org/abs/2304.04556). As written, it seems the submission presents this as a novel contribution, which is misleading.\n- Line 181: this is a proof sketch, not a full proof. I could not find the complete proof in the appendix.\n- The proof of Theorem 4 is also only a sketch. For instance, see the last two lines.\n\nExperiments:\n\n- There should be experiments specifically validating the theoretical bounds. As is, the experiments are rather qualitative.\n- I had to make some manual adjustments to get the provided code running. Including a `setup.py` would be helpful.\n- I ran your code for the gradient approximation in Figure 3. On my laptop, the approximated gradient takes approximately 100 times longer to compute than standard gradients. Did you observe similar behavior? In which settings do you expect it to offer speed advantages?\n- Similarly, I profiled the code from Appendix G against standard attention on a single GPU, using $B$, $H$, $N$, $D = 256$, $8$, $500$, $32$, and $k = 5$. I observed a runtime of 0.5831 seconds for kNN Attention versus 0.0061 seconds for Traditional Attention. If you do not observe similar results, could you provide code showing that your method leads to a speed advantage?", "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": "di8aTdTUl3", "reviewer": "Reviewer_HavR", "summary": "This paper studies the theoretical aspects of k-Nearest-Neighbor (kNN) attention (Roy et al., 2021) and addresses the challenge of quadratic complexity in self-attention. \n\nIt reformulates self-attention as expectations over softmax distributions, leveraging Lazy Gumbel Sampling (Mussmann et al., 2017) for efficient approximation. Then, novel sub-quadratic algorithms are introduced for approximating self-attention gradients using Markov Chain methods.", "questions": "1) Could you please clarify the theoretical novelty? For example:\n- Is Theorem 4 a direct consequence of Median-Of-Means Boosting?\n- Does Theorem 8 (main result) rely heavily on Mussmann et al. (2017)?\n\n2) Where is the proof of Theorem 2? Is it a direct consequence of another specific theorem?  \n3) In Theorem 4, what does $O(T)$ refer to? A discussion on the different parameters in the theorem is needed.  \n4) In Theorem 28, how large can $\\rho$ be? Does it approach 1?  \n5) Where is the proof of Theorem 29? Is it a direct consequence of (Alman & Song, 2024a; b)?  \n\n6) For notation, please clarify Gumbel $(a, b)$ and Bin $(a, b)$.  \n7) In my opinion, the statement and proof of Theorem 5 are vague. Are you trying to say there exists an index \\(\\hat{j}\\) that gives the maximum, and can you provide a proof for this?  \n8) In Lines 184–189, what is the Moment Generating Function of the Gumbel distribution used in the proof of Lemma 6?  \n9) Where is the proof of Theorem 7? What do you mean by \"it is easy to derive from Algorithm 1\"?  \n10) In Theorem 7, how strong is the assumption $k = \\sqrt{n}$?  \n\n11) Algorithm 2 takes $\\epsilon$ and $\\delta$ as parameters, but I couldn't find them in the algorithm's steps. How are they used? Also, what is the range of these parameters?  \n\n12) Theorem 8 is unclear. Why do we need $k$ and $\\ell$ to satisfy two inequalities and then set them equal from another inequality? Why not use only the equality case? Considering $\\epsilon$ and $\\delta$ between 0 and 1, for larger $\\ell$, we have:\n$\nk \\geq \\sqrt{\\frac{8n^2 \\varepsilon^{-2} \\log(4/\\delta)}{\\ell}}.\n$\nHowever, this assumption seems unrealistic as $k$ scales with $n$.\n\n13) Some parts of Algorithm 3 need clarification. In Line 2, \"lg\" should be \"log\". Why should $N$ be selected in that way? This requires more explanation. Also, what is $1^n$?  \n\n14) In Line 246, which algorithm are you referring to?  \n15) It is unclear how the total runtime in Line 224 was obtained. Is this result similar to Theorem 4 under comparable assumptions?  \n\n16) What is Figure 3(a)? I couldn’t find it.  \n17) What is the error definition in the discussion on \"Efficiency of kNN Attention\" (Lines 421–425)?  \n18) What do you mean by convex and non-convex cases? In both settings, the attention approximation problem is non-convex due to softmax and the product of $Q$, $K$, and $V$.  \n\n19) In the experiments section, you refer to “kNN” and “our algorithms” when comparing them with exact gradients or attention mechanisms. It is still unclear what \"kNN\" and \"our algorithms\" refer to in the context of this paper.  \n\n20) What specific gains can be observed from the experiments in Sections 4.2 and 4.3? Should we expect time improvements in these sections, for example, in the NanoGPT case? How can we precisely measure the efficiency of the proposed methods, especially for self-attention gradient computation?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "- The paper presents a theoretical analysis of kNN attention (Roy et al., 2021), connecting it to Gumbel noise sampling and Markov Chain methods for efficient gradient approximation. \n\n- I think the paper is well-structured. \n\n- Additionally, the author provides empirical experiments that demonstrate scalability and effectiveness in reducing computational and memory overhead.", "weaknesses": "- The authors need to clarify their contributions in the introduction and compare the theoretical novelty to Mussmann et al. (2017). The differences between this work and prior works are not clear until one reads the algorithms and theorems in detail. \n\n- Many theorems are presented without proofs or discussions (please see my questions below). This is problematic as even if the proofs are clear to the authors, they should provide proper references or detailed discussions on the theorems. \n\n- The literature contains many randomized attention methods, such as Nyströmformer and Skyformer. These should be discussed in the related works. Adding benchmarks from these methods would also be useful.\n\n- In the experiments section, some parts need clarification (please see my questions). For example, the authors used \"kNN\" as a legend, but it is unclear whether this refers to standard kNN or their modified version. Similarly, in several places, they mention \"our algorithm\" without specifying which algorithm they are referring to.", "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": "0o10z4DF5h", "reviewer": "Reviewer_AXc3", "summary": "This paper provides a theoretical analysis of $k$-NN attention, which is often explored as a method to reduce the quadratic time complexity of self-attention in Transformers. \nSpecifically, it demonstrates that by combining $k$-NN with lazy Gumbel sampling, an unbiased estimate of self-attention can be obtained within an $\\epsilon$-multiplicative error with sub-quadratic time and space complexity. \nAdditionally, the paper proposes a method for approximating the gradient of self-attention within $\\epsilon$-additive error using random walk simulation, also achieving sub-quadratic time complexity.", "questions": "- In Sections 2.2 and 2.3, two approximation methods are proposed. Can their performance be compared directly?\n- In Section 4.1, the experiments are conducted with the matrices $Q,K$, and $V$ sampled from a uniform distribution. How would the performance change if $Q,K$, and $V$ were sampled from more biased distributions, such as those encountered in real-world data?\n- In Figure 3(a), why does the computational time for $k=n^{1/4}$  outperform $k=\\sqrt{n}$? Also, what does \"Brute Force\" refer to in this figure?\n- In Section 4.2, would it be possible to include a comparison of computation speeds?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper excellently ties the use of $k$-NN in self-attention to the context of lazy Gumbel sampling by framing self-attention as an expectation calculation.\n- Clear pseudocode is provided for each algorithm.\n- The authors have made their experimental code publicly available.", "weaknesses": "- **Lack of originality**\n  + At the beginning of Section 2.1, the authors claim that the first contribution of this paper is interpreting the softmax computation in self-attention as an expectation calculation. However, this interpretation itself is not a novel idea. For example, see [1].\n  + Much of Section 2 simply applies results from [2] to the computation of self-attention. The authors should cite [2] in Theorems 5, 7, and 8, as well as in Lemma 6.\n- **Presentation issues**\n  + The reference to [2] is listed as an arXiv preprint, but it was accepted at Uncertainty in Artificial Intelligence 2017.\n  + In equation (3), the dot is used inconsistently between $q_i^\\top \\cdot k_k$ and $q_i^\\top k_s$.\n  + The time complexity in Theorem 2 should be clarified as time complexity in expectation.\n  + The citation format on line 97 of page 2 should be consistent with others.\n  + \"BIN\" in Algorithm 1 is not defined.\n  + The term \"kNN index\" lacks sufficient explanation.\n  + The sentence \"if we assume that...\" on line 222 of page 5 was unclear to me; further clarification would be appreciated.\n  + It would be helpful to clearly indicate which parts of Algorithm 2 constitute pre-processing.\n  + Theorem 8 seems to rely on Theorem 3.5 from [2], which requires the assumption that $V_{sj}$ is bounded. If, like Theorem 2, the assumption $\\\\|V\\\\|_{\\infty} = O(\\log n)$ is used, this should be explicitly stated in the theorem.\n  + The reference to Figure 3(a) and (b) is ambiguous; it seems to refer to the figure at the top of page 9. If so, the figure at the bottom of page 9 should be renumbered as Figure 4.\n  + Numbered and unnumbered equations are mixed inconsistently.\n\n[1] Kratsios, Universal Regular Conditional Distributions. 2021.  \n[2] Mussmann et al., Fast Amortized Inference and Learning in Log-linear Models with Randomly Perturbed Nearest Neighbor Search. 2017.", "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": "6LA8nYdTVW", "reviewer": "Reviewer_jxN9", "summary": "This paper provides a theoretical analysis of the KNN-nearest-neighbor sparse attention techniques for efficient attention computations. In that setting, every token attends only to its k nearest neighbors. By leveraging developed theoretical framework, the Authors propose a novel algorithm for the sub-quadratic time approximation of the self-attention gradients for efficient Transformer-training (default computations involving attention modules in Transformers require quadratic time in the sequence length, a prohibitively large time complexity for longer input sequences). The conducted analysis leverages an interpretation of the attention mechanism as computing the average value vector with respect to the softmax distribution defined on all the keys. Experimental evaluations confirms Authors' theoretical findings.", "questions": "1. The Authors cast it in the paper as an open question, yet I wonder whether they have any intuition how the optimal values of k can be derived in the more mathematically principal way (rather than just empirically), by leveraging the theoretical framework that was already developed in the paper. \n\n2. Did the Authors try to apply proposed in the paper sub-quadratic attention-gradient algorithm for other modalities than text, e.g. in the context of the ViTs ?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Great paper, that sheds new light on sparse attention mechanisms leveraging the notion of the kNN-graph. The probabilistic interpretation of the attention mechanism is actually well-known in the literature, yet it is very elegantly applied here to conduct a rigorous theoretical analysis of the method. New algorithm for the sub-quadratic computations of the attention gradients is yet another contribution of very practical impact. The idea to approximate the expectation coming from the probabilistic interpretation of the attention module via lazy Gumbel Noisy sampling is yet another beautiful insight that the paper provides.", "weaknesses": "The paper might further benefit from placing newly developed sparse attention mechanisms in the context of other efficient attention methods, e.g. those based on low-rank linear attention. It is also not clear how to choose the optimal value k, since, as the Authors explain, practically optimal value k is often significantly smaller than \\sqrt{n}.", "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)