$k$NN Attention Demystified: A Theoretical Exploration for Scalable Transformers

## Abstract

Abstract Despite their power, Transformers (Vaswani, 2017 ) face challenges with long sequences due to the quadratic complexity of self-attention. To address this limitation, methods like k 𝑘 k italic_k -Nearest-Neighbor ( k 𝑘 k italic_k NN) attention have been introduced (Roy et al., 2021 ) , enabling each token to attend to only its k 𝑘 k italic_k closest tokens. While k 𝑘 k italic_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 𝑘 k italic_k NN attention, reformulating self-attention as expectations over softmax distributions and leveraging lazy Gumbel sampling (Mussmann et al., 2017 ) with k 𝑘 k italic_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.

## 1 Introduction

Transformer models have become the dominant neural architecture across language, vision, and other domains (Vaswani, ; Dosovitskiy et al., ).
However, scaling them to handle larger input sequences remains a significant challenge (Tay et al., ), primarily due to the quadratic complexity of computing self-attention.
Overcoming this limitation is crucial for advancing neural networks.
Extending context length would enable Transformers to tackle complex tasks like book summarization (Kryściński et al., ) and time-series forecasting (Wen et al., ; Zeng et al., ; Zhou et al., ).
Furthermore, improving attention efficiency would reduce the computational burden of training, making these models more accessible.
Bridging this “compute divide” is vital for democratizing AI (Ahmed and Wahed, ).

Efficient computation of self-attention has been a focal point of research
in recent years (Fournier et al., ).
Flash Attention (Dao et al., ) and related work
(Saha and Ye, ) optimize the exact calculation of
attention by minimizing wasted computation during GPU I/O operations.
However, most approaches focus on approximating the attention function.
Sparse Transformers improve efficiency by allowing each token to attend
to only a small subset of tokens (Meister et al., ).
These subsets are identified through deterministic methods
(Child et al., ; Guo et al., ; Soldaini and Moschitti, ; Li et al., ; Qiu et al., ; Beltagy et al., ),
randomized algorithms (Kitaev et al., ; Han et al., ; Zandieh et al., ; Pagliardini et al., ),
or adaptive techniques (Correia et al., ).
Additionally, self-attention is often approximated using low-rank
matrices and kernel methods (Wang et al., ; Tay et al., ; Xiong et al., ; Katharopoulos et al., ; Choromanski et al., ).
On the negative side, recent fine-grained complexity reductions indicate
that achieving a good approximation with sub-quadratic time is not feasible across
all scenarios (Keles et al., ; Alman and Song, ).

In this work, we focus on sparse attention methods where each token
vector $q_{i}\in\mathbb{R}^{d}$ attends to the $k$ tokens
$k_{j}\in\mathbb{R}^{d}$ with the largest inner products $q_{i}^{T}k_{j}$
(Gupta et al., ; Wang et al., ), a paradigm we refer
to as $k$NN Attention.
The Routing Transformer (Roy et al., ) was an early example,
using $k$-means clustering to ensure each query only attends to keys
within the same cluster.
Memorizing Transformers (Wu et al., ) later extended this
approach by leveraging $k$NN search within a stored memory, enabling
models to memorize new data during inference.
More recently, Unlimiformer models (Bertsch et al., ) have
improved efficiency by using a single $k$NN data structure
(or index) across all attention heads and layers.

Previous works have empirically shown that $k$NN Attention not only improves
computational efficiency, but also enhances model architectures and capabilities.
However, a rigorous theoretical analysis of $k$NN Attention is still lacking.
Key questions remain unresolved, including the precise approximation guarantees
it offers, the optimal value of $k$, and how to extend the method
to approximate the backward pass.

##### Notation

Let $Q,K,V\in\mathbb{R}^{n\times d}$ be our query, key and value matrices.
Let $q_{i}=Q_{i,:}\in\mathbb{R}^{d}$ be $i$-th row of $Q$ written as a column vector.
We will also denote the $j$-th column of $Q$ by $Q_{:,j}$.
We define $A:=QK^{T}\in\mathbb{R}^{n\times n}$ to be the attention matrix, and $O=\text{softmax}(A)\cdot V\in\mathbb{R}^{n\times d}$ to be the output of the attention function.
The softmax function is applied row-wise to $A$ and is defined as a vector valued function $\sigma:\mathbb{R}^{n}\to\mathbb{R}^{n}$:

$$ $\sigma(y_{1},...,y_{n})_{i}=\frac{\exp(y_{i})}{\sum_{s=1}^{n}\exp(y_{s})}$ $$

We also let $[n]:=\{1,2,...,n\}$ and use the notation $\text{polylog}(n)$ as a substitute of $\log^{k}(n)$ for some arbitrary constant $k\in\mathbb{Z}^{+}$ that is independent of $n$.
Finally, we use the $\widetilde{O}$ notation to hide polylogarithmic factors. We will often make use of the following boosting lemma:

###### Lemma 1 (Median-Of-Means Boosting, Chakrabarti ( 2020 ) ) .

If $\widehat{Q}$ is an unbiased estimator of some statistic,
then one can obtain an $(\varepsilon,\delta)$-multiplicative estimate of that statistic by suitably combining $K:=\frac{C}{\varepsilon^{2}}\frac{\text{Var}[\widehat{Q}]}{\mathbb{E}[\widehat
{Q}]^{2}}\ln\frac{2}{\delta}$ independent samples of $\widehat{Q}$, where $C$ is a universal constant.

For a comprehensive outline of preliminary results and theory, please refer to Appendix A.

###### Lemma 1 (Median-Of-Means Boosting, Chakrabarti ( 2020 ) ) .

### 1.1 Our Contributions

#### 1.1.1 A Theoretical Framework for k 𝑘 k italic_k NN Attention

Our work provides a theoretical framework to explain both the efficiency and
effectiveness of $k$NN Attention.
Our framework reformulates self-attention as
expectations over softmax distributions.
These expectations are approximated by sampling from each
distribution in sublinear time using Lazy Gumbel Noise Sampling.
By connecting $k$NN, $k$-Maximum Inner Product Search (MIPS), and Gumbel noise sampling, we develop a
new sub-quadratic self-attention approximation algorithm aligning with
the $k$NN Attention paradigm, as summarized in the following informal theorem:

###### Theorem 2 .

Let $Q,K,V\in\mathbb{R}^{n\times d}$ and $\varepsilon,\delta$ be positive constants. Assume $||V||_{\infty}=O(\log n)$. Then $k$NN-Attention as presented in Algorithm 2
with $k=\sqrt{n}$ outputs a matrix $\widehat{O}\in\mathbb{R}^{n\times d}$
such that:

$$ $\displaystyle|\widehat{O}_{ij}-O_{ij}|\leq\varepsilon O_{ij}$ (1) $$

for all $(i,j)\in[n]\times[d]$ with probability at least $1-\delta$ and
in sub-quadratic time and space.

###### Theorem 2 .

#### 1.1.2 Approximating the Backward Pass

Our framework can be extended to solve the problem of approximating attention gradients.
Even though backpropagation is the main memory bottleneck for large models,
few methods approximate attention gradients directly.
Alman and Song’s work () is most relevant, deriving
inapproximability results for certain parameter regimes.

We present new approximation algorithms for self-attention gradients using
$k$NN search.
A key challenge is the need to multiply by the transpose of a stochastic matrix,
which disrupts our expectation-based reformulation.
To address this, we use a Markov-Chain sampling technique, treating the attention
matrix as a transition matrix and applying a single-step iteration.
Our main theorem can be informally stated as follows:

###### Theorem 3 .

Let $\phi$ be a scalar loss function and
$\partial\phi/\partial O\in\mathbb{R}^{n\times d}$.
Then, under certain assumptions on the $||\cdot||_{\infty}$ norms of $Q,K,V$,
there exist sub-quadratic time algorithms that output estimates
$\widehat{D}^{Q},\widehat{D}^{K},\widehat{D}^{V}\in\mathbb{R}^{n\times d}$
for which with probability at least $1-\delta$ it holds that:

$$ $\displaystyle||\widehat{D}^{Q}-\partial\phi/\partial Q||_{\infty}\leq e_{Q},\, ||\widehat{D}^{K}-\partial\phi/\partial K||_{\infty}\leq e_{K}\,\text{ and }\, ||\widehat{D}^{V}-\partial\phi/\partial V||_{\infty}\leq e_{V}$ (2) $$

where $e_{Q},e_{K},e_{V}$ are explicit error parameters that can roughly
be bounded by $O(\varepsilon n\cdot\text{polylog}(n))$

Algorithm
4 computes $\widehat{D}^{V}$, while Algorithms for $\widehat{D}^{K}$ and $\widehat{D}^{Q}$ can be found in Appendices E and F.

###### Theorem 3 .

## 2 k 𝑘 k italic_k NN Attention as an Approximation Algorithm

### 2.1 Reformulating Self-Attention as Expectation

Our first contribution is viewing the self-attention output as a matrix of
expectations under various softmax distributions.
Let $D_{i}$ be the softmax distribution defined by $D_{i}(j)\propto\exp(q_{i}^{T}\cdot k_{j})$ over $[n]$. Then, notice that we can write:

$$ $\displaystyle O_{ij}$ $\displaystyle=\sum\limits_{k=1}^{n}\frac{\exp(q_{i}^{T}\cdot k_{k})}{\sum_{s=1 }^{n}\exp(q_{i}^{T}k_{s})}\cdot V_{kj}=\sum\limits_{k=1}^{n}D_{i}(k)\cdot V_{ kj}=\mathop{\mathbb{E}}_{k\sim D_{i}}\left[V_{kj}\right]$ (3) $$

Thus, to approximate $O_{ij}$, we have to estimate the expected value in Equation 3.
Let $k$ be sampled according to $D_{i}$.
Then, the estimator $\widehat{O}_{ij}=V_{kj}$ is unbiased, as $\mathop{\mathbb{E}}_{k\sim D_{i}}[\widehat{O}_{ij}]=O_{ij}$. By imposing certain assumptions on the entries of the V matrix, we can bound the
variance of $\widehat{O}_{ij}$ and use boosting to obtain explicit error
guarantees:

###### Theorem 4 .

Suppose $||V||_{\infty}\leq B=O(\log(n))$ and assume that for any $i\in[n]$ we can sample from $D_{i}$ in time $O(T)$. Then, there exists an algorithm to output a matrix $\widehat{O}\in\mathbb{R}^{n\times d}$
such that:

$$ $\displaystyle|\widehat{O}_{ij}-O_{ij}|\leq\varepsilon O_{ij}$ (4) $$

for all $(i,j)\in[n]\times[d]$ with probability at least $1-\delta$,
where $\varepsilon,\delta>0$ are constants. The algorithm runs in $O(nd\cdot T\cdot\varepsilon^{-2}\log(nd/\delta)\log n)$ time.

###### Proof.

Given that $\widehat{O}_{ij}$ is an unbiased estimator of $O_{ij}$, we can utilize Lemma 1 to get an
$(\varepsilon,\delta)$-multiplicative estimator for $O_{ij}$.
To determine a sufficient number of samples of $\widehat{O}_{ij}$, we first bound the variance of our estimator:

$$ $\displaystyle\text{Var}\left[\widehat{O}_{ij}\right]\leq\mathop{\mathbb{E}}_{k \sim D_{i}}\left[V_{kj}^{2}\right]=\sum\limits_{k=1}^{n}D_{i}(k)V_{kj}^{2}\leq B \cdot O_{ij}$ (5) $$

Then, the number of samples required is:

$$ $\displaystyle O\left(\varepsilon^{-2}\cdot\log(1/\delta)\cdot\text{Var}\left[ \widehat{O}_{ij}\right]\cdot{\mathbb{E}\left[\widehat{O}_{ij}\right]^{-2}} \right)=O\left(\varepsilon^{-2}\cdot\log(1/\delta)\log n\right)$ (6) $$

due to our assumption on $||V||_{\infty}$.
To ensure that all $nd$ elements of $O$ are approximated within the desired
guarantees, we have to set $\delta^{\prime}:=\delta/(nd)$ and union-bound over
all $nd$ elements of $O$.
Since each sample requires $O(T)$ time, we arrive at the desired time complexity.
∎

###### Theorem 4 .

###### Proof.

### 2.2 Efficient Sampling from D i subscript 𝐷 𝑖 D_{i} italic_D start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT via Lazy Gumbel Sampling

Theorem 4 previously assumed we could directly sample from the distribution $D_{i}$. The Lazy Gumbel Sampling method proposed by Mussmann et al. () provides a way to sample from each $D_{i}$ in sublinear time, even with limited knowledge of $D_{i}$. However, there is an initial pre-processing step that takes a bit more than linear time across all the distributions.

Fix some $i\in[n]$ and let $Z_{ij}=q_{i}^{T}k_{j}$. In the Gumbel Max Trick (Lemma 19),
we form the random variables $N_{ij}=Z_{ij}+G_{ij}$ where
$G_{ij}\sim\text{Gumbel}(0,1)$ for all $j\in[n]$ and sample
$\arg\max N_{ij}$. This is equivalent to sampling $j\in[n]$ from the softmax distribution over the $Z_{ij}$ scores.
Mussmann et al. () observed that if we have the top $k$ $Z_{ij}$ values in a set $S_{i}$ and add
Gumbel noise to just them, then for any $j\notin S_{i}$ to be ultimately picked, its Gumbel noise $G_{ij}$ must be quite large.
We can use the concentration properties of the Gumbel distribution to
argue that in expectation we only need to sample $\frac{n}{k}$
elements not in $S_{i}$.
Setting $k=\sqrt{n}$ allows us to balance the two, resulting in a
sublinear time algorithm for sampling from $D_{i}$. An illustration of the idea can be seen in Figure 1, as it was presented in Mussmann et al. ().

Figure: Algorithm 1 Lazy Gumbel Sampling from $D_{i}$, for some $i\in[n]$

We can see that this method samples exactly from $D_{i}$:

###### Theorem 5 (Correctness of Algorithm 1 ) .

After running Algorithm 1, it holds that:

$$ $\displaystyle\widehat{j}=\arg\max\limits_{j\in[n]}\{q_{i}^{T}k_{j}+G_{ij}\}$ (7) $$

where $G_{ij}\sim\text{Gumbel}(0,1)$. In other words, $\widehat{j}$ is
sampled according to $D_{i}$.

###### Proof.

The only way that we do not find the maximum is if one of the points in $[n]\setminus(S_{i}\cup T_{i})$ are the true maximum. However those points (by construction)
have Gumbel noise at most $B$, so they cannot be the overall maximum.
∎

Figure: Figure 1: Lazy Gumbel sampling
Refer to caption: extracted/5985875/lazy_gumbel_3.png

In Appendix B, we show that the expected number $m$ of large Gumbels is at most $n/k$.
Our simplified proof uses the Gumbel distribution’s Moment Generating Function,
rather than the original exponential-based analysis.

###### Lemma 6 .

The following holds:

$$ $\displaystyle\mathbb{E}\left[m\right]\leq\frac{n}{k}$ (8) $$

Due to Lemma 6, we see that we need to set
$k=\sqrt{n}$ to optimize our overall time complexity.
We have the following theorem, which follows easily from the pseudocode of
Algorithm 1 and Lemma 6:

###### Theorem 7 .

Let $k=\sqrt{n}$. Suppose that we are able to retrieve the set $S_{i}$ in
$f(n,k)$ time.
Then, we can use Algorithm 1 to sample from $D_{i}$ in
$O(\sqrt{n}+f(n,\sqrt{n}))$ time in expectation.

###### Theorem 5 (Correctness of Algorithm 1 ) .

###### Proof.

###### Lemma 6 .

###### Theorem 7 .

#### 2.2.1 Obtaining the top k 𝑘 k italic_k inner products

Algorithm 1 relies on obtaining the set $S_{i}$ of the top $\sqrt{n}$ inner products $q_{i}^{T}k_{j}$ for each $i\in[n]$ in sub-quadratic $f(n,\sqrt{n})$ time. Since the $k_{j}$ vectors are fixed, while the $q_{i}$ vectors act as queries, this setup is known as the $k$-Maximum Inner Product Search Problem (MIPS)

The $k$-MIPS problem can be reduced to the $k$NN problem using a transformation proposed by Neyshabur and Srebro (). We add an extra dimension to normalize all key vectors. Specifically, the inner product $q_{i}^{T}k_{j}$ can be expressed as:

$$ $\displaystyle q_{i}^{T}k_{j}=\frac{1}{2}\left(||q_{i}||_{2}^{2}+||k_{j}||_{2}^ {2}-||q_{i}-k_{j}||_{2}^{2}\right)$ (9) $$

If the norms $||k_{j}||_{2}$ are the same across all $j$, the problem reduces to finding the $k$ nearest neighbors to $q_{i}$. To enforce this, we define:

$$ $\displaystyle\left(k^{\prime}_{j}\right)^{T}=\left[k_{j}^{T},\sqrt{M-||k_{j}|| _{2}^{2}}\right]$ (10) $$

so that $||(k^{\prime})_{j}||_{2}=M$ for all $j\in[n]$, where $M$ is a previously known upper bound. When querying with $q_{i}$, we use:

$$ $\displaystyle(q^{\prime}_{i})^{T}=\left[q_{i}^{T},0\right]$ (11) $$

This transformation preserves the inner products, allowing us to solve the $k$NN problem for $q^{\prime}_{i}$. We can then use a $k$NN index $H$ to preprocess $K$ and query it with each $q^{\prime}_{i}$ to construct $S_{i}$ for all $i\in[n]$. We remain agnostic to the specific $k$NN index one could use for this algorithm(^1^11For a specific construction with precise theoretical guarantees that uses Locality Sensitive Hashing (LSH), please refer to Appendix C.), but if we assume that the construction runtime is slightly larger than linear and the query time slightly larger than $k$, then $k$NN attention techniques have total runtime of $\approx\widetilde{O}(dn^{3/2}\cdot\varepsilon^{-2}\log(1/\delta))$ time and space.

Figure: Algorithm 2 $k$NN Attention

### 2.3 k 𝑘 k italic_k NN Attention without Median-of-Means

This section describes a simpler algorithm for computing the expected value needed for self-attention. The algorithm still uses $k$NN indices to find the top $k$ inner products per query, but usually outperforms Algorithm 2 in practice, due to its amenity for hardware-accelerated vectorization, and is thus our preferred implementation for experiments(^2^22See Appendix G for a PyTorch implementation of this algorithm.).

Building on Mussmann et al. (), the algorithm estimates $\mathbb{E}_{k\sim D_{i}}[V_{kj}]$ using set $S_{i}$ by sampling $\ell$ additional vectors outside $S_{i}$ (set $T_{i}$) and upweighting them in the expectation sum, as follows:

$$ $\displaystyle\widehat{O}_{ij}=\frac{\sum_{s\in S_{i}}e^{q_{i}^{T}k_{s}}\cdot V _{sj}+\frac{n-k}{\ell}\sum_{s\in T_{i}}e^{q_{i}^{T}k_{s}}\cdot V_{sj}}{\sum_{s \in S_{i}}e^{q_{i}^{T}k_{s}}+\frac{n-k}{\ell}\sum_{s\in T_{i}}e^{q_{i}^{T}k_{s }}}$ (12) $$

The quality of this estimator and the optimal choices for $k$ and $\ell$ are derived as follows:

###### Theorem 8 .

The estimator $\widehat{O}_{ij}$ satisfies the following error guarantee with probability at least $1-\delta$:

$$ $\displaystyle\left|\widehat{O}_{ij}-O_{ij}\right|=O(\varepsilon)$ $$

if the following two conditions hold:
$k^{2}\ell\geq 8n^{2}\varepsilon^{-2}\log\left(4/\delta\right)$ and $k\ell\geq 2n\varepsilon^{-2}\log\left(2/\delta\right)$. Setting $k=\ell=O\left(n^{2/3}\varepsilon^{-1}\sqrt{\log(1/\delta)}\right)$
gives us an $\widetilde{O}\left(dn^{5/3}\varepsilon^{-1}\sqrt{\log(1/\delta)}\right)$ algorithm for estimating self-attention within additive error $O(\varepsilon)$, assuming an efficient $k$NN implementation.

###### Proof.

The proof of the additive error guarantee can be found in Mussmann et al. ().
∎

###### Theorem 8 .

###### Proof.

## 3 Approximating the Attention Gradients

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 Appendix D.

###### Lemma 9 (Attention Gradients) .

Let $Q,K,V\in\mathbb{R}^{n\times d}$. Let $P:=\text{softmax }(QK^{T})\in\mathbb{R}^{n\times n}$ be the normalized attention matrix. Let $\phi$ be a scalar function of $O$ and $D^{O}=\partial\phi/\partial O\in\mathbb{R}^{n\times d}$. Similarly define $D^{Q},D^{K}$ and $D^{V}$. The following relationships hold:

$$ $\displaystyle D^{V}$ $\displaystyle=P^{T}\cdot D^{O}$ (13) $\displaystyle D^{Q}_{ij}$ $\displaystyle=\sum\limits_{k=1}^{n}P_{ik}\left(D^{P}_{ik}-\langle D^{P}_{i,:}, P_{i,:}\rangle\right)K_{kj}$ (14) $\displaystyle D^{K}_{ij}$ $\displaystyle=\sum\limits_{k=1}^{n}P_{ki}\left(D^{P}_{ki}-\langle D^{P}_{k,:}, P_{k,:}\rangle\right)Q_{kj}$ (15) $$

where $D^{P}_{ij}=\partial\phi/\partial P_{ij}=\langle D^{O}_{i,:},V_{j,:}\rangle$.

Clearly, calculating $D^{Q},D^{K}$ and $D^{V}$ naively requires storing $P$, which requires $O(dn^{2})$ time.

###### Lemma 9 (Attention Gradients) .

### 3.1 Estimating D V superscript 𝐷 𝑉 D^{V} italic_D start_POSTSUPERSCRIPT italic_V end_POSTSUPERSCRIPT using Random Walk Simulations

We now give an algorithm for estimating $D^{V}$. Suppose we want to calculate the $j$-th column of $D^{V}$:

$$ $\displaystyle D^{V}_{:,j}=P^{T}\cdot D^{O}_{:,j}$ (16) $$

for $j\in[d]$. Fix $\vv{x_{j}}:=D^{O}_{:,j}\in\mathbb{R}^{n\times 1}$ and suppose that $\vv{x_{j}}\geq 0$. We will relax this assumption in Section 3.1.1. Then, $\vv{y_{j}}:=\vv{x_{j}}/||\vv{x_{j}}||_{1}$ is a distribution over the universe $[n]$. Imagine a random walk over $[n]$ with transition matrix $P$ and initial distribution $\vv{y_{j}}$. Then:

$$ $\displaystyle\vv{\pi_{j}}:=P^{T}\cdot\vv{y_{j}}$ (17) $$

is the distribution after one step in the process. Thus, we can estimate $\vv{\pi_{j}}$ with Markov Chain simulations, by first picking an item $i\in[n]$ from the distribution $\vv{y_{j}}$, and then picking another item $k\in[n]$ with probability $P_{ik}$. We make $N$ independent length-$1$ random walks like this and let:

$$ $\displaystyle X_{v}^{(j,s)}=\begin{cases}1,&\text{ if the $s$-th walk ends up in state $v$}\\ 0,&\text{ otherwise}\end{cases}$ $$

We know that $\mathbb{E}[X_{v}^{(j,s)}]=\pi_{j}(v)$ for all $s\in[N]$. Thus, we can form a boosted estimator:

$$ $\displaystyle\widehat{p_{j}}(v)=\frac{1}{N}\sum\limits_{s=1}^{N}X_{v}^{(j,s)}$ (18) $$

Figure: Figure 2: A single-step Markov Chain sample.
Refer to caption: extracted/5985875/markov_chain_3.png

This estimator is unbiased due to linearity of expectation, so we can use the Hoeffding bound to ensure that our empirical distribution is close to the true distribution as long as we take enough samples:

$$ $\displaystyle\Pr\left[|\widehat{p_{j}}(v)-\pi_{j}(v)|\geq\varepsilon\right] \leq 2\exp(-2N\varepsilon^{2})$ $$

We set the probability of failure to $1/(dn^{2})$ so that we can union bound over all $v\in[n]$ and all $j\in[d]$. It follows that we require:

$$ $\displaystyle N=\Theta(\varepsilon^{-2}\ln(nd))$ (19) $$

Of course, we need to scale $\vv{\pi_{j}}$ back to recover $D^{V}_{:,j}$. We define:

$$ $\displaystyle\widehat{D}^{V}_{:,j}=||\vv{x_{j}}||_{1}\cdot\widehat{p_{j}}$ (20) $$

Then we get that with probability at least $1-1/n$ it holds for all $j\in[d]$ that:

$$ $\displaystyle\left|\left|\widehat{D}^{V}_{:,j}-D^{V}_{:,j}\right|\right|_{ \infty}=||\vv{x_{j}}||_{1}\cdot\left|\left|\widehat{p_{j}}-\vv{\pi_{j}}\right| \right|_{\infty}\leq\varepsilon||\vv{x_{j}}||_{1}$ (21) $$

#### 3.1.1 Relaxing the non-negativity assumption

We now relax the non-negativity constraint on $\vv{x_{j}}$, accepting some approximation error. Since normalizing $\vv{x_{j}}$ with its L1 norm fails if $\vv{x_{j}}$ has negative entries, we adopt a numerical stability technique to ensure we get a valid probability distribution even when $\vv{x_{j}}$ contains negative entries. Let

$$ $\displaystyle M_{j}:=-\min\limits_{\begin{subarray}{c}v\in[n]\\ (\vv{x_{j}})_{v}\leq 0\end{subarray}}(\vv{x_{j}})_{v}$ (22) $$

be the absolute value of the most negative entry of $\vv{x_{j}}$. If $\vv{x_{j}}\geq 0$, then set $M_{j}=0$. Now, if $\vv{M_{j}}:=M_{j}\cdot 1^{n}\in\mathbb{R}^{n\times 1}$, then $\vv{x_{j}^{\prime}}=\vv{x_{j}}+\vv{M_{j}}\geq 0$. Therefore, we can estimate

$$ $\displaystyle\widehat{p^{\prime}}_{j}\approx\vv{\pi_{j}^{\prime}}:=P^{T}\cdot \vv{x_{j}^{\prime}}$ (23) $$

using our Markov Chain method. Going back to our original goal of estimating $\vv{\pi}_{j}$, we have:

$$ $\displaystyle\vv{\pi_{j}}$ $\displaystyle:=P^{T}\cdot\vv{x_{j}}$ $\displaystyle=P^{T}\cdot(\vv{x_{j}^{\prime}}-\vv{M_{j}})$ $\displaystyle=\vv{\pi_{j}^{\prime}}-P^{T}\cdot\vv{M_{j}}$ $\displaystyle=\vv{\pi_{j}^{\prime}}-M_{j}\cdot P^{T}\cdot 1^{n}$ $$

where $1^{n}$ is the all $1$-s vector. So then we only need to additionally estimate $P^{T}\cdot 1^{n}$. This can be done in the same fashion only once, as a pre-processing step. Specifically, suppose that we estimate $P^{T}\cdot 1^{n}$ as $\widehat{s}$. We know from our prior analysis that using $\Theta(\varepsilon^{-2}\log n)$ random walks we get an estimate $\widehat{s}$ such that:

$$ $\displaystyle\left|\left|\widehat{s}-P^{T}\cdot 1^{n}\right|\right|_{\infty} \leq\varepsilon n$ (24) $$

Putting it all together, our final estimator is then:

$$ $\displaystyle\widehat{p}_{j}:=\widehat{p^{\prime}}_{j}-M_{j}\cdot\widehat{s}$ (25) $$

Eventually, the total error for estimating $\vv{\pi_{j}}$ becomes:

$$ $\displaystyle\left|\left|\widehat{D}^{V}_{:,j}-D^{V}_{:,j}\right|\right|_{ \infty}=\left|\left|\widehat{p}_{j}-\vv{\pi_{j}}\right|\right|_{\infty}$ $\displaystyle=\left|\left|\widehat{p^{\prime}}_{j}-M_{j}\cdot\widehat{s}-\vv{ \pi^{\prime}_{j}}+M_{j}\cdot P^{T}\cdot 1^{n}\right|\right|_{\infty}$ (26) $\displaystyle\leq\left|\left|\widehat{p^{\prime}}_{j}-\vv{\pi^{\prime}_{j}} \right|\right|_{\infty}+M_{j}\cdot\left|\left|P^{T}1^{n}-\widehat{s}\right| \right|_{\infty}$ (27) $\displaystyle\leq\varepsilon\left|\left|\vv{x^{\prime}_{j}}\right|\right|_{1}+ \varepsilon M_{j}\cdot n$ (28) $\displaystyle=\varepsilon\langle x_{j},1^{n}\rangle+2\varepsilon nM_{j}$ (29) $$

where the first inequality follows from the triangle inequality and the last equality follows from

$$ $\displaystyle\left|\left|\vv{x^{\prime}_{j}}\right|\right|_{1}=\sum\limits_{k= 1}^{n}\left|(x_{j})_{k}+M_{j}\right|=\sum\limits_{k=1}^{n}\left[(x_{j})_{k}+M_ {j}\right]=\langle x_{j},1^{n}\rangle+nM_{j}$ (30) $$

We present the entirety of the method as Algorithm 4.

Figure: Algorithm 3 Estimating $P^{T}x$, with query access to $P\in\mathbb{R}^{n\times n}$ a stochastic matrix

Figure: Algorithm 4 Estimating $D^{v}V$

#### 3.1.2 Runtime analysis

For each $j\in[d]$ we take $N=O(\varepsilon^{-2}\log n)$ samples. We can take one sample in $O(nd)$ time. In addition, we must pre-calculate the sums $\langle x_{j},1^{n}\rangle+nM_{j}=nM_{j}+\sum_{k=1}^{n}D^{O}_{kj}$ for all $j\in[d]$, which takes $O(nd)$ time. As a result, we arrive at the following theorem:

###### Theorem 10 .

Given $Q,K,V$ and $D^{O}$, Algorithm 4 calculates ${\partial\phi}/{\partial V_{ij}}=D^{V}_{ij}$ for all $(i,j)\in[n]\times[d]$ within an additive approximation error of

$$ $\displaystyle e_{V}=\varepsilon\cdot\langle D^{O}_{:,j},1^{n}\rangle+2n \varepsilon M_{j},\,\text{where }M_{j}:=-\min_{i\in[n],D^{O}_{ij}\leq 0}D^{O}_ {ij}$ (31) $$

with probability at least $1-\frac{1}{n}$. The time complexity is $O(nd^{2}\varepsilon^{-2}\log n)$.

###### Remark 11 .

Note that Algorithm 4 does not materialize the $P$ matrix. Instead it accesses its elements by using $Q$ and $K$ in $O(d)$ time per element.

###### Theorem 10 .

###### Remark 11 .

## 4 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 $k$NN Attention in approximating the attention function. We randomly sample matrices $Q,K,V\in\mathbb{R}^{n\times d}$ from a uniform distribution over $[-B,B]^{n\times d}$ and assess the approximation quality on these inputs. Our focus is on the “classic” $k$NN Attention estimator (Theorem 8) with $\lambda=1$, as used in implementations like Bertsch et al. () and Wu et al. (). We vary $k$ to study how the error decreases as $k$ increases and compare the efficiency to the naive $O(n^{2})$ attention, expecting notable performance gains. This experiment is implemented in PyTorch, running on a MacBook Air with an M3 CPU and 8GB of RAM.

##### Efficiency of kNN Attention

Our experiments confirm $k$NN Attention’s superior speed, demonstrating sub-quadratic scaling. With a batch size of $1$ and $H=10$ attention heads, it handles self-attention for $n=10^{6}$, while the naive method runs out of memory beyond $n\geq 20000$. Increasing $k$ further leads to memory errors for $n\geq 50000$, highlighting kNN Attention’s memory efficiency. Detailed results are in Figure LABEL:fig:efficiency.

##### Role of k 𝑘 k italic_k in the Approximation Error

We investigate the impact of $k$ on the approximation error, predicting that error increases as $k$ decreases. The experiment confirms this, showing that for $k\geq n^{1/8}$, the error is minimal. Our theory suggests a threshold closer to $\sqrt{n}$ which indicates that the optimal $k$ may vary by dataset. The results appear in Figure LABEL:fig:error. Interestingly, the error is more pronounced for small values of both $B$ and $k$, potentially due to the limited approximation power when $k$ 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 a normal distribution, as this strategy aligns with typical neural network weight initialization strategies, and approximate $D^{Q}$ and $D^{V}$ 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 $\alpha$ is varied between 0.05 and 0.5, while the error parameter is fixed at $\varepsilon=0.05$ and the confidence parameter at $\delta=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 PyTorch’s autograd to compute the exact attention gradients. As shown in Figure LABEL:fig:grad-desc-experiment, our approximation closely matches the expected results in the convex case but deviates from the optimal convergence in the non-convex case. A 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 $k$NN attention into LLMs to study its impact on training and inference. While previous work has explored $k$NN indices for efficient fine-tuning and training (Bertsch et al., ; Wu et al., ), our goal is to understand how Transformer LLMs respond to attention function approximation, linking it to our theory and providing practical guidelines. The architecture and training methods are adapted from nanoGPT (Karpathy, ), and our experiments are conducted on an NVIDIA L40 GPU with 48GB of memory.

Our first experiment trains a mini character-level Transformer on a small Shakespeare dataset, replacing attention with $k$NN attention. We compare training and validation perplexity to the exact method. Results in Figure LABEL:fig:shakespeare_experiment show that $k$NN Attention maintains a small perplexity gap. However, as overfitting occurs, the perplexity difference widens, possibly due to increasing maximum approximation error. Future work could explore the impact of larger $k$ values on perplexity.

We also experiment with fine-tuning a large pre-trained LLM using $k$NN attention, in the hopes that the approximation will not severly degrade the model’s quality. We manage to fine-tune GPT-2 XL(^3^33$1.61$B parameters, see Radford et al. ()) on a Shakespeare dataset - a task typically infeasible on a single L40 GPU due to memory constraints with quadratic attention. Examples from prompting this model can be found in Appendix H.

## 5 Conclusion

In this work, we developed a theoretical framework for leveraging $k$NN techniques to design
efficient 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 $k$NN approximations
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 $k$ value is often significantly smaller than the predicted $\sqrt{n}$.

In conclusion, by applying sublinear algorithm techniques, our work provides a solid foundation
for making Transformers more scalable and efficient while identifying critical areas for further
research in LLM approximation.

## Reproducibility Statement

To aid in the reproduction of our experiments, we include our code in the repository at [https://github.com/sansui-123/knn_attention](https://github.com/sansui-123/knn_attention). 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 6.

## Appendix

In the following sections we deposit theoretical results, proofs and algorithms that are missing from the main paper, either due to space constraints or for the sake of clarity.

## Appendix A Preliminaries

### A.1 Self-Attention and Approximation

We start by defining self-attention and some of its variants.

###### Definition 12 (Self-Attention) .

Let $Q,K,V\in\mathbb{R}^{n\times d}$. We can think of these matrices as a collection of $n$ $d$-dimensional query, key and value vectors respectively. We define the self-attention function as follows:

$$ $\displaystyle O(Q,K,V)=D^{-1}AV$ (32) $$

where $A=\exp(QK^{T})\in\mathbb{R}^{n\times n}$ is the attention matrix and $D=\text{diag}(A1^{n})$. $D$ effectively implements taking a row-wise softmax of the entries of $QK^{T}$. Many modern implementations of attention consider causal attention, in which we mask away the upper-triangular entries of $A$.

###### Remark 13 (Normalization by $\sqrt{d}$ ) .

In most implementations we divide $QK^{T}$ by $\sqrt{d}$ (Vaswani, ) because it reduces the variance of each element of $A$ had $Q,K,T$ been selected from a uniform distribution. We will omit this technicality because it does not affect our algorithmic techniques. For the rest of this paper we will assume that $d^{-1/2}$ has been pre-normalized into $K$.

###### Remark 14 (Dropout) .

Many attention implementations also use dropout. Every entry of $A$ will be masked to 0 0 with probability $p$, where $p$ is set to a small constant, like $0.1$.

###### Definition 15 ( $(\varepsilon,\delta)$ -estimators) .

Let $X$ be a statistic and $\widehat{X}$ be an estimator we have for it. $\widehat{X}$ is an $(\varepsilon,\delta)-$additive estimator if with probability at least $1-\delta$ it holds that

$$ $\left|X-\widehat{X}\right|\leq\varepsilon$ $$

Respectively, we call the estimator multiplicative if

$$ $\left|X-\widehat{X}\right|\leq\varepsilon X$ $$

We will often make use of the following boosting lemma from the theory of randomized algorithms:

###### Lemma 16 (Median-Of-Means Amplification Technique, (Chakrabarti, 2020 ) ) .

If $\widehat{Q}$ is an unbiased estimator of some statistic,
then one can obtain an $(\varepsilon,\delta)$-multiplicative estimate of that statistic by suitably combining

$$ $K:=\frac{C}{\varepsilon^{2}}\ln\frac{2}{\delta}\frac{\text{Var}[\widehat{Q}]}{ \mathbb{E}[\widehat{Q}]^{2}}$ $$

independent samples of $\widehat{Q}$, where $C$ is a universal constant.

###### Definition 12 (Self-Attention) .

###### Remark 13 (Normalization by d 𝑑 \sqrt{d} square-root start_ARG italic_d end_ARG ) .

###### Remark 14 (Dropout) .

###### Definition 15 ( ( ε , δ ) 𝜀 𝛿 (\varepsilon,\delta) ( italic_ε , italic_δ ) -estimators) .

###### Lemma 16 (Median-Of-Means Amplification Technique, (Chakrabarti, 2020 ) ) .

### A.2 Gumbel Noise

The Gumbel Distribution will be useful for sampling from softmax distributions. We define it below:

###### Definition 17 (Gumbel Distribution) .

The Gumbel distribution with mean $\mu$ and parameter $\beta$ has the following probability density function:

$$ $\displaystyle\text{Gumbel}(\mu,\beta)(x)=\frac{1}{\beta}e^{-e^{-(x-\mu)/\beta}}$ (33) $$

We will make use of the following properties of the Gumbel Distribution:

###### Lemma 18 (Gumbel Distribution Properties, (Chattamvelli and Shanmugam, 2021 ) ) .

The following are true:

- •
The mean of a Gumbel distribution is $\mu+\beta\gamma$(^4^44 $\gamma\approx 0.577$ is the Euler-Mascheroni constant.)
- •
The moment generating function (MGF) of the Gumbel$(\mu,\beta)$ distribution is:
$\displaystyle M(t)=\Gamma(1-\beta t)e^{\mu t}$
(34)
where $\Gamma$ is the Gamma function.
- •
We can easily
sample a Gumbel random variable of mean $\mu$ and parameter $\beta$ by using the uniform distribution:
$\displaystyle X=\mu-\beta\ln(-\ln(U)),\quad\text{ where }U\sim\text{unif}([0,1])$
(35)
- •
Let $M_{1},...,M_{n}$ be $(\mu,\beta)$ independent Gumbel random variables. Then,
$M:=\max_{i\in[n]}M_{i}$
is a $(\mu+\beta\ln n,\beta)$ Gumbel random variable.

Next, we present the well-known Gumbel Max Trick, an alternative way to sample from a softmax distribution:

###### Lemma 19 (Gumbel-Max-Trick, (Huijben et al., 2022 ) ) .

Let $x_{1},...x_{n}$ be real numbers and consider the softmax categorical distribution $p$ where

$$ $p_{i}=\frac{\exp(x_{i})}{\sum\limits_{k=1}^{n}\exp(x_{i})}$ $$

Consider sampling $n$ Gumbel random variables $G_{1},...,G_{n}\sim\text{Gumbel}(0,1)$ and let

$$ $\widehat{i}\in\arg\max\limits_{i\in[n]}\{x_{i}+G_{i}\}$ $$

Then $\widehat{i}$ is distributed according to $p$.

###### Definition 17 (Gumbel Distribution) .

###### Lemma 18 (Gumbel Distribution Properties, (Chattamvelli and Shanmugam, 2021 ) ) .

###### Lemma 19 (Gumbel-Max-Trick, (Huijben et al., 2022 ) ) .

### A.3 Locality Sensitive Hashing

In our theoretical exposition we will make extensive use of schemes for approximate nearest neighbor search. A very successful such suite of algorithms with a long history (Andoni et al., , ) of provable theoretical guarantees is Locality Sensitive Hashing (LSH):

###### Theorem 20 (Existence of LSH, (Gionis et al., 1999 ; Mussmann et al., 2017 ) ) .

Let $V\subseteq U$ be a set of size $n$ with a similarity
measure $\text{Sim}(\cdot,\cdot)$. Consider a hash family $H$ such that for scalars $S_{1}>S_{2}$ and $p_{1}>p_{2}$:

- •
For any $x,y\in V$ where $\text{Sim}(x,y)\geq S_{1}$, $\Pr_{h\in H}[h(x)=h(y)]\geq p_{1}$
- •
For any $x,y\in V$ where $\text{Sim}(x,y)\leq S_{2}$, $\Pr_{h\in H}[h(x)=h(y)]\leq p_{2}$

This is called an $(S_{1},S_{2},p_{1},p_{2})$-Locality Sensitive Hash Family. Given such a family, one can construct a data structure which, given any query $q\in U$, does the following with high probability: if there exists some point $v\in V$ with $\text{Sim}(v,q)\geq S_{1}$, it returns a point $v^{\prime}\in V$ with $\text{Sim}(v^{\prime},q)\geq S_{2}$. If no point $v\in V$ exists with $\text{Sim}(v,q)\geq S_{2}$, it returns a negative answer $\perp$. Further, this can be done with $\widetilde{O}(n^{\rho})$ query time and $\widetilde{O}(n^{1+\rho})$ space where $\rho=\log p_{1}/\log p_{2}<1$.

LSH also finds numerous applications in solving the Maximum Inner Product Search Problem (MIPS) as shown in Neyshabur and Srebro (), Shrivastava and Li (), and others.

###### Theorem 20 (Existence of LSH, (Gionis et al., 1999 ; Mussmann et al., 2017 ) ) .

### 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 21 (Random Walks) .

Consider a state space $V=[n]$ and a weighted complete graph on $V$ with weights $w$ in $[0,1]$ such that for all $u\in V$

$$ $\sum\limits_{v\in V}w_{uv}=1$ $$

This graph represents a random walk with transition matrix $P\in[0,1]^{n\times n}$, where $P_{ij}=w_{ij}$. $P$ is a stochastic matrix because its rows sum to $1$. In a random walk, we start at some vertex and choose a 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 a distribution $p$ over $V$ and we do a single step in the random walk, we can obtain the resulting distribution by multiplying the original distribution with $P^{T}$:

###### Lemma 22 (Single Step Random Walk Transition) .

Consider a distribution $p\in\Delta(n)$(^5^55$\Delta(n):=\{x\in\mathbb{R}^{n}\mid x\geq 0,||x||_{1}=1\}$ is the probability simplex over $[n]$.) over $V$. Then, the quantity $P^{T}\cdot p$ gives the distribution over $V$ after one step of the random walk.

###### Proof.

Let $q$ be the distribution after one step. Let $v\in V$. We have by law of total probability that:

$$ $q(v)=\sum\limits_{u\in V}p(u)\cdot w_{uv}=(P^{T}p)_{v}$ $$

∎

Finally, we state a well-known concentration result about independent random variables:

###### Lemma 23 (Hoeffding Bound) .

Let $X_{1},...,X_{n}$ be independent random variables where $a_{i}\leq X_{i}\leq b_{i}$ almost surely. Let $S_{n}:=X_{1}+\cdots+X_{n}$. We have that:

$$ $\Pr\left[\left|S_{n}-\mathbb{E}[S_{n}]\right|\geq t\right]\leq 2\exp\left(- \frac{2t^{2}}{\sum\limits_{i=1}^{n}(b_{i}-a_{i})^{2}}\right)$ $$

###### Definition 21 (Random Walks) .

###### Lemma 22 (Single Step Random Walk Transition) .

###### Proof.

###### Lemma 23 (Hoeffding Bound) .

## Appendix B Proof of Lemma 6

We prove Lemma 6. Our proof deviates from the proof of Mussmann et al. () in that it uses a MGF-based argument, which we believe is cleaner.

###### Lemma 24 (Reminder) .

In the context of Algorithm 1, we have that:

$$ $\displaystyle\mathbb{E}\left[m\right]\leq\frac{n}{k}$ (36) $$

###### Proof.

By Lemma 18, we can generate Gumbel$(0,1)$ random variables as follows: Let $U_{j}$ be uniform in $[0,1]$. Then:

$$ $\displaystyle G_{ij}=-\ln(-\ln(U_{j}))$ (37) $$

is distributed according to $\text{Gumbel}(0,1)$. We want $G_{ij}>B$, which implies that:

$$ $\displaystyle-\ln(-\ln(U_{j}))>B\iff U_{j}>\exp(-\exp(-B))$ (38) $$

So the number of points for which the Gumbel noise exceeds $B$ is distributed according to the Binomial distribution with parameters $n-k$ and $1-\exp(-\exp(-B))$. If we condition on $M:=\max\limits_{j\in S_{i}}\{q_{i}^{T}k_{j}+G_{ij}\}$, we have that:

$$ $\displaystyle\mathbb{E}\left[m\mid M\right]$ $\displaystyle=(n-k)(1-\exp(-\exp(-B)))$ (39) $\displaystyle\leq n\exp(-B)$ (40) $$

where the last inequality follows by $e^{-x}\geq 1-x$:

$$ $1-\exp(-\exp(-B))\leq 1-(1-\exp(-B))=\exp(-B)$ $$

Now we can bound $\mathbb{E}[n\exp(-B)]$ by using the MGF of the Gumbel distribution (see Lemma 18). Let $M^{\prime}:=\max_{j\in S_{i}}G_{ij}$. Recall by Lemma 18 that $M^{\prime}$ is a Gumbel random variable with $\mu_{M^{\prime}}=\log k$ and $\beta_{M^{\prime}}=1$. Let $f_{M^{\prime}}(t)=\mathbb{E}[e^{tM^{\prime}}]$ be its moment generating function. We know that:

$$ $\displaystyle f_{M^{\prime}}(t)=\Gamma(1-t)\cdot e^{(\log k)t}=\Gamma(1-t) \cdot k^{t}$ (41) $$

This allows us to write:

$$ $\displaystyle\mathbb{E}[n\exp(-B)]$ $\displaystyle=n\cdot\mathbb{E}[\exp(-B)]$ (42) $\displaystyle=n\cdot\mathbb{E}[\exp(S_{\min}-M)]$ (43) $\displaystyle=n\cdot\mathbb{E}[\exp(S_{\min}-\max\limits_{j\in S_{i}}\{Z_{ij}+ G_{ij}\})]$ (44) $\displaystyle\leq n\cdot\mathbb{E}[\exp(S_{\min}-\min_{j\in S_{i}}Z_{ij}-M^{ \prime}]$ (45) $\displaystyle=n\cdot\mathbb{E}[\exp(-M^{\prime})]$ (46) $\displaystyle=n\cdot f_{M^{\prime}}(-1)$ (47) $\displaystyle=\frac{n}{k}$ (48) $$

where inequality 45 follows because

$$ $\max_{j\in S_{i}}\{Z_{ij}+G_{ij}\}\geq\min_{j\in S_{i}}\{Z_{ij}\}+\max_{j\in S _{i}}G_{ij}=S_{\min}+M^{\prime}$ $$

For a quick proof of this statement, let $\widehat{j}:=\arg\min_{j\in S_{i}}Z_{ij}$ and $\widetilde{j}:=\arg\max_{j\in S_{i}}G_{ij}$. Also let $j^{*}:=\arg\max_{j\in S_{i}}\{Z_{ij}+G_{ij}\}$. Then we have:

$$ $\displaystyle\max\limits_{j\in S_{i}}\{Z_{ij}+G_{ij}\}$ $\displaystyle=Z_{ij^{*}}+G_{ij^{*}}$ (49) $\displaystyle\geq Z_{i\widetilde{j}}+G_{i\widetilde{j}}$ (50) $\displaystyle\geq Z_{i\widehat{j}}+G_{i\widetilde{j}}$ (51) $\displaystyle=S_{\min}+M^{\prime}$ (52) $$

Finally 48 follows because $\Gamma(2)=2!=1$. Now, via law of total expectation we finally get:

$$ $\displaystyle\mathbb{E}[m]=\mathbb{E}_{M}\left[\mathbb{E}\left[m\mid M\right] \right]\leq\mathbb{E}_{M}\left[\frac{n}{k}\right]=\frac{n}{k}$ (53) $$

∎

###### Lemma 24 (Reminder) .

###### Proof.

## Appendix C k 𝑘 k italic_k NN Attention via Concentric LSH

In the main paper, we abstracted away the specific $k$NN method used to obtain the top-$k$ key vectors $k_{j}$ for every query vector $q_{i}$. 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 2, this method could substitute the $k$NN index $H$.

This approach in question was proposed by Mussmann et al. () and it uses a 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. ():

###### Definition 25 (Approximate $k$ -MIPS) .

We say that a set $S_{i}$ is an approximate top-$k$ inner product solution if $|S_{i}|=k$ and there exists a constant $c$ such that:

$$ $\displaystyle\max\limits_{j\notin S_{i}}q_{i}^{T}k_{j}-\min\limits_{j\in S_{i} }q_{i}^{T}k_{j}<c$ (54) $$

If we are able to generate an approximate solution $S_{i}$ instead of an exact one, we have to lower our threshold $B$ to $M-S_{\min}-c$ in Algorithm 1, which in turns implies that $E[m]\leq\sqrt{n}\cdot e^{c}$. This remains sublinear in $n$ because $c$ is a constant.

The solution to the approximate version of the problem is constructed using LSH. Specifically, we build a sequence of $O(\text{polylog}(n,d))$ LSH data structures, each with concentric approximation radii, and hash all the key vectors $k_{j}\in\mathbb{R}^{d}$ into them. For each query vector $q_{i}$, we hash it across all these data structures and identify the first pair of consecutive LSH structures, $D_{i}$ and $D_{i+1}$, where $D_{i+1}$ contains more than $\sqrt{n}$ points in the buckets corresponding to $q_{i}$, while $D_{i}$ contains fewer than $\sqrt{n}$ points. For further details, see Mussmann et al. (). The following theorem ultimately holds:

###### Theorem 26 ( Mussmann et al. ( 2017 ) ) .

Let $0<\rho<1$ be a constant. There exists an algorithm for solving the approximate version of $k$-MIPS on any single query $q$ with probability at least $1-1/n^{2}$ by using an explicit concentric LSH construction. The algorithm takes $O(dn^{1+\rho}\cdot\text{polylog}(n,d))$ pre-processing time/space, and $O(\sqrt{n}+n^{\rho}\cdot\text{polylog}(n,d))$
time/space per query.

Figure: Figure 4: An illustration of the concentric LSH construction of Mussmann et al. () In the $D_{i+1}$ band we find at least $\sqrt{n}$ points and in the $D_{i}$ band we find fewer than $\sqrt{n}$ points.
Refer to caption: extracted/5985875/concentric_lsh_illustration_2.png

###### Remark 27 (The role of $\rho$ ) .

The choice of $\rho<1$ allows us to compute $S_{i}$ in sublinear time for each $i\in[n]$. The value of $\rho$ is determined by the radii gaps in the concentric construction. Our algorithm for computing $S_{i}$ is sublinear in $n$ because $\rho<1$. Depending on the particular input dataset, we could have $\rho\leq 1/2$, which which case $f(n,\sqrt{n})=\widetilde{O}(\sqrt{n})$ in the context of Theorem 6.

###### Definition 25 (Approximate k 𝑘 k italic_k -MIPS) .

###### Theorem 26 ( Mussmann et al. ( 2017 ) ) .

###### Remark 27 (The role of ρ 𝜌 \rho italic_ρ ) .

### C.1 A complete algorithm based on solving k 𝑘 k italic_k -MIPS

We now have a complete algorithm to estimate self-attention with provable guarantees that is based on the solving $k$-MIPS problem for every query vector $q_{i}$. If we combine the boosted estimator approach of Theorem 4 with the Lazy Gumbel Sampling Technique of Algorithm 1 and the $k$-MIPS LSH technique of Theorem 26, we arrive at the following theorem. We give pseudocode for the resulting algorithm, as Algorithm 5:

###### Theorem 28 .

Let $\varepsilon>0$ and $\delta>0$ be small positive constants. There exists an algorithm that can estimate Self-Attention in the same way as Theorem 4 and fail with probability at most $\delta+1/n$. The algorithm’s time complexity is shown in the table below, where $\rho\in(0,1)$ is a fixed constant.

|  | Pre-Processing | Main Computation |
| --- | --- | --- |
| Complexity | $\widetilde{O}(dn^{1+\rho})$ | $\widetilde{O}\left(n^{1+\max\{1/2,\rho\}}\cdot d\cdot\varepsilon^{-2}\log(1/<br>\delta)\right)$ |

Figure: Algorithm 5 Approximating Self-Attention using concentric LSH $k$-MIPS solver

###### Proof.

Let $k=\sqrt{n}$. Suppose we construct the concentric LSH data structure according to Theorem 26. This takes $\widetilde{O}(dn^{1+\rho})$ time. Let us condition on the event that for all queries $q_{i}$ the data structure provides a correct approximate answer $S_{i}$ to the $k$-MIPS problem. This happens with probability at least $1-1/n$ by union bound over all $n$ queries. Now we use our sets $S_{i}$ in Algorithm 1 to sample from $D_{i}$. Since retrieving $S_{i}$ takes $f(n,k)=O(\sqrt{n}+n^{\rho}\cdot\text{polylog}(n))$ time and space, Theorem 7 dictates that sampling from $D_{i}$ also takes $\widetilde{O}(n^{\max\{1/2,\rho\}})$ time and space. Thus, in the context of Theorem 4 we have that $T=\widetilde{O}(n^{\max\{1/2,\rho\}})$. Substituting back gives us the desired runtime and failure probability guarantees.
∎

###### Theorem 28 .

###### Proof.

## Appendix D Derivation of the Self-Attention Gradients

Suppose we have a scalar function $\phi$ that represents the loss when training our neural network after computing the output $O$: $\ell=\phi(O)$. Suppose that we have calculated $\frac{\partial\phi}{\partial O_{ij}}$ for all $i\in[n],j\in[d]$ and stored it in an matrix $D^{O}\in\mathbb{R}^{n\times d}$. 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. ().

##### Calculating ∂ ϕ ∂ V i ⁢ j italic-ϕ subscript 𝑉 𝑖 𝑗 \frac{\partial\phi}{\partial V_{ij}} divide start_ARG ∂ italic_ϕ end_ARG start_ARG ∂ italic_V start_POSTSUBSCRIPT italic_i italic_j end_POSTSUBSCRIPT end_ARG

All these calculations just use the chain rule. One can simply draw a tree of dependencies and use it to perform the derivation. $\phi$ depends on $O_{ij}$ and $O_{ij}$ depends on all $V_{rj}$, so:

$$ $\frac{\partial\phi}{\partial V_{ij}}=\sum\limits_{r=1}^{n}\frac{\partial\phi}{ \partial O_{rj}}\cdot\frac{\partial O_{rj}}{\partial V_{ij}}=\sum\limits_{r=1} ^{n}D^{O}_{rj}\frac{\partial O_{rj}}{\partial V_{ij}}$ $$

Now, we calculate that:

$$ $\frac{\partial O_{rj}}{\partial V_{ij}}=\frac{\partial}{\partial V_{ij}}\sum \limits_{k=1}^{n}P_{rk}V_{kj}=P_{ri}$ $$

so that gives:

$$ $\displaystyle\frac{\partial\phi}{\partial V_{ij}}=\sum\limits_{r=1}^{n}D^{O}_{ rj}P_{ri}=\sum\limits_{r=1}^{n}P^{T}_{ir}D^{O}_{rj}$ (55) $$

Thus, we can write the result succinctly:

$$ $\displaystyle D^{V}=P^{T}\cdot D^{O}$ (56) $$

##### Calculating ∂ ϕ ∂ Q i ⁢ j italic-ϕ subscript 𝑄 𝑖 𝑗 \frac{\partial\phi}{\partial Q_{ij}} divide start_ARG ∂ italic_ϕ end_ARG start_ARG ∂ italic_Q start_POSTSUBSCRIPT italic_i italic_j end_POSTSUBSCRIPT end_ARG

To do this, we will first calculate $\frac{\partial\phi}{\partial P_{ij}}$ and $\frac{\partial\phi}{\partial S_{ij}}$, where $S=QK^{T}$.

- •
First, each $O_{ij}$ depends on all $P_{ik}$, so the chain rule gives:
$\frac{\partial\phi}{\partial P_{ij}}=\sum\limits_{k=1}^{d}\frac{\partial\phi}{
\partial O_{ik}}\cdot\frac{\partial O_{ik}}{\partial P_{ij}}=\sum\limits_{k=1}
^{d}D^{O}_{ik}\frac{\partial O_{ik}}{\partial P_{ij}}$
We can calculate that:
$\frac{\partial O_{ik}}{\partial P_{ij}}=V_{jk}$
and so:
$\displaystyle D^{P}_{ij}=\frac{\partial\phi}{\partial P_{ij}}=\sum\limits_{k=1
}^{d}D^{O}_{ik}V_{jk}=\langle D^{O}_{i,:},V_{j,:}\rangle$
(57)
for all $i\in[n],j\in[n]$.
- •
Now recall that $P_{ij}=\frac{\exp(S_{ij})}{L_{i}}$, so $P_{ij}$ depends on all $S_{ik}$ for $k=1,...,n$. Thus:
$\displaystyle\frac{\partial\phi}{\partial S_{ij}}=\sum\limits_{k=1}^{n}\frac{
\partial\phi}{\partial P_{ik}}\cdot\frac{\partial P_{ik}}{\partial S_{ij}}$
$\displaystyle=\sum\limits_{k=1}^{n}D^{P}_{ik}\cdot\frac{\partial P_{ik}}{
\partial S_{ij}}$
$\displaystyle=D^{P}_{ij}\cdot\frac{\partial P_{ij}}{\partial S_{ij}}+\sum
\limits_{k=1,k\neq j}^{n}D^{P}_{ik}\cdot\frac{\partial P_{ik}}{\partial S_{ij}}$
We now calculate seperately the two cases by using the quotient rule:
–
$k\neq j$:
$\displaystyle\frac{\partial P_{ik}}{\partial S_{ij}}=\frac{\partial}{\partial S
_{ij}}\frac{\exp(S_{ik})}{\sum\limits_{r=1}^{n}\exp(S_{ir})}$
$\displaystyle=-\exp(S_{ik})\cdot\frac{\exp(S_{ij})}{\left(\sum\limits_{r=1}^{n
}\exp(S_{ir})\right)^{2}}$
$\displaystyle=-P_{ik}P_{ij}$
–
$k=j$:
$\displaystyle\frac{\partial P_{ij}}{\partial S_{ij}}=\frac{\partial}{\partial S
_{ij}}\frac{\exp(S_{ij})}{\sum\limits_{r=1}^{n}\exp(S_{ir})}$
$\displaystyle=\frac{\exp(S_{ij})\sum\limits_{r=1}^{n}\exp(S_{ir})-\exp(S_{ij})
\exp(S_{ij})}{\left(\sum\limits_{r=1}^{n}\exp(S_{ir})\right)^{2}}$
$\displaystyle=P_{ij}-P_{ij}^{2}$
Now we can put it all together:
$\displaystyle\frac{\partial\phi}{\partial S_{ij}}$
$\displaystyle=D^{P}_{ij}\cdot\frac{\partial P_{ij}}{\partial S_{ij}}+\sum
\limits_{k=1,k\neq j}^{n}D^{P}_{ik}\cdot\frac{\partial P_{ik}}{\partial S_{ij}}$
$\displaystyle=D^{P}_{ij}\cdot(P_{ij}-P_{ij}^{2})-\sum\limits_{k=1,k\neq j}^{n}
D^{P}_{ik}\cdot P_{ik}P_{ij}$
$\displaystyle=D^{P}_{ij}\cdot P_{ij}-\sum\limits_{k=1}^{n}D^{P}_{ik}\cdot P_{
ik}P_{ij}$
$\displaystyle=P_{ij}\left(D^{P}_{ij}-\langle D^{P}_{i,:},P_{i,:}\rangle\right)$

Now finally, for $i\in[n],j\in[d]$, $Q_{ij}$ influences $S_{ik}$ for all $k\in[n]$, so:

$$ $\displaystyle\frac{\partial\phi}{\partial Q_{ij}}$ $\displaystyle=\sum\limits_{k=1}^{n}\frac{\partial\phi}{\partial S_{ik}}\frac{ \partial S_{ik}}{\partial Q_{ij}}$ (58) $\displaystyle=\sum\limits_{k=1}^{n}P_{ik}\left(D^{P}_{ik}-\langle D^{P}_{i,:}, P_{i,:}\rangle\right)K_{kj}$ (59) $$

##### Calculating ∂ ϕ ∂ K i ⁢ j italic-ϕ subscript 𝐾 𝑖 𝑗 \frac{\partial\phi}{\partial K_{ij}} divide start_ARG ∂ italic_ϕ end_ARG start_ARG ∂ italic_K start_POSTSUBSCRIPT italic_i italic_j end_POSTSUBSCRIPT end_ARG

We know that $K_{ij}$ influences $S_{ki}$ for $k\in[n]$, so:

$$ $\displaystyle\frac{\partial\phi}{\partial K_{ij}}$ $\displaystyle=\sum\limits_{k=1}^{n}\frac{\partial\phi}{\partial S_{ki}}\frac{ \partial S_{ki}}{\partial K_{ij}}$ (60) $\displaystyle=\sum\limits_{k=1}^{n}P_{ki}\left(D^{P}_{ki}-\langle D^{P}_{k,:}, P_{k,:}\rangle\right)Q_{kj}$ (61) $$

## Appendix E Estimating D Q superscript 𝐷 𝑄 D^{Q} italic_D start_POSTSUPERSCRIPT italic_Q end_POSTSUPERSCRIPT

In this section we give an efficient algorithm for estimating $D^{Q}$. This algorithm is based on our $k$NN-Attention framework. Recall that we found that:

$$ $D^{Q}_{ij}=\sum\limits_{k=1}^{n}P_{ik}\left(D^{P}_{ik}-\langle D^{P}_{i,:},P_{ i,:}\rangle\right)K_{kj}$ $$

We can write this expression as an expectation with respect to the distribution $D_{i}$:

$$ $\displaystyle\frac{\partial\phi}{\partial Q_{ij}}$ $\displaystyle=\mathbb{E}_{k\sim D_{i}}\left[D^{P}_{ik}K_{kj}\right]-\mathbb{E} _{k\sim D_{i}}\left[K_{kj}\cdot\mathbb{E}_{s\sim D_{i}}[D^{P}_{is}]\right]$ (62) $\displaystyle=\underbrace{\mathbb{E}_{k\sim D_{i}}\left[D^{P}_{ik}K_{kj}\right ]}_{E_{1}}-\underbrace{\mathbb{E}_{k\sim D_{i}}\left[K_{kj}\right]}_{E_{2}} \cdot\underbrace{\mathbb{E}_{s\sim D_{i}}[D^{P}_{is}]}_{E_{3}}$ (63) $$

This allows us to use any of our softmax expectation estimators. We choose the Median-Of-Means estimator for the purposes of a clean analysis. We just have to do it three times and ensure that the terms we take expectations over are efficiently computable. Indeed, because

$$ $\displaystyle D^{P}_{ik}=\langle D^{O}_{i,:},V_{k,:}\rangle,$ (64) $$

we can compute all three of those expectations in sublinear time! Let $\widehat{E_{1}},\widehat{E_{2}},\widehat{E_{3}}$ be the estimates we produce. Then, almost identically to the error analysis we did for the forward pass, we get an $(\varepsilon,\delta)$-additive estimate for $E_{i}$, where $i\in\{1,2,3\}$(^6^66In Theorem 4 we used a multiplicative approximation. To get the additive approximation guarantee we need $O(\varepsilon^{-2}\log(1/\delta)\cdot\text{Var}[\widehat{O}_{ij}])$ samples, where $\text{Var}[\widehat{O}_{ij}]\leq B^{2}=O(\text{polylog}(n))$.)

$$ $\displaystyle\Pr\left[|\widehat{E_{1}}-E_{1}|\geq\varepsilon\right]\leq\frac{ \delta}{3}$ (65) $\displaystyle\Pr\left[|\widehat{E_{2}}-E_{2}|\geq\varepsilon\right]\leq\frac{ \delta}{3}$ (66) $\displaystyle\Pr\left[|\widehat{E_{3}}-E_{3}|\geq\varepsilon\right]\leq\frac{ \delta}{3}$ (67) $$

And so, putting these three together and using the union bound we get that with probability at least $1-\delta$ it holds that:

$$ $\displaystyle\left|\widehat{E_{1}}-\widehat{E_{2}}\cdot\widehat{E_{3}}-E_{1}+E _{2}\cdot E_{3}\right|$ $\displaystyle\leq\left|\widehat{E_{1}}-E_{1}\right|+\left|\widehat{E_{2}}\cdot \widehat{E_{3}}-E_{2}\cdot E_{3}\right|$ (68) $\displaystyle=\left|\widehat{E_{1}}-E_{1}\right|+\left|\widehat{E_{2}}\cdot \widehat{E_{3}}-\widehat{E_{2}}\cdot E_{3}+\widehat{E_{2}}\cdot E_{3}-E_{2} \cdot E_{3}\right|$ (69) $\displaystyle\leq\left|\widehat{E_{1}}-E_{1}\right|+\widehat{E_{2}}\left| \widehat{E_{3}}-E_{3}\right|+E_{3}\left|\widehat{E_{2}}-E_{2}\right|$ (70) $\displaystyle\leq\varepsilon+\varepsilon\cdot\widehat{E_{2}}+\varepsilon E_{3}$ (71) $\displaystyle\leq\varepsilon+\varepsilon(E_{2}+\varepsilon)+\varepsilon E_{3}$ (72) $\displaystyle=\varepsilon+\varepsilon^{2}+\varepsilon(E_{2}+E_{3})$ (73) $$

In order to bound the variance of our estimators, we need to assume some bounds on the inputs, analogously to $||V||_{\infty}\leq B=O(\lg n)$ from Theorem 4. First, we assume that $||K||_{\infty}\leq B_{K}=O(\text{polylog}(n))$. Second, we have that $||D^{P}||_{\infty}\leq dB\cdot||D^{O}||_{\infty}=O(\text{polylog}(n))$ if $d=O(\log n)$. This also gives that $||D^{P}\circ K||_{\infty}\leq B_{K}\cdot||D^{P}||_{\infty}=O(\text{polylog}(n))$.

These assumptions are reasonable within the context of the hardness results proved for the attention mechanism and the computation of its gradients(^7^77Another motivation for assumming an upper bound on the norm of $D^{O}$ is to avoid the phenomenon of exploding gradients in training neural networks.) (Alman and Song, , ). Given these assumptions, we can also bound the error more compactly. Starting from Equation 73, we get:
$e_{Q}\leq O(\varepsilon)+\varepsilon(B_{K}+dB\cdot||D^{O}||_{\infty})=O(
\varepsilon\cdot\text{polylog}(n))$. As a result, we arrive at the following theorem:

###### Theorem 29 .

Assume that $||K||_{\infty}=O(\text{polylog}(n)),d=O(\log n)$ and $||D^{O}||_{\infty}=O(\text{polylog}(n))$. There exists a sub-quadratic algorithm that takes as input $Q,K,V,D^{O}\in\mathbb{R}^{n\times d}$ and outputs a matrix $\widehat{D}^{Q}\in\mathbb{R}^{n\times d}$ such that:

$$ $\displaystyle\left|\left|\widehat{D}^{Q}-D^{Q}\right|\right|_{\infty}\leq O( \varepsilon\cdot\text{polylog}(n))$ (74) $$

This algorithm is shown as Algorithm 6.

Figure: Algorithm 6 Estimating $D^{Q}$

###### Theorem 29 .

## Appendix F Estimating D K superscript 𝐷 𝐾 D^{K} italic_D start_POSTSUPERSCRIPT italic_K end_POSTSUPERSCRIPT

Finally, we turn to estimating $D^{K}$. Our earlier calculations show that

$$ $\frac{\partial\phi}{\partial K_{ij}}=\sum\limits_{k=1}^{n}P_{ki}\left(D^{P}_{ ki}-\langle D^{P}_{k,:},P_{k,:}\rangle\right)Q_{kj}$ $$

We can break up this sum into two terms:

$$ $\displaystyle\frac{\partial\phi}{\partial K_{ij}}=\underbrace{\sum\limits_{k=1 }^{n}P_{ki}D^{P}_{ki}Q_{kj}}_{A_{ij}}-\underbrace{\sum\limits_{k=1}^{n}P_{ki} \langle D^{P}_{k,:},P_{k,:}\rangle\cdot Q_{kj}}_{B_{ij}}$ (75) $$

We will estimate both terms separately:

### F.1 Estimating A i ⁢ j subscript 𝐴 𝑖 𝑗 A_{ij} italic_A start_POSTSUBSCRIPT italic_i italic_j end_POSTSUBSCRIPT

For $i\in[n]$ and $j\in[d]$, we have:

$$ $\displaystyle A_{ij}=\sum\limits_{k=1}^{n}P_{ki}D^{P}_{ki}Q_{kj}$ $\displaystyle=\sum\limits_{k=1}^{n}P_{ki}Q_{kj}\cdot\langle D^{O}_{k,:},V_{i,:}\rangle$ (76) $\displaystyle=\sum\limits_{k=1}^{n}P^{T}_{ik}\cdot Y_{kj}^{(i)}$ (77) $$

where $Y_{kj}^{(i)}:=Q_{kj}\cdot\langle D^{O}_{k,:},V_{i,:}\rangle$. So we can write:

$$ $\displaystyle A_{ij}=\underbrace{(P^{T})_{i,:}}_{1\times n}\cdot\underbrace{Y^ {(i)}_{:,j}}_{n\times 1}$ (78) $$

We will use our familiar Markov Chain estimation method from Algorithm 3 to calculate this quantity. However, in this case we only care about estimating the $i$-th entry in the vector $(P^{T})\cdot Y^{(i)}_{:,j}$, which we can do by performing $O(\lg n\cdot\varepsilon^{-2})$ simulations. Ultimately, by following the same analysis as in Algorithm 3, we are able to estimate $A_{ij}$ with probability at least $1-\frac{1}{n}$ and error:

$$ $\displaystyle\left|\widehat{A}_{ij}-A_{ij}\right|$ $\displaystyle\leq\varepsilon\langle Y^{(i)}_{:,j},1^{n}\rangle+2\varepsilon nM _{j}^{(i)}$ (79) $\displaystyle=\varepsilon\sum\limits_{k=1}^{n}Q_{kj}\cdot\langle D^{O}_{k,:},V _{i,:}\rangle+2\varepsilon nM_{j}^{(i)}$ (80) $$

where

$$ $M_{j}^{(i)}=-\min\limits_{\begin{subarray}{c}k\in[n]\\ Y^{(i)}_{kj}\leq 0\end{subarray}}Y^{(i)}_{kj}$ $$

###### Remark 30 .

Because we would need to calculate all $n^{2}$ values of $M_{j}^{(i)}$, we will instead use a single upper bound $M\geq M_{j}^{(i)}$ for all $(i,j)\in[n]\times[d]$ for this algorithm. We assume that we know a large enough $M$ in advance and that $M=O(\text{polylog}(n))$.

In the next paragraphs, we will tackle some implementation issues that arise in this approach. We did not see these issues when estimating $D^{V}$, and because they make the algorithm a lot more complicated, we left them for last.

###### Remark 30 .

##### Pre-calculating the normalizing factors

We need to pre-calculate the normalizing sums $N_{j}^{(i)}=\langle Y^{(i)}_{:,j},1^{n}\rangle+nM$ for all $(i,j)\in[n]\times[d]$. Naively, it takes $O(n^{2}d)$ time to calculate all those sums. However, with some preprocessing we can take the time down to $O(nd^{2})$. First, observe that we have:

$$ $\displaystyle N_{j}^{(i)}=nM+\langle Y^{(i)}_{:,j},1^{n}\rangle$ $\displaystyle=nM+\sum\limits_{k=1}^{n}Q_{kj}\cdot\langle D^{O}_{k,:},V_{i,:} \rangle=nM+\langle V_{i,:},\sum\limits_{k=1}^{n}Q_{kj}\cdot D^{O}_{k,:}\rangle$ (81) $$

We can thus first pre-compute the $d$ vectors $\overrightarrow{E_{j}}=\sum_{k=1}^{n}Q_{kj}\cdot D^{O}_{k,:}\in\mathbb{R}^{d}$ for each $j\in[d]$ in $O(nd^{2})$ time. Then, for each $i\in[n]$ and $j\in[d]$, we can produce $N_{j}^{(i)}$ in $O(d)$ time by using Equation 81, bringing the total time complexity to $O(nd^{2})$.

##### Sampling according to Y : , j ( i ) + M ⋅ 1 n subscript superscript 𝑌 𝑖 : 𝑗 ⋅ 𝑀 superscript 1 𝑛 Y^{(i)}_{:,j}+M\cdot 1^{n} italic_Y start_POSTSUPERSCRIPT ( italic_i ) end_POSTSUPERSCRIPT start_POSTSUBSCRIPT : , italic_j end_POSTSUBSCRIPT + italic_M ⋅ 1 start_POSTSUPERSCRIPT italic_n end_POSTSUPERSCRIPT efficiently

Unfortunately, because we are now estimating $A_{ij}$ individually for all $(i,j)\in[n]\times[d]$, we cannot spend $O(n)$ time to generate each sample. We need to generate samples in sublinear time with some pre-processing. This seems intuitively difficult at first because we have $O(nd)$ distributions over $[n]$ and each distribution requires $\Omega(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_{1},...,p_{n}]$:

- 1.
Compute the cumulative sums $s_{i}=\sum_{k=1}^{i}p_{i}$. We know that $s_{1}=p_{1}$ and $s_{n}=1$.
- 2.
Pick some $x\sim\text{Unif}(0,1)$ uniformly at random from $(0,1)$.
- 3.
Find the interval $[p_{i},p_{i+1}]$ for $i\in[1,n-1]$ in which $x$ falls in. That is, find the smallest $i$ for which $x\leq s_{i}$. We can do this in $O(\lg n)$ time using binary search.
- 4.
Output $i$.

Figure: Figure 5: An illustration of the CDF sampling method: We form the CDF and then sample an index by choosing $x\in(0,1)$ and using binary search to find the corresponding bucket.
Refer to caption: extracted/5985875/cdf_sampling_new.png

It is easy to see that this method outputs a value $i$ with probability $p_{i}$. If we applied this method naively we would still take $O(n^{2}d)$ time because we’d have to calculate all the cumulative sums. However, the inner product structure again comes to our rescue:

$$ $\displaystyle Y^{(i)}_{kj}=\langle V_{i,:},Q_{kj}\cdot D^{O}_{k,:}\rangle$ (82) $$

So, we can create $d$ cumulative sum tables $\Sigma_{j}$ for $j\in[d]$, each of which stores $n$ cumulative-sum $\mathbb{R}^{d}$ vectors as follows:

$$ $\displaystyle(\Sigma_{j})_{\ell}=\sum\limits_{s=1}^{\ell}Q_{sj}\cdot D^{O}_{s, :}\in\mathbb{R}^{d},\,\forall\ell\in[n]$ (83) $$

This requires $O(nd^{2})$ time and space to construct. Now, in order to sample with probability proportional to $Y^{(i)}_{kj}+M$ given that we know $N_{j}^{(i)}$, we sample $x_{ij}\sim\text{Unif}(0,1)$ and perform binary search to find the interval $x_{ij}$ belongs to. At that point, we can calculate the $O(\log n)$ necessary cumulative sums in $O(d)$ time each by using our pre-processing:

$$ $\displaystyle\sum\limits_{s=1}^{\ell}\left(Y^{(i)}_{kj}+M\right)=kM+\langle V_ {i,:},(\Sigma_{j})_{\ell}\rangle$ (84) $$

This allows us to sample in $O(d\lg n)$ time after a $O(nd^{2})$ pre-processing. Our algorithm in total is included as part of Algorithm 7.

##### Sampling with respect to D i subscript 𝐷 𝑖 D_{i} italic_D start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT

Again, we 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-processing time: the Lazy-Gumbel Sampling method. We will omit the pre-processing details in the algorithm pseudocode.

Figure: Algorithm 7 Estimating $D^{K}$ – Part 1: Computing $A$

### F.2 Estimating B i ⁢ j subscript 𝐵 𝑖 𝑗 B_{ij} italic_B start_POSTSUBSCRIPT italic_i italic_j end_POSTSUBSCRIPT

For $(i,j)\in[n]\times[d]$, we first have:

$$ $\displaystyle B_{ij}$ $\displaystyle=\sum\limits_{k=1}^{n}P_{ki}\cdot\langle D^{P}_{k,:},P_{k,:} \rangle\cdot Q_{kj}$ (85) $\displaystyle=\sum\limits_{k=1}^{n}P_{ki}X_{kj}$ (86) $$

where $X_{kj}=\langle D^{P}_{k,:},P_{k,:}\rangle\cdot Q_{kj}$. Notice that $X_{kj}$ takes $O(nd)$ time to naively compute, so we will first approximate it with $\widehat{X}_{kj}$. Observe that:

$$ $\displaystyle X_{kj}$ $\displaystyle=Q_{kj}\cdot\langle D^{P}_{k,:},P_{k,:}\rangle$ (87) $\displaystyle=Q_{kj}\cdot\sum\limits_{s=1}^{n}D^{P}_{ks}\cdot P_{ks}$ (88) $\displaystyle=Q_{kj}\cdot\mathbb{E}_{s\sim D_{k}}[D^{P}_{ks}]$ (89) $\displaystyle=\mathbb{E}_{s\sim D_{k}}[Q_{kj}\cdot D^{P}_{ks}]$ (90) $\displaystyle\approx\widehat{X}_{kj}$ (91) $$

Let us approximate $\mathbb{E}_{s\sim D_{k}}[Q_{kj}\cdot(D^{p})_{ks}]\approx\widehat{X}_{kj}$ using the Lazy Gumbel Sampling and Median-Of-Means method. This allows us to get for all $(k,j)\in[n]\times[d]$ with probability at least $1-\delta$ that:

$$ $\displaystyle\left|\widehat{X}_{kj}-X_{kj}\right|\leq\varepsilon$ (92) $$

###### Remark 31 .

To have an $o(n)$ bound for the variance, we have to assume (again) that $||X||_{\infty}=O(\text{polylog}(n))$. This follows from the assumption that $||Q||_{\infty}=O(\text{polylog}(n))$ and $||D^{P}||_{\infty}=O(\text{polylog}(n))$. The latter follows from $||D^{O}||_{\infty}=O(\text{polylog}(n))$. So the assumptions here are the same as in Theorem 29.

Now we can define:

$$ $\displaystyle\widehat{B}_{ij}$ $\displaystyle=\sum\limits_{k=1}^{n}P_{ki}\widehat{X}_{kj}$ (93) $$

We can bound the error of this approximation using the triangle inequality:

$$ $\displaystyle\left|B_{ij}-\widehat{B}_{ij}\right|$ $\displaystyle=\left|\sum\limits_{k=1}^{n}P_{ki}(\widehat{X}_{kj}-X_{kj})\right|$ (94) $\displaystyle\leq\sum\limits_{k=1}^{n}P_{ki}\left|\widehat{X}_{kj}-X_{kj}\right|$ (95) $\displaystyle\leq\varepsilon\sum\limits_{k=1}^{n}P_{ki}$ (96) $\displaystyle=\varepsilon\langle P_{:,i},1^{n}\rangle$ (97) $$

Now the problem is calculating $\widehat{B}$. Note that we can write:

$$ $\displaystyle\widehat{B}=P^{T}\cdot\widehat{X}$ (98) $$

Finally, this takes us back to the calculation of $D^{V}$. We can use the exact same Markov Chain method and get a final approximation $\widetilde{B}$ so that with probability at least $1-\frac{1}{n}$ it holds that:

$$ $\displaystyle\left|\widetilde{B}_{ij}-\widehat{B}_{ij}\right|\leq\varepsilon \langle\widehat{X}_{:,j},1^{n}\rangle+2\varepsilon nM^{(X)}_{j}$ (99) $$

where

$$ $M^{(X)}_{j}:=-\min_{\begin{subarray}{c}k\in[n]\\ \widehat{X}_{kj}\leq 0\end{subarray}}\widehat{X}_{kj}$ $$

Then the overall error can be bounded as follows:

$$ $\displaystyle\left|\widetilde{B}_{ij}-{B_{ij}}\right|$ $\displaystyle\leq\left|\widetilde{B}_{ij}-\widehat{B}_{ij}\right|+\left| \widehat{B}_{ij}-{B_{ij}}\right|$ (100) $\displaystyle\leq\varepsilon\langle P_{:,i},1^{n}\rangle+\varepsilon\langle \widehat{X}_{:,j},1^{n}\rangle+2\varepsilon nM^{(X)}_{j}$ (101) $\displaystyle\leq\varepsilon\langle P_{:,i},1^{n}\rangle+\varepsilon\langle{X_ {:,j}},1^{n}\rangle+\varepsilon^{2}n+2\varepsilon nM^{(X)}_{j}$ (102) $\displaystyle=\varepsilon\langle P_{:,i}+X_{:,j},1^{n}\rangle+\varepsilon^{2}n +2\varepsilon nM_{j}^{(X)}$ (103) $$

To wrap up our implementation details, we can calculate the required normalization sums as follows:

$$ $\displaystyle\langle\widehat{X}_{:,j},1^{n}\rangle=\sum\limits_{k=1}^{n} \widehat{X}_{kj}=\sum\limits_{k=1}^{n}Q_{kj}\widehat{D}_{k}$ (104) $$

We can do this in $\approx\widetilde{O}(dn^{3/2})$ time if we precompute in advance

$$ $\displaystyle\widehat{D}_{k}:=\langle D^{P}_{k,:},P_{k,:}\rangle$ (105) $$

using Lazy Gumbel Sampling for all $k\in[n]$. Further, each element $\widehat{X}_{ij}$ can be computed in $\approx\widetilde{O}(\sqrt{n})$ time as well in a similar fashion. Finally, $M^{(X)}_{j}$ can also be calculated in such time. Our algorithm is given below as Algorithm 8. By combining algorithms 7 and 8 we arrive at the following theorem for Algorithm 9:

Figure: Algorithm 8 Estimating $D^{K}$ – Part 2: Computing $B$

###### Theorem 32 .

There exists an algorithm that approximates $D^{K}$ on inputs $Q,K,V,D^{O}$ under our standard assumptions such that the estimate $\widehat{D}^{K}$ satisfies:

$$ $\displaystyle\left|\left|\widehat{D}^{K}_{:,j}-D^{K}_{:,j}\right|\right|_{ \infty}\leq\varepsilon\langle P_{:,i}+X_{:,j},1^{n}\rangle$ $\displaystyle+\varepsilon^{2}n+2\varepsilon nM_{j}^{(X)}$ $\displaystyle+\varepsilon\sum\limits_{k=1}^{n}Q_{kj}\cdot\langle(D^{o})_{k,:}, V_{i,:}\rangle+2\varepsilon nM$ $$

where:

$$ $\displaystyle M\geq M_{j}^{(i)}:=-\min\limits_{\begin{subarray}{c}k\in[n]\\ Y^{(i)}_{kj}\leq 0\end{subarray}}Y^{(i)}_{kj}\,\text{ and }M^{(X)}_{j}:=-\min_ {\begin{subarray}{c}k\in[n]\\ \widehat{X_{kj}}\leq 0\end{subarray}}\widehat{X_{kj}}$ (106) $$

under our previous definitions for all $j\in[d]$. The algorithm runs in sub-quadratic time and space and succeeds with probability $\geq 1-\delta$.

###### Proof.

We know that $D^{K}_{ij}=A_{ij}-B_{ij}$. We have that $\widehat{D}^{K}_{ij}=\widehat{A}_{ij}-\widetilde{B}_{ij}$ and that:

$$ $\displaystyle|\widehat{A}_{ij}-A_{ij}|$ $\displaystyle\leq\varepsilon\sum\limits_{k=1}^{n}Q_{kj}\cdot\langle(D^{o})_{k, :},V_{i,:}\rangle+2\varepsilon nM$ (107) $\displaystyle|\widetilde{B}_{ij}-B_{ij}|$ $\displaystyle\leq\varepsilon\langle P_{:,i}+X_{:,j},1^{n}\rangle+\varepsilon^{ 2}n+2\varepsilon nM_{j}^{(X)}$ (108) $$

Thus by the triangle inequality we get the desired error guarantee.
∎

Figure: Algorithm 9 Estimating $D^{K}$: Putting it all together

###### Remark 31 .

###### Theorem 32 .

###### Proof.

## Appendix G Vectorized Implementation of the Forward Pass

We present the vectorized implementation of $k$NN Attention that we used in our experiments. This is based on Theorem 8.

Figure: Listing 1: kNN Attention - Vectorized

## Appendix H Samples from Fine-Tuned GPT2-XL

Below we give some samples from our fine-tuned GPT2-XL using fast approximate kNN attention:

- •
Prompt: My favorite thing in the world is: Just Cause 3.
And I couldn’t stop playing it.
It’s a 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 a bunch of different characters. It’s epic, and it’s open world, and you can play as a bounty hunter, or you can play as a doctor, or a cowboy, or you can play as a soldier.
You have a buddy system, and you can go to another character’s story mission if you want.
This is a great game, and I’m excited to see what the developers do with it.
- •
Prompt: I like being in nature because it’s so natural, and because it’s always changing. Nature is a friend and a companion, and I’m so thankful every time I’m able to go out into the outdoors.
The second thing that I 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. I don’t have to think about what I’m wearing, or what I’m doing, and I can just be in nature in my element. I also feel great being out in nature because 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.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Please see weaknesses section.
soundness: 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.
The 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.

The 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.
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.
weaknesses: I don’t have many weaknesses to point out, but I do have some questions/points I’d like to raise/discuss:


—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?


—(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?

—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.

—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.


—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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: Please see the questions in the previous part. 

In addition: 

- 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?
soundness: 3
strengths: - The problem studied in the paper is well defined in the abstract. 
- Understanding the approximation abilities of sparse attention models, of which kNN attention are a special case, is a significant research question. 
- Similarly, proposing alternatives to computationally costly gradient computations in Transformers could have a strong impact in the community.
- The code is provided which enables to reproduce the experimental results.
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.

More precisely, the submission is organized as follows:

- In part 1, the theoretical results are outlined in brief.
- 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$.
  - 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.
  - Section 2.3 introduces an alternative method for computing kNN attention outputs, designed to be more compatible with modern hardware.
- 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.
- 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.
weaknesses: I do not believe this paper to be ready to be published at ICLR for the following reasons:

Overall clarity:

The 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. 

I am adding specific suggestions / questions for improvement below:

- It should be clearly mentioned that Part 1 sketches the results that are later detailed in Parts 2 and 3. 
- 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. 
- 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. 
- 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.
- 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.
- 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?
- In line 129, what does $T$ represent?
- In line 120, the term $k_k$ in equation (3) is confusing.
- 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.

Theoretical Contributions:

- In general, I find the results challenging to interpret. How do these results compare to previous work? What are typical values?
- 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.
- Line 181: this is a proof sketch, not a full proof. I could not find the complete proof in the appendix.
- The proof of Theorem 4 is also only a sketch. For instance, see the last two lines.

Experiments:

- There should be experiments specifically validating the theoretical bounds. As is, the experiments are rather qualitative.
- I had to make some manual adjustments to get the provided code running. Including a `setup.py` would be helpful.
- 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?
- 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?

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 1
questions: 1) Could you please clarify the theoretical novelty? For example:
- Is Theorem 4 a direct consequence of Median-Of-Means Boosting?
- Does Theorem 8 (main result) rely heavily on Mussmann et al. (2017)?

2) Where is the proof of Theorem 2? Is it a direct consequence of another specific theorem?  
3) In Theorem 4, what does $O(T)$ refer to? A discussion on the different parameters in the theorem is needed.  
4) In Theorem 28, how large can $\rho$ be? Does it approach 1?  
5) Where is the proof of Theorem 29? Is it a direct consequence of (Alman & Song, 2024a; b)?  

6) For notation, please clarify Gumbel $(a, b)$ and Bin $(a, b)$.  
7) 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?  
8) In Lines 184–189, what is the Moment Generating Function of the Gumbel distribution used in the proof of Lemma 6?  
9) Where is the proof of Theorem 7? What do you mean by "it is easy to derive from Algorithm 1"?  
10) In Theorem 7, how strong is the assumption $k = \sqrt{n}$?  

11) 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?  

12) 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:
$
k \geq \sqrt{\frac{8n^2 \varepsilon^{-2} \log(4/\delta)}{\ell}}.
$
However, this assumption seems unrealistic as $k$ scales with $n$.

13) 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$?  

14) In Line 246, which algorithm are you referring to?  
15) It is unclear how the total runtime in Line 224 was obtained. Is this result similar to Theorem 4 under comparable assumptions?  

16) What is Figure 3(a)? I couldn’t find it.  
17) What is the error definition in the discussion on "Efficiency of kNN Attention" (Lines 421–425)?  
18) 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$.  

19) 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.  

20) 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?
soundness: 2
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. 

- I think the paper is well-structured. 

- Additionally, the author provides empirical experiments that demonstrate scalability and effectiveness in reducing computational and memory overhead.
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. 

It 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.
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. 

- 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. 

- 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.

- 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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - In Sections 2.2 and 2.3, two approximation methods are proposed. Can their performance be compared directly?
- 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?
- 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?
- In Section 4.2, would it be possible to include a comparison of computation speeds?
soundness: 3
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.
- Clear pseudocode is provided for each algorithm.
- The authors have made their experimental code publicly available.
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. 
Specifically, 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. 
Additionally, 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.
weaknesses: - **Lack of originality**
  + 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].
  + 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.
- **Presentation issues**
  + The reference to [2] is listed as an arXiv preprint, but it was accepted at Uncertainty in Artificial Intelligence 2017.
  + In equation (3), the dot is used inconsistently between $q_i^\top \cdot k_k$ and $q_i^\top k_s$.
  + The time complexity in Theorem 2 should be clarified as time complexity in expectation.
  + The citation format on line 97 of page 2 should be consistent with others.
  + "BIN" in Algorithm 1 is not defined.
  + The term "kNN index" lacks sufficient explanation.
  + The sentence "if we assume that..." on line 222 of page 5 was unclear to me; further clarification would be appreciated.
  + It would be helpful to clearly indicate which parts of Algorithm 2 constitute pre-processing.
  + 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.
  + 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.
  + Numbered and unnumbered equations are mixed inconsistently.

[1] Kratsios, Universal Regular Conditional Distributions. 2021.  
[2] Mussmann et al., Fast Amortized Inference and Learning in Log-linear Models with Randomly Perturbed Nearest Neighbor Search. 2017.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
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. 

2. 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 ?
soundness: 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.
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.
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}.

Meta-review: The paper addresses an important and relevant topic of quadratic complexity of attention in Transformers, and provides a theoretical analysis behind the empirical performance of KNN attention. The paper is well-written and the ideas are clearly organized. The authors effectively combine several well-known ideas, such as decomposing the attention matrix as expectations over softmax, using median-of-means boosting to achieve high-probability approximation bounds and runtime complexity bounds, and employing the Gumbel-Max trick and the concentration properties of the Gumbel distribution. The paper proposes new sub-quadratic approximation algorithms for attention gradients, inspired by approximation and randomized algorithm techniques.

Decision: Accept (Poster)