HSR-Enhanced Sparse Attention Acceleration

Introduction

Large Language Models (LLMs) have showcased remarkable capabilities across various applications, including context-aware question answering, content generation, summarization, and dialogue systems, among others . Long-context tasks of LLMs have gained more and more attention. Several LLMs extend their context length to $128$K tokens, such as Yarn , GPT-4 , Claude 3.5 , Llama 3.1 , Phi-3.5 , Mistral Nemo , etc. A bottleneck for long-context tasks is the computational cost of the attention mechanism in LLMs. The key to LLM success is the transformer architecture , wildly used in various practical scenarios , whose critical component is the attention mechanism. Let $n$ be the data length, $m$ be the length of query tokens, and $d$ be the feature dimensionAs $d$ is always fixed in practice, there is no need to scale up $d$ in analysis. Thus, in this work, we always assume $d$ is a small constant.. The conventional attention uses Softmax activation and is defined as follows:

In practical LLM applications, there are two scenarios for attention computation depending on the context length $n$ and query length $m$. The first case, $m=\Theta(1)$, represents the iterative text generation based on the pre-computed Key Value Cache (KV), which stores the intermediate attention key and value matrices. The second case, $m=\Theta(n)$, represents the full self-attention computation before text generation or the cross-attention computation. However, in both cases, when the context window $n$ becomes larger, the running time will increase correspondingly, i.e., it will be linear and quadratic in $n$ for $m=\Theta(1)$ and $m=\Theta(n)$, respectively. Thus, reducing the running time of attention computations with long context input becomes essential to minimize response latency and increase throughput for LLM API calls.

In this work, we introduce novel methods to reduce the running time complexity for both cases, i.e., $m=\Theta(1)$ and $m=\Theta(n)$. Our approach is inspired by the inherent sparsity found within attention mechanisms. Numerous prior studies have highlighted the significant sparsity in the attention matrix . This manifestation of sparsity in Softmax attention is that a large number of attention scores, i.e., $QK^{\top}$, concentrate on a small number of entries, which is known as “massive activation”. Due to this nature, Softmax attention can be accelerated by only calculating the entries that contain large attention scores, introducing negligible approximation errors .

To expedite the computation, the critical task is to identify the large/non-zero entries for Softmax/ReLU attention, respectively. To do so, we utilize the half-space reporting (HSR) data structure, which is introduced in to address the half-space range reporting problem. This is a fundamental problem in computational geometry and can be formally defined as follows:

To the best of our knowledge, this is the first work incorporating the HSR data structure with attention computation, to reduce the running time complexity with the help of the sparsity within the attention mechanisms.

Theoretically, we provide rigorous proofs for reducing the computational time (1) for ReLU attention generation from $O(mn)$ to $O(mn^{4/5})$ (Algorithm 2 and Theorem 4.1); (2) for full ReLU attention computation from $O(mn)$ to $O(mn^{1-1/\lfloor d/2\rfloor}+mn^{4/5})$ (Algorithm 3 and Theorem 5.1), without incurring any approximation error in both cases.

We achieve the same running time speed up for the conventional Softmax attention, and we give rigorous theoretical proofs to ensure that the resulting approximation error remains negligible (Theorem 4.2, 5.2 and Theorem 4.3).

We conduct empirical experiments on prominent LLMs to verify the approximation error associated with Softmax attention utilizing “massive activated” entries only. The results show that the error using a few top entries is already insignificant, consistent with our theoretical analysis.

Roadmap. Section 2 presents related work. In Section 3, we introduce essential concepts and key definitions used this paper. In Section 4, we present our main results, i.e., guarantees on run time reduction and approximation error. In Section 5, we introduce the extension of our method on full attention computation. In Section 6, we provide a brief summary of the techniques used in our proof. In Section 7, we provide our empirical results of evaluating three mainstream LLMs with Softmax attention with top-$r$ indices on different $r$. In Section 8, we discuss the potential of extending our method to other activation functions. In Section 9, we concludes our algorithm and contributions.

Related Work

Attention acceleration for long context input. Long context window is essential for transformer based LLMs in many downstream tasks. However, due to the quadratic time complexity associated with self-attention mechanisms, transformers are usually hard to inference efficiently. Numerous methods have been proposed to enhance the inference efficiency. One approach involves using alternative architectures as proxies for attention to support faster inference, such as Mamba , PolySketchFormer , Hopfield Models and Linearizing Transformers . However, the broad applicability of these methods across different applications and modalities remains to be fully validated. Another line of research focuses on approximating attention matrix computation . Nevertheless, these methods often rely on assumptions that may not be practical. For instance, some approaches use polynomial methods to approximate the exponential function, which requires all entries to be bounded by a small constant. However, our HSR-enhanced attention framework is designed based on practical observation and validated by empirical support. These advancements in attention mechanisms and transformer efficiency not only improve general model performance but also play a crucial role in enhancing in-context learning capabilities, where models leverage information from the immediate context to perform tasks without fine-tuning. We refer the readers to some other related works .

ReLU attention. ReLU attention is an innovative mechanism that employs the ReLU activation function in place of the traditional Softmax function for attention computation. Previous studies have highlighted the promise potential of ReLU attention in various domains. From empirical side, has demonstrated that incorporating ReLU as the activation function in vision transformers enhances performance on downstream tasks. has shown that transformers equipped with ReLU attention outperform those with Softmax attention, particularly when dealing with large key-value memory in machine translation tasks. From theoretical side, the scale-invariant property of ReLU attention facilitates the scalability of transformer networks. Furthermore, have shown that the inherent properties of ReLU attention contribute positively to the learning process of transformer models. Another key advantage of ReLU attention is that the ReLU function effectively sets all negative values to zero, allowing us to bypass these non-contributory elements during attention computation, thereby reducing the running time of attention computation. Importantly, omitting these zero and negative entries does not introduce any error into the final output of the ReLU attention mechanism.

Half-space reporting (HSR) data structure. The Half-Space Reporting (HSR) data structure, initially proposed by , was developed to address the half-space range reporting problem. The expedited range query capability inherent to HSR has been demonstrated to significantly enhance computational efficiency across a variety of tasks, as evidenced by numerous previous works in the literature. Studies such as and have applied HSR to facilitate solving general linear programming (LP) problems. Another line of research has highlighted HSR’s potential in expediting the training process of contemporary neural networks . There is also a collection of research that concentrates on leveraging HSR for the advancement of solutions to geometric and graphical challenges .

Preliminary

In Section 3.1, we introduce notations used in the paper. In Section 3.2, we introduce a modified version of Softmax attention that operates on a specific subset of indices. It defines the top-$r$ nearest neighbors Softmax attention, which focuses on the most relevant entries in the attention matrix. In Section 3.3, we describe the massive activation property for attention mechanisms. In Section 3.4, we present a data structure for efficiently solving the half-space range reporting problem.

2 Softmax Attention with Index Set

Recall that we have already provided the definition of ReLU attention in Definition 1.2. Here, we present the key concepts of Softmax attention. For Softmax attention, since we only calculate the “massive activated” entries to get our approximated results, we introduce the formal definition:

as the submatrix of $K$ and $V$, i.e., whose row index is in $R$ or $\overline{R}$, respectively.

In this work, we consider calculating the Softmax attention on the “massive activation” index set, where we define the “massive activation” index set as the top-$r$ indices. We introduce our definition for top-$r$ indices of Softmax attention as follows:

3 Massive Activation

Now, we introduce our observations on the properties of the attention scores (the inner products of query vectors and key vectors). This further facilitates the error analysis of the top-$r$ indices Softmax attention. To begin with, we provide the definition of the massive activation property as follows:

The top-$n^{\gamma}$ entries are massive, i.e., $\frac{1}{n^{\gamma}\cdot\|q\|_{2}}\sum_{i\in\mathsf{NN}(n^{\gamma},q,K)}\langle q,K_{i}\rangle\geq\beta_{1}\log(n)$.

The remaining terms are upper bounded, i.e, $\forall i\in[n]\setminus\mathsf{NN}(n^{\gamma},q,K)$, $\frac{1}{\|q\|_{2}}\langle q,K_{i}\rangle\leq\beta_{2}\log(n)$.

An intuitive understanding of Definition 3.3 is that, the summation of “massive activated” entries dominates the summation of all entries, and the entries we ignored only contributes little to the final summation. Therefore, it is reasonable for us to omit those non “massive activated” entries.

There are many distributions satisfying the property in Definition 3.3, such as (1) $K$ drawing from any subexponential distribution, e.g., multivariate Laplace distributions, (2) $K$ drawing from any mixture of Gaussian distribution with $n^{1-\gamma}$ Gaussian clusters.

4 Half-Space Reporting (HSR) Data Structure

We restate the result from for solving the half-space range reporting problem. The interface of their algortihm can be summarized as in Algorithm 1. Intuitively, the data-structure recursively partitions the set $S$ and organizes the points in a tree data-structure. Then for a given query $(a,b)$, all $k$ points of $S$ with $\operatorname{sgn}(\langle a,x\rangle-b)\geq 0$ are reported quickly. Note that the query $(a,b)$ here defines the half-space $H$ in Definition 1.3. We summarize the time complexity of HSR data structure as follows:

Part 1. ${\cal T}_{\mathsf{init}}(n,d)=O_{d}(n\log n)$, ${\cal T}_{\mathsf{query}}(n,d,k)=O(dn^{1-1/\lfloor d/2\rfloor}+dk)$.

Part 2. ${\cal T}_{\mathsf{init}}(n,d)=O(n^{\lfloor d/2\rfloor})$, ${\cal T}_{\mathsf{query}}(n,d,k)=O(d\log(n)+dk)$.

Main Results on Attention Generation

In this section, we present our key findings regarding attention generation, $m=\Theta(1)$, for both ReLU and Softmax attention mechanisms. Across both scenarios, we have reduced the time complexity from a naive $O(mn)$ to $O(mn^{4/5})$. Specifically, for the ReLU attention model, we have managed to accelerate the processing time without introducing any approximation errors. In the case of Softmax attention, our technique results in only an insignificant approximation error.

We begin with introducing our result on ReLU attention generation as follows:

Theorem 4.1 shows that our Algorithm 2 accelerates the running time of ReLU attention generation from naive $O(mn)$ to $O(mn^{4/5})$, which is a significant speed up when the KV Cache is large. The at least $1-\delta$ success probability originates from the sparsity analysis of ReLU attention (Lemma 6.1), where with probability at least $1-\delta$, we have the number of non-zero entries of each row of the attention matrix is at most $n^{4/5}$.

Then, we move on to presenting our result on Softmax attention generation. Our results consist two parts: the improved running time of Softmax attention generation, and the error analysis of Softmax attention with index set. Firstly, we introduce our result about the imporved running time of Softmax attention generation as follows:

Theorem 4.2 demonstrates that if we choose the threshold $b$ satisfying $R=\mathsf{NN}(n^{4/5},q,K)$, we can achieve a significant running time improve of the Softmax attention generation.

It is evident that this method introduces an approximation error due to the exclusion of certain entries. Nevertheless, under mild assumptions about the distribution of the attention scores, we demonstrate that this approximation error is indeed negligible. The proof’s intuitive explanation lies in the fact that the majority of attention scores are focused on the small subset of entries that we retain. We organize our result as follows:

Theorem 4.3 presents the error of Softmax attention with index set is relatively small. Consequently, omitting the remaining less significant entries is a justifiable compromise.

With mild assumptions on $V$, we can have more precious results from Theorem 4.3. For example, if the entries in $V$ conform to subgaussian distribution with constant variance, we have $\|V\|_{\infty}=O(\log(n))$ with high probability.

Extension on Full Attention Computation

In this section, we extend our results to full attention computation scenario, where the number of queries and keys is proportional, i.e., $m=\Theta(n)$. Essentially, the full attention computation is beneficial in practical applications, particularly within the context of cross-attention computations. For ReLU attention, we leverage Part 1 result of Corollary 3.5 to accelerate the identification of non-zero entries (activated entries). We introduce our result on ReLU attention as follows:

In Theorem 5.1, we improve the running time of full ReLU attention computation from $O(n^{2})$ to $O(n^{2-1/\lfloor d/2\rfloor}+n^{1+4/5})$, which is a notable uplift of the running time when $n$ is extremely large.

Then, we present our result on Softmax attention. Intuitively, we use the Part 1 result of Corollary 3.5 to identify those “massive activated” entries (top-$r$ indices) within the attention matrix of Softmax attention, and calculate the Softmax attention with top-$r$ indices. We organize our result as follows:

Theorem 5.2 demonstrates our $O(n^{2-1/\lfloor d/2\rfloor}+n^{1+4/5})$ running time on Softmax full attention computation, which improves from naive running time $O(n^{2})$.

Technical Overview

In Section 6.1, we introduce our analysis about the sparsity in the ReLU attention mechanism. In Section 6.2, we present our results of two general attention frameworks. In Section 6.3, we provide our error analysis of Softmax attention with index set. We have shown that with mild assumption on the distribution of attention scores, the error of Softmax attention with index set is negligible.

Intuitively, the ReLU activation will deactivate some key and query pairs. We introduce the results of employing the concentration inequality to quantitatively analyze the number of non-zero entries.

2 General Attention Frameworks

First, we introduce our general framework for attention generation computation. Here, we use the Part 1 result of the HSR data structure. As for this framework is designed for the attention generation task, the key matrix $K$ is fixed in each inference. Therefore, in the Init procedure, we initialize the HSR data structure with the key matrix $K$. Then, in each inference, we use the same HSR data structure to answer the query from each row of the query matrix $Q$. We provide the result of this general attention generation framework as follows.

Part 1. The Init procedure runs in $O(n^{\lfloor d/2\rfloor})$ time.

Part 2. For each query, the Inference procedure runs in $O(mn^{4/5})$ time.

The general framework for full attention computation is quite different from the previous one. Namely, we choose the Part 2 result of the HSR data structure. Since in each inference, both the query matrix $Q$ and the key matrix $K$ differ from any other inference, we first initialize the HSR data structure with the key matrix $K$. Then for each row of the query matrix $Q$, we query the HSR data structure to find the activated entries.

3 Error Analysis of Softmax Attention with Top-r𝑟r Indices

Calculating the Softmax attention on the “massive actavted” index set will introduce approximation error. In the following Lemma, we analyze the quantity of this approximation error. Here, we use $\alpha$ to denote the summation of all entries activated by $\exp(x)$ function, and we use $\overline{\alpha}$ to denote the summation of those entries which are excluded from “massive activated” index set. We provide the general error bound of Softmax attention with index set as follows.

Note that Lemma 6.4 only provides a general error analysis of Softmax attention with index set. Under mild assumption on the distribution of attention scores, we show that this error is actually very small. For more details, please refer to Theorem 4.3.

Experiments

In this section, we present our empirical results of evaluating three mainstream LLMs with Softmax attention with top-$r$ indices on different $r$, showing that the results of the experiments are consistent with our theoretical analysis.

Datasets. To estimate the approximation error of the Softmax attention with “massive activation” entries, we conduct experiments on the PaulGrahamEssays datasets from LLMTest-NeedleInAHaystack . Specifically, for each article in the dataset, we first input $2^{15}=32768$ tokens to the LLMs, then generate $1024$ tokens.

Models. To demonstrate the generalization of our approximation error bound, we conducted experiments on three mainstream large models: LLaMA 3.1 8B Instructhttps://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct , Mistral Nemo 12B Instructhttps://huggingface.co/mistralai/Mistral-Nemo-Base-2407 , and Phi 3.5 Mini 3.8B Instructhttps://huggingface.co/microsoft/Phi-3.5-mini-instruct .

Results. The experiments are conducted on the setting discussed in previous paragraphs. We evaluated the performance of three mainstream LLMs using Softmax attention with top-$r$ indices. In particular, we chose $r$ from the set $\{2^{2},2^{4},2^{6},2^{8},2^{10},2^{12},2^{15}\}$. As depicted in Figure 2, a significant increase in the perplexity (drop in performance) of LLMs is observed only when $r$ falls below $2^{4}$. This suggests that the “massive activated” tokens are predominantly found within the top-$2^{4}$ entries. In comparison to the total of $2^{15}$ entries, the “massive activated” entries constitute a relatively minor fraction. The observed results align with our theoretical analysis, confirming that the approximation error of the Softmax attention mechanism with top-$r$ indices is insignificant for larger values of $r$.

Discussion and Future Work

The sparsity within neural networks arises primarily from the incorporation of non-linear activation functions. These non-linear functions determine the mechanism or circuit of the neural networks, e.g., the induction head in transformers . Gaining insight into these non-linear layers not only enhances our understanding of how neural networks work but also paves the way for optimizing training and inference. We hope our analysis may inspire efficient neural network architecture design. This work represents the initial point of this envisioned blueprint. We concentrate on analyzing the combinations of LLMs and fundamental non-linear activation functions—ReLU and Softmax, which are most relevant to contemporary applications. By analyzing these functions, we aim to demonstrate to the research community that a thorough examination of a model’s non-linear characteristics can significantly enhance the running time complexity of neural networks.

Conclusion

This work investigates the exploitation of the intrinsic sparsity present in both ReLU and Softmax attention mechanisms to decrease the computational complexity of full attention computation and attention generation scenarios. Specifically, we employ the Half-Space Reporting (HSR) data structure to accelerate the process of identifying non-zero or “massive activated” entries within ReLU and Softmax attentions, respectively. Importantly, our approach does not import any errors to ReLU attention, and it results in only a negligible approximation error for Softmax attention.

Acknowledgement

Research is partially supported by the National Science Foundation (NSF) Grants 2023239-DMS, CCF-2046710, and Air Force Grant FA9550-18-1-0166.

References

Appendix A Preliminary

In this section, we display more fundamental concepts. In Section A.1, we introduce several important probability properties and bounds. In Section A.2, we detail the time complexity and performance of half-space reporting (HSR) data structures.

We state several fundamental properties and bounds for some common distributions.

For $i\in[d]$, let $x_{i}$ denote the $i$-th entry of $x$.

Suppose for $i\in[d]$, $y_{i}\sim\mathcal{N}(0,\sigma^{2})$.

Then the inner product of $x$ and $y$, $\langle x,y\rangle$ conforms Gaussian distribution $\mathcal{N}(0,\|x\|_{2}^{2}\sigma^{2})$. Namely, we have $\langle x,y\rangle\sim\mathcal{N}(0,\|x\|_{2}^{2}\sigma^{2})$.

For any $i,j\in[n],i\neq j$, $y_{i}$ and $y_{j}$ are independent.

Then, for any $i,j\in[n],i\neq j$, $\langle x,y_{i}\rangle$ and $\langle x,y_{j}\rangle$ are independent.

We provide tail bounds for chi-square and Gaussian distributed random variables:

Let $X\sim\mathcal{X}_{k}^{2}$ be a chi-squared distributed random variable with $k$ degrees of freedom. Each one has zero means and $\sigma^{2}$ variance.

Suppose we have a random variable $x\sim\mathcal{N}(\mu,\sigma)$.

where the first step, the second step follows from basic algebra, the third step follows from that the inequality holds for any $\lambda>0$, and the fourth step follows from Markov’s inequality.

Then we consider the numerator and we use $y=x-\mu$ to simplify the calculation, we have

where the first step follows from the definition of the moment generating function, the second and the third steps follow from basic algebra, and the fourth step follows from the property of the probability density function.

where the first step follows from Eq. (A.1) and Eq.(A.1), the second step follows from the calculation of infimum. ∎

The Bernstein’s inequality for bounding sums of independent random variables is:

A.2 Half-Space Reporting (HSR) Data Structures

The time complexity of the HSR data structure is:

Let $d$ be a fixed constant. Let $t$ be a parameter between $n$ and $n^{\lfloor d/2\rfloor}$. There is a dynamic data structure for half-space reporting that uses $O_{d,\epsilon}(t^{1+\epsilon})$ space and pre-processing time, $O_{d,\epsilon}(\frac{n}{t^{1/\lfloor d/2\rfloor}}\log n+k)$ time per query where $k$ is the output size and $\epsilon>0$ is any fixed constant, and $O_{d,\epsilon}(t^{1+\epsilon}/n)$ amortized update time.

Part 1. ${\cal T}_{\mathsf{init}}(n,d)=O_{d}(n\log n)$, ${\cal T}_{\mathsf{query}}(n,d,k)=O(dn^{1-1/\lfloor d/2\rfloor}+dk)$.

Part 2. ${\cal T}_{\mathsf{init}}(n,d)=O(n^{\lfloor d/2\rfloor})$, ${\cal T}_{\mathsf{query}}(n,d,k)=O(d\log(n)+dk)$.

Appendix B Full ReLU Attention Computation

In this section, we focus on optimizing the standard ReLU attention calculation. By leveraging a HSR data structure and assuming sparsity, the time complexity can be reduced to $O(n^{1+4/5}d)$.

Assume each entry of $K$ is from Gaussian ${\cal N}(0,\sigma_{k}^{2})$, and each entry of $Q$ is from Gaussian ${\cal N}(0,\sigma_{q}^{2})$.

Let $\delta\in(0,1)$ denote the failure probability.

Let $\sigma_{a}=4\cdot(1+d^{-1}\log(m/\delta))^{1/2}\cdot\sigma_{q}\sigma_{k}$.

Let hsr data structure be defined as Part 1 in Corollary A.7.

There exists an algorithm (Algorithm 3), with at least $1-\delta$ probability, computes full attention of $Q,K,V$ in $O(mn^{1-1/\lfloor d/2\rfloor}+mn^{4/5})$ time.

The running time for Inference procedure can be written as

The first term ${\cal T}_{\mathsf{init}}(n,d)$ corresponds to the initialization of the hsr data structure. Since we use Part 1 result from Corollary A.7, the running time for initialization is ${\cal T}_{\mathsf{init}}(m,d)=O_{d}(m\log m)$.

The second term $\sum_{i=1}^{m}{\cal T}_{\mathsf{query}}(n,d,\widetilde{k}_{i})$ comes from the HSR query operation (Line 11). Since we use Part 1 result from Corollary A.7, we have

where the first step follows from ${\cal T}_{\mathsf{query}}(n,d,\widetilde{k}_{i})=O(dn^{1-\lfloor d/2\rfloor}+d\widetilde{k}_{i})$ (Part 1 of Corollary A.7), the second step follows from with high probability $\widetilde{k}_{i}$ at most $n^{4/5}$ (Lemma D.3).

The third term $O(\sum_{i=1}^{m}\widetilde{k}_{i})$ corresponds to calculating $A_{j,i}$ (Line 13). By Lemma D.3, we have the third term is $O(mn^{4/5})$.

The fourth term $O(\sum_{i=1}^{m}\widetilde{k}_{i})$ corresponds to calculating $D^{-1}AV$. Since for $i$-th row of $A$, there are $\widetilde{k}_{i}$ non-zero entries. Therefore, it takes $O(\sum_{i=1}^{m}\widetilde{k}_{i})$ time for calculating $D^{-1}A$. Therefore, it takes $O(d\sum_{i=1}^{m}\widetilde{k}_{i})$ time to calculate $D^{-1}AV$. By Lemma D.3, with high probability, $\widetilde{k}_{i}$ is at most $n^{4/5}$. Therefore, we have the third term as $O(mn^{4/5}d)$.

To sum up, the overall running time is $O(mn^{1-1/\lfloor d/2\rfloor}d+mn^{4/5}d)$. ∎

We can now derive a more specific result for the full ReLU attention computation:

Let ReLU attention be defined as Definition 1.2.

Assume each entry of $K$ is from Gaussian ${\cal N}(0,\sigma_{k}^{2})$, and each entry of $Q$ is from Gaussian ${\cal N}(0,\sigma_{q}^{2})$.

Let $\delta\in(0,1)$ denote the failure probability.

Let $\sigma_{a}=4\cdot(1+d^{-1}\log(m/\delta))^{1/2}\cdot\sigma_{q}\sigma_{k}$.

There exists an algorithm (Algorithm 3), with probability at least $1-\delta$, takes $O(n^{2-1/\lfloor d/2\rfloor}d+n^{1+4/5}d)$ time to compute the full ReLU attention of $Q,K,V$.

By Lemma B.1, we have that the FullAttentionComputation data structure (Algorithm 3) can run Inference to calculate the ReLU attention, in $O(m^{1-\lfloor d/2\rfloor}nd+mn^{4/5}d)$ time.

Then, we have the ReLU attention calculation run in $O(n^{1+4/5}d)$ time. ∎

Appendix C ReLU Attention Generation

In this section, we present a theoretical analysis of the time complexity of ReLU attention generation using a HSR data structure.

Assume each entry of $K$ is from Gaussian ${\cal N}(0,\sigma_{k}^{2})$, and each entry of $Q$ is from Gaussian ${\cal N}(0,\sigma_{q}^{2})$.

Let $\delta\in(0,1)$ denote the failure probability.

Let $\sigma_{a}=4\cdot(1+d^{-1}\log(m/\delta))^{1/2}\cdot\sigma_{q}\sigma_{k}$.

Let hsr data structure be defined as Part 2 in Corollary A.7.

Then, there exists an algorithm (Algorithm 2), with at least $1-\delta$ probability, has the following performance:

Part 1. The Init procedure runs in $O(n^{\lfloor d/2\rfloor})$ time.

Part 2. For each query, the Inference procedure runs in $O(mn^{4/5}d)$ time.

The Init procedure only runs the initialization of the HSR data structure. Since we use Part 2 result from Corollary A.7, the running time of Init procedure is ${\cal T}_{\mathsf{init}}(n,d)=O(n^{\lfloor d/2\rfloor})$.

The running time for Inference procedure can be written as

The first term $\sum_{i=1}^{m}{\cal T}_{\mathsf{query}}(n,d,\widetilde{k}_{i})$ corresponds to the HSR query operation (Line 16). Since we use the Part 2 result from Corollary A.7, we have

where the first step follows from ${\cal T}_{\mathsf{query}}(n,d,k)=O(d\log n+dk)$ in Part 2 of Corollary A.7, the second step follows from with high probability, $\widetilde{k}_{i}$ is at most $n^{4/5}$ (Lemma D.3), the third step follows from $\log n<n^{4/5}$.

The second term $O(d\sum_{i=1}^{m}\widetilde{k}_{i})$ corresponds to calculating $A_{i,j}$ (Line 18). There are $m$ iterations, and in each iteration, it calculates $\widetilde{k}_{i}$ entries of $A$. Then, the second term is $O(d\sum_{i=1}^{m}\widetilde{k}_{i})$. By Lemma D.3, with high probability, $\widetilde{k}_{i}$ is at most $n^{4/5}$. Therefore, we have the second term as $O(mn^{4/5}d)$.

Similar to the proof of Lemma B.1 this term is $O(mn^{4/5}d)$.

To sum up, we have the overall running time for Inference procedure is $O(mn^{4/5}d)$. ∎

We now derive a comprehensive sparsity analysis for the ReLU attention mechanism:

Let ReLU attention be defined as Definition 1.2.

Assume each entry of $K$ is from Gaussian ${\cal N}(0,\sigma_{k}^{2})$, and each entry of $Q$ is from Gaussian ${\cal N}(0,\sigma_{q}^{2})$.

Let $\delta\in(0,1)$ denote the failure probability.

Let $\sigma_{a}=4\cdot(1+d^{-1}\log(m/\delta))^{1/2}\cdot\sigma_{q}\sigma_{k}$.

There exists an algorithm (Algorithm 2), with at least $1-\delta$ probability, takes $O(mn^{4/5}d)$ time to generate the answer.

We make use of the AttentionGeneration data structure (Algorithm 2) in Lemma C.1.

For calculating the attention between $q_{i}$ and $\{k_{1},k_{2},\cdots,k_{i-1},k_{i}\}$, it takes $O(i\cdot d)$ time.

There are $m$ queries in total. The running time for $m$ queries is $O(mn^{4/5}d+m^{2}d)$.

Since we have $n\gg m$, the overall running time for the generation is $O(mn^{4/5}d)$. ∎

Appendix D Sparsity Analysis

To begin our analysis, we first examine the application of Bernstein’s inequality to the matrix $K$:

Let the ReLU attention be defined as Definition 1.2.

Assume each entry of $K$ is from Gaussian ${\cal N}(0,\sigma_{k}^{2})$

Let $\sigma_{a}=\|x\|_{2}\sigma_{k}/\sqrt{d}$.

Then, we can show that, with probability at least $1-\exp(-\Omega(n\cdot\exp(-\frac{b^{2}}{2\sigma_{a}^{2}})))$, the number of non-zero entries $\widetilde{k}_{i}$ is at most $2n\cdot\exp(-\frac{b^{2}}{2\sigma_{a}^{2}})$. Namely, we have

Let $r_{i}\in\{0,1\}$ be the indicator function of $\langle x,K_{i,*}\rangle$. Then, we have $\widetilde{k}_{i}=\sum_{j=1}^{n}r_{j}$.

Since $r_{i}$ is an indicator function, then we have

By assumption, we have $K_{i,j}\sim\mathcal{N}(0,\sigma_{k}^{2})$.

Let $\sigma_{a}=\|x\|_{2}\cdot\sigma_{k}/\sqrt{d}$.

By the property of Gaussian distribution (Fact A.1), we have $\langle x,K_{i,*}\rangle\sim\mathcal{N}(0,d\cdot\sigma_{a}^{2})$ and $\langle x,K_{i,*}\rangle/\sqrt{d}\sim\mathcal{N}(0,\sigma_{a}^{2})$.

For any $i,j\in[n]$, by Fact A.2, we have $\langle x,K_{i,*}\rangle$ and $\langle x,K_{j,*}\rangle$ are independent, which implies $r_{i}$ and $r_{j}$ are independent.

By the tail bound of Gaussian distribution (Fact A.4), we have

Since we have $\widetilde{k}_{i}=\sum_{j=1}^{n}r_{j}$, by Eq. (3), we have

Let $k_{0}:=n\cdot\exp(-\frac{b^{2}}{2\sigma_{a}^{2}})$. By the Bernstein inequality (Lemma A.5), we have

Then, we reach our conclusion: with probability at least $1-\exp(-\Omega(n\cdot\exp(-\frac{b^{2}}{2\sigma_{a}^{2}})))$, the number of non-zero entries in each row of the attention matrix $A$ is bounded by $\widetilde{k}_{i}\leq 2n\cdot\exp(-\frac{b^{2}}{2\sigma_{a}^{2}})$.

We turn our attention to bounding $\|x\|_{2}$:

Assume each entry of $Q$ is from $\mathcal{N}(0,\sigma_{q}^{2})$.

Then, we can show that, for $t\geq 0$ with probability $1-\exp(-t)$, $\|x\|_{2}$ is at most $\sqrt{3}\cdot(d+t)^{1/2}\cdot\sigma_{q}$. Namely, we have

By the assumption, we have $x_{i}\sim\mathcal{N}(0,\sigma_{q}^{2})$.

Since $\|x\|_{2}^{2}=\sum_{i=1}^{d}x_{i}^{2}$, by Chi-square tail bound (Lemma A.3), we have

Since we have $2\sqrt{dt}\leq d+t$, Eq. (5) implies

We can now present our formal sparsity analysis, which builds upon the previous lemmas:

Let the ReLU attention be defined as Definition 1.2.

Assume each entry of $K$ is from Gaussian ${\cal N}(0,\sigma_{k}^{2})$, and each entry of $K$ is from Gaussian ${\cal N}(0,\sigma_{q}^{2})$.

Let $\delta\in(0,1)$ denote the failure probability.

Let $\sigma_{a}=4\cdot(1+d^{-1}\log(m/\delta))^{1/2}\cdot\sigma_{q}\sigma_{k}$.

Then, we can show that, with probability at least $1-\delta$, for all $i\in[m]$, the number of non-zero entries of the $i$-th row $\widetilde{k}_{i}$ is at most $2n^{4/5}$.

This proof follows from applying union bound on Lemma D.1 and Lemma D.2.

We choose $t=d+\log(m/\delta)$. Then, Eq. (6) implies

Let $\sigma_{a}=\|x\|_{2}\cdot\sigma_{k}/\sqrt{d}$. By Eq.(7), we have $\sigma_{a}=4\cdot(1+d^{-1}\log(m/\delta))^{1/2}\cdot\sigma_{q}\sigma_{k}$.

Let $b=\sigma_{a}\cdot\sqrt{0.4\log n}$. Then, Eq. (8) implies

Taking union bound over Eq. (7) and Eq. (9), we have

where the first step follows from the union bound, the second step follows from basic algebra, the third step follows from Eq. (10).

Taking the union bound over $m$ rows in $A$, then we can show that with probability $1-\delta$, for all rows of $A$, that row’s $\widetilde{k}_{i}$ is at most $2n^{4/5}$.

Appendix E Running Time of Softmax Attention

In this section, we provide our results on reducing the running time of Softmax attention. We begin with introducing our result on Softmax attention generation.

The Softmax attention generation scenario can be proved by substituting the ReLU attention $\mathsf{Attn}_{r}$ (Definition 1.2) with Softmax attention with index set $\widehat{\mathsf{Attn}}_{s}$ (Definition 3.2) in Algorithm 2 and Theorem 4.1. ∎

Then, we move on to our result on Softmax full attention computation.

The Softmax full attention computation scenario can be proved by substituting the ReLU attention $\mathsf{Attn}_{r}$ (Definition 1.2) with Softmax attention with index set $\widehat{\mathsf{Attn}}_{s}$ (Definition 3.2) in Algorithm 3 and Theorem 5.1. ∎

Appendix F Error Analysis of Softmax Attention

In this section, we provide an error analysis of the Softmax attention mechanism, deriving error bounds for the general case and a specific case with the massive activation property.

The following lemmas establish error bounds for Softmax attention when using index sets, formalizing the approximation error in attention computation.

Let $\alpha,\overline{\alpha}$ and $\widehat{\mathsf{Attn}}_{s}$ be defined as Definition 3.2.

where the first step is by Definition 3.2, the second step is by triangle inequality, the third step is by $\|uV\|_{\infty}\leq\|u\|_{1}\cdot\|V\|_{\infty}$ for any vector $u$ and conformable matrix $V$, and the fourth step is by definition of $\widehat{\alpha}$ and $\overline{\alpha}$, i.e., $\widehat{\alpha}=\langle\widehat{u},{\bf 1}_{r}\rangle=\|\widehat{u}\|_{1}$ (note that each entry of $\widehat{u}$ is positive), the fifth step is by $\max\{\|\widehat{V}\|_{\infty},\|\overline{V}\|_{\infty}\}=\|V\|_{\infty}$, the sixth step in by simple calculation and the last step is by $\widehat{\alpha}+\overline{\alpha}=\alpha$. ∎

Building on this, we now present a more specific error analysis incorporating the massive activation property:

Let $\gamma\in$, $\beta_{1}\geq\beta_{2}\geq 0$.

Let the Softmax attention with index set $\widehat{\mathsf{Attn}}_{s}$ be defined as Definition 3.2.

Let $\mathsf{NN}(r,q,K)\subseteq[n]$ denote the indices of top-$r$ entries of $qK$.

Let $R=\mathsf{NN}(n^{\gamma},q,K)\subseteq[n]$, where $|R|=n^{\gamma}$.

Assume the query $q$ and key cache $K$ have $(\gamma,\beta_{1},\beta_{2})$ massive activation property.

Let $\alpha,\overline{\alpha},\widehat{\alpha}$ be defined in Definition 3.2. By Lemma F.1, we have

where the first step is by Definition of $\widehat{\alpha}$, the second step is by Definition 3.3 and Jensen inequality, and the last step is by simple calculation.

where the first step is by Definition of $\overline{\alpha}$, the second step is by Definition 3.3, and the last step is by simple calculation.

Finally, we finish the proof by the fact $\widehat{\alpha}+\overline{\alpha}=\alpha$. ∎