HyperAttention: Long-context Attention in Near-Linear Time
Authors: Insu Han
Rajesh Jayaram
Amin Karbasi
Vahab Mirrokni
David Woodruff
Amir Zandieh
Venue: ICLR 2024 poster (2024)
Keywords: self-attention
locality sensitive hashing
long-range context
Primary area: general machine learning (i.e., none of the above)

Abstract: We present an approximate attention mechanism named `HyperAttention` to address the computational challenges posed by the growing complexity of long contexts used in Large Language Models (LLMs). 
Recent work suggests that in the worst-case scenario, the quadratic time is necessary unless the entries of the attention matrix are bounded or the matrix has low stable rank. 
We introduce two parameters which measure: (1) the max column norm in the normalized attention matrix, and (2) the ratio of row norms in the unnormalized attention matrix after detecting and removing large entries. We use these fine-grained parameters to capture the hardness of the problem. 
Despite previous lower bounds, we are able to achieve a linear time sampling algorithm even when the matrix has unbounded entries or a large stable rank, provided the above parameters are small.
HyperAttention features a modular design that easily accommodates integration of other fast low-level implementations, particularly FlashAttention. 
Empirically, employing Locality Sensitive Hashing (LSH) to identify large entries, HyperAttention outperforms existing methods, giving significant speed improvements compared to state-of-the-art solutions like FlashAttention. 
This development presents substantial implications for enabling LLMs to handle significantly larger contexts.

Published as a conference paper at ICLR 2024
HYPERArTENTION: LONG-CONTEXT ATTENTION IN NEAR-LINEAS TIME
Insu Han
Rajesh Jayaram Google Research rk rkjayaram@goog e. com
Amin Karbasi
Yale University
Yale University, Google Research amin karbas @yale le edu
insu. www.hannnnll... edu
Vahab Mirrokni
David P. Woodruff CMU, Google Research dwoodruf@cs edu
Amir Zandich
Google Research
Independent Researcher amir 2ed5 28 gmai com
mizrokni@google... com
ABSTRACT
We present an approximate attention mechanism named "HyperAttention" to ad- dress the computational challenges posed by the growing complexity of long con- texts used in Large Language Models (LLMs). Recent work suggests that in the worst-case scenario, quadratic time is necessary unless the entries of the atten- tion matrix are bounded or the matrix has low stable rank. We introduce two parameters which measure: (I) the max column norm in the normalized attention matrix, and (2) the ratio of row norms in the unnormalized attention matrix af- ter detecting and removing large entries. We use these fine-grained parameters capture the hardness of the problem. Despite previous lower bounds, we are able achieve linear time sampling algorithm even when the matrix has unbounded entries or large stable rank, provided the above parameters are small. HyperAt- tention features a modular design that easily accommodates integration of other fast low-level implementations, particularly FlashAttention. Empirically, employ- ing Locality Sensitive Hashing (LSH) to identify large entries, HyperAttention outperforms existing methods, giving significant speed improvements compared to state-of-the-an solutions like FlashAttention. We validate the empirical perfor- mance of HyperAttention For example, HyperAttention makes the inference time of ChatGLM3 50% faster 32k context length while perplexity increases from 5.6 to 6.3. On larger context length, e.g., 131k, with causal masking, HyperAttention offers 22-fold speedup on a single attention layer.
variety of different long-context length datasets.
INTRODUCTION
Transformers (Vaswani et al., 2017) have been successfully applied to wide variety of learning tasks in areas such as natural language processing (Devlin et al., 2018; Yang et al., 2019; Brown al., 2021),
al., 2020; Raffel al., 2020), computer vision (Carion et al., 2020; Dosovitskiy and time series forecasting (Zhou et al., 2021). Despite their success, these models face serious scalability limitations because naïve exact computation of their attention layers incurs quadratic (in the sequence length) runtime and memory complexities. This presents a fundamental challenge for scaling transformer models longer context lengths.
Various approaches have been explored to tackle the quadratic-time attention layer, with one no- table direction focusing on approximating intermediate matrices in attention layers. Methods for doing this include approximations by sparse matrices (Kitaev et al., 2020; Daras et al., 2020; Roy et al., 2021; Sun et al., 2021; Ding et al., 2023; Han et al., 2023), low-rank matrices (Choromanski et al., 2021; Katharopoulos et al., 2020; Kacham et al., 2023), or combination of both (Chen al., 2021b; Zaheer et al., 2020; Chen et al., 2021a; Dass et al., 2022). These methods aim to provide various components of attention, but none of them provide end-to-end ap- proximations of the full dot-product attention. Moreover, none of these works support the use of causal masking, which is crucial part of modern transformer architecture. In ongoing work,
faster approximation

---

Published as a conference paper at ICLR 2024
Kacham et al. (2023) bypass the hardness of softmax by switching to polynomial functions and ap- plying fast polynomial sketching methods (Ahle et al., 2020; Woodruff & Zandich, 2020; 2022). On the negative side, recent theoretical bounds suggest that entry-wise approximations to the attention matrix are impossible in sub-quadratic time general (Alman & Song, 2023).
Nevertheless, recently, KDEFormer (Zandieh al., 2023) has provided provable approximation in subquadratic time, under the assumption that the entries of the attention matrix are bounded. Theo- retically, KDEFormer runs in roughly 1n11T))) time; it employs kernel density estimation (KDE) to approximate column norms, allowing to compute probabilities with which to sample columns of the attention matrix. However, the current algorithms for KDE are lacking practical efficiency (Charikar al., 2020), and even in theory, there is gap between the runtime of KDEFormer and the theoretically feasible (nn) time algorithms. In (Alman & Song, 2023), the authors demonstrated that under the same assumption of bounded entries, nearly linear time Onn (no)) algorithm is possible. However, their algorithm also involves using the polynomial method to approximate the softmax and is likely impractical (e.g., it was not empirically evaluated by the authors). In this work, we provide an algorithm which achieves the best of both worlds, being both a (1) practically efficient algorithm that (2) achieves the best possible near-linear time guarantee. Additionally, our approach supports casual masking, which was not possible via previous works.
1.1 PROBLEM STATEMENT
The dot-product atlention (Vaswani et al., 2017) involves processing three input matrices: Q (queries), K (keys), V (values), all of size n x d, where n is the number of tokens in the input sequence and d is the dimension of latent representations. This process outputs the following:
Att = D-'AV
Here, matrix A : exp (QKT) is defined the element wiiee exponential of QKT Additionally, D is an x n diagonal matrix derived from the sum of rows of A. Di,, A for [n]. In this context, matrix A is referred to as the "attention matrix", and D A is called the "softmax matrix". It is important note that calculating the attention matrix A directly requires ((n2d) operations, and storing it consumes (nn²) memory. Consequently, a straightforward computation of Att demands runtime of 52(n2d) and (n22) memory.
Our objective is to efficiently approximate the output matrix Att while retaining its spectral prop- erties. Our strategy involves designing an efficient estimator for the diagonal scaling matrix D in near-linear time. Additionally, we aim to quickly approximate the matrix product of the softmax matrix A and value matrix V through subsampling. To be more specific, our objective is to find a sampling matrix m with a limited number m = (l)) of rows, along with a diagonal matrix D E paxx, such that the following bound on the operator norm of the error is met:
Att D-'AS SV op
A
lp |
1.2 OUR CONTRIBUTIONS
We show that efficiently solving the matrix multiplication component of the attention approximation problem in Eq. (1) can be achieved by defining the sampling matrix S based on the row norms of V. The more challenging aspect lies in obtaining reliable spectral approximation for the diagonal matrix D. In recent result, Zandieh et al. (2023) effectively leverages fast KDE solvers to attain high-quality approximation of D. However, we streamline the KDEformer procedure and demon- strate that uniform sampling is sufficient achieve the desired spectral guarantee, eliminating the need for importance sampling based on kernel densities. This significant simplification allows us to develop a practical and provably linear time algorithm.
In contrast to prior work (Alman & Song, 2023; Zandieh et al., 2023), our approach does not necessi- tate bounded entries bounded stable rank. Furthermore, the fine-gra finengrined parameters introduce to analyze the time complexity may remain small even when the entries in the attention matrix or the stable rank are large.
2

---

Published as a conference paper at ICLR 2024
Our work is inspired by the hard instance of Alman & Song (2023) for showing quadratic time lower bounds. Such instances have one randomly placed large entry in each row of the attention matrix. Our algorithm has initial phase where find large entries of the attention matrix in black box manner, such as by using Locality Sensitive Hashing (Kitaev et al., 2020), or a possibly learned CountSketch applied the attention matrix (Charikar et al., 2002; Li et al., 2023a), or just known heavy entry pattern (Chen et al., 202 We assume these procedures are fast, and that after removing the heavy entries, two parameters in the resulting attention matrix are small: (I) the max column fn-norm, and (2) the ratio of row norms in the un-normalized attention matrix.
Prior work of Zandieh et al. (2023) used KDE identify columns in the attention matrix with large norm and to perform approximate matrix product with the value matrix by sampling such columns. As mentioned, finding such columns requires at least O(n 1.13) time. Instead, we observe that by doing a one-sided sampling from the squared row norms of V. we can avoid the use of KDEs and achieve the same spectral norm guarantee in terms of the stable rank. Although our algorithm is simple and just samples by the row norms of the value matrix (or even samples uniformly in practice), the main technical challenge is that we do not know the row norms of the attention matrix needed in order to normalize it and produce proper factorization of it. This is reminiscent of the quadratic time hard instance of (Alman & Song, 2023) where we may not be able to find a heavy entry in a row easily, and thus cannot normalize by its norm in the attention matrix. Our parameters (1) and (2) above allow us to argue that the heavy entries, if they exist, are not distributed in the worst possible way.
Empirically, HyperAttention demonstrates significant speed improvements, achieving over a 50x acceleration in forward and backward propagation for sequence lengths of n 13lk. When deal- ing with causal masking, the method still delivers a substantial 5x speedup. Moreover, when our approach is applied to pretrained LLMs, e.g., chatglm3-6b-3 m2-33kkkk6 (Du et al., 2022) and evaluated on long-context benchmark datasets, so-called LongBench (Bai et al., 2023), it maintains performance levels that closely match those of the original models, even without the need for fine-tuning. Fur- thermore, we investigate task-specific evaluations and discover summarization and code completion tasks are more robust to approximate attention layers than question answerings.
PRELIMINARIES
We make use of the Hamming sorted LSH, variant of angular sensiit hashing introduced in the work by Zandich et al. (2023). In this variant, the hash buckets are arranged in order of their Hamming distances. This LSH variant is particularly well-suited for designing GPU-friendly algorithms aimed at identifying dominant entries within the attention matrix A. In the context of Hamming sorted LSH, if let H Rd + [B] be a hash function with B buckets drawn from an LSH family, then the collision probability PHи [H(g) = H({)) is "roughly" proportional to (q, k). A very useful property of this LSH variant is that its buckets are ordered in such a way that geometrically adjacent buckets have consecutive buckets. We provide the following definition.
Definition 1 (Hamming sorted LSH, Definition 7.3 of (Zandieh et al., 2023)). For positive integer r, there exists an LSH function H : Rdd + [2"), such that for any x., y E R2d its collision probability is Pr(H(()) H(v) =
B(a.y) where (x. u) == cos -1
xx' TTy a lill Furthermore, this LSH
function hashes similar points to adjacent buckets. Specifically, the probability that two points end up in adjacent buckets is given by Pr (H((x) H(y) + (mod 2") 28(2.)) (1 0(z.0)
Using this LSH function, as demonstrated by Zandieh et al. (2023), we can sort keys and queries within an attention layer in such a way that large entries get shifted towards the diagonal of the attention matrix. Subsequently, these significant entries in the attention matrix can be captured by computing equal-sized blocks along the diagonal. This approach aligns with the block- memooyy ac- cess patterns of modern hardware and can be efficiently parallelized through batching across blocks.
3

---

Published as a conference paper at ICLR 2024
11
102
bg
10g
ky Kig K7 Ni
Ko Ry
Ko
91
22
92
PK
95
P
44
43
PQ
96
PQ
07
Q9
99
91
A
APQ
M O A
Figure I: How sortLSH finds large entries of A: (Left) Keys and queries undergo hashing using the Hamming ordered LSH H(-). (Middle) Keys and queries are rearranged based their hash buckets. Attention matrix after applying these row and column permutations is denoted as P Large entries of Apo Pk concentrated around the diagonal blocks. (Right) rows and columns permutations are reversed on the attention matrix and M R.
A is highlighted.
Algorithm 1: sortLSH for locating large entries of A
1: input: matrices Q, K E Rexd, and block size b
2: Let H((-) be Hamming sorted LSH as per Definition and hash rows Q, K
3: Let Pk,, PQ E Sym(n) be permutations satisfying Pk (i) < Pk( if H(K,,) S H()),)
and (1) < PQ(() H (QQ..) I H(Qj.:)
4: return Mask matrix M H E {0, 1]mx defined as M. ( (PQ())/)=PPX(I1) 1
3 ALGORITHM
To obtain a spectral guarantee when approximating Att, our initial step involves producing a 1 approximation of the diagonal entries in the matrix D. Subsequently, we approximate the matrix product between D 'A and V via sampling according to the squared row eg-norms of V.
Estimating Our procedure for approximating consists of two steps. Initially, we identify the dominant entries within the attention matrix using an algorithm rooted in the Hamming sorted LSH, as defined in Definition 1. The second step revolves around randomly selecting a small subset of keys K. We will demonstrate that under certain mild assumptions about matrices A and D, this simple approach allows us to establish spectral bounds on the estimated matrix. Our aim is to find sufficiently precise approximate matrix D that satisfies:
|| (D-t-D-1 A
125 D Al
Our assumption is that the column norms of the softmax matrix exhibit relatively uniform distri- bution. To be more precise, assume that for any E [ni] there exists some (1) such that
D-'A 'A- e 2 < a. It's worth noting that our assumption more general in comparison to the bounded input entries assumption made in (Alman & Song, 2023). In fact, if their assumption holds,
implies that | A ! <
n o for all E [n]. In Section 4.3, we empirically compute a
to be the maximum of the squared normss of the columns in D A and verify that it is indeed sublinear n.
The first step of our empirical algorithm involves identifying large entries of the attention matrix A through hashing keys and queries into uniformly-sized buckets using the Hamming sorted I.SH, which we refer to sortLSH. This process is detailed in Algorithm and is visually illustrated in Fig. I. Note that we also mention other was of identifying large patterns, such as checking for a known heavy hitter pattern, using CountSketch which describe more below.
Algorithm returns a sparse mask designed to isolate the dominant entries of the attention matrix. Given this mask, compute approximation of the matrix D in Algorithm that satisfies the spectral guarantee in Eq. (2). This algorithm accomplishes this by combining the attention values corresponding to the mask with randomly chosen subset of columns from the attention matrix.

---

Published as a conference paper at ICLR 2024
Algorithm 2: ApproxD for estimating diagonal matrix D
1: input: matrices Q, K Rexd, large entries mask M H E (0,1)*X, parameters > 0,
ęT, a > K²K, and integer m
2: Randomly choose subset T S [r]] with cardinality - =
3: T < maxjer (1- M e (.))
festimate of maximum un-masked row sum}
4: Generate m L..dd sample &,, l2................... lm ~ (n])
5: for i [n] do
6: C. C &'
(MM exp(K + +/K)) {(capped) row-sum of masked entries)
Flog
{d, below is the (capped) row-sum estimator for the un-masked entries}
7:
di + Em (1 M, min (exp (1Qr.. Ke,a)),
8: d; + (M. exp K (Q,))) + max (d;, T/K)
{full row-sum estimate)
9: return diagonal matrix D diag( {d. Hin)
Algorithm 3: HyperAttention: attention mechanism in near-linear time
1: input: matrices Q, K, V E Rxxd, mask matrix M E {0, 1)mxn, and parameter
notry
2: Run Algorithm 2 and let D <- APPRoXD (Q, K, M _n(()) E, no(1) d ddd.'"' no())
3: Let mmxx be an i.i.d. sampling matrix based on squared row norms of as Lemma
4: return D and S
The assumptions of Lemma are used ensure that the variance of the estimator is small, and the same complexity of the algorithm increases as a function of the parameters a, We remark that our algorithm is versatile and can function effectively with predefined mask that specifies the positions of dominant entries within the attention matrix, mirroring the approach taken in (Chen et al., 2021a). The main guarantee provided by this algorithm is given Lemma 1.
Lemma (Approximating D). For any Q, K E RRadd let A D Rax" be the diagonal matrix with Dz, ||A Additionally, suppose that
exp(QKT) Also let
max4e(n) | A 112 For any mask matrix M E f0, 1]nxn let define the condition mAX (1 MR A..))
number k = minyetu (1- -MM A,,)" If 1fm= =
(==1a² logn). the output D of Algorithm satisfies
Eq. (2) with probability at least poly(a))
Approximating the product of softmax matrix D- D-LA A and values matrix V. Given D that meets the spectral approximation conditions as in Eq. (2), we can achieve the spectral constraint in Eq. (1), by finding sampling matrix that satisfies the following condition,
D-1AST SV D-LAV op -
We can efficiently find sampling matrix S mg with small number of rows that satisfies Eq. (3) by leveraging well-established techniques in Approximate Matrix Multiplication (AMM).
Lemma 2. For any matrices D. Raxn, V E Rxd consider a sampling matrix S E mxn
constructed as follows: first generate i.i.d. samples f1,... Lm E [n]] according squared
IIV...ll then let the 7 row of S be vm- (ttr), If m I
[|VI|
norms of matrix V. i.e.,
Ω. (e-2d ( srank((n(( D-1A)) for some > 0, the following holds with probability least 0.99:
ASS SV
D-1 -V---1
AV
|
Se D D-
||Vllpp
The above result is standard and for proof refer to (Drineas & Kannan, 2001).
Main Theorem. Now, we can integrate the subroutines for approximating the diagonal D and approximating the matrix product between D-1 A and values matrix V. With this, we introduce the
5

---

Published as a conference paper at ICLR 2024
HyperAttention, efficient algorithm that can approximate the attention mechanism with spectral guarantees as per Eq. (1) in near-linear time. Our Algorithm 3 takes as input a mask M H that defines the positions of dominant entries within the attention matrix. This mask can be generated using the sort SH algorithm (Algorithm I), or it can be a predefined mask similar to the approach taken in (Chen al., 2021a). The large entries mask M is assumed to be sparse by design and its number of nonzero entries is bounded () n + We now introduce main theorem which will be proved in Appendix A.
Theorem 1 (HyperAttention guarantee), For any matrices Q.K. V E Raxd any mask matrix 1 M E {0, (0.1)"x,, and parameter no(1), let A = exp( (OKQQQQ) and let D E Raxn be the diag-
onal matrix with Ds,; = ||A..:I. max || D-
200)) and maxiem M, A...) (1- J.
S
no(l) then with probability at least 0.98 the outputs S, D of Algorithm satisfy the spectral condi- tion as in Eq. (1). Moreover, this algorithm's 's runtime is O(d n +o(1)) + d unz(M" M)).
Note that even if M is not given to us, but M H can be found in d-n ++((1) time, the theorem holds. We also give examples when this possible by using Hamming sorted LSH, which our experiments are based on, or using the kepth of (Larsen et al., 2016) which is based on CountSketch (Charikar et al., 2002) but also gives a fast recovery time. In the supplementary we show:
Corollary 1 (HyperAttention with sortLSH), Suppose all preconditions of Theorem I hold. Further, suppose the mask matrix M E {0, jaxR is generated as in Algorithm I with block size b n1) and = loog2 n in Definition I. We further assume there are most n + pairs (i.j) with 8(Q K,,)) [1 (1--(()), where O is as in Definition 1. Then with probability 1 - 1/n the M we find in Algorithm I has at most n. + non-zero entries and with probability at least .98, the outputs S, D of Algorithm 3 satisfy Eq. (1) and the overall runtime is O(d n n)(+(())).
We note the assumption on the angles of the rows of Q and K in Corollary 1 is satisfied if most rows are drawn uniformly at random from a d-dimensional sphere, since in this case they will be nearly orthogonal, i.c., have angle at most 5 (1 o(]) with high probability, However, the corollarly also allows n +011) pairs of rows to have arbitrary angle, which may be more realistic.
Corollary (HyperAttention with ExpanderSketch). Suppose all preconditions of Theorem 1 hold. Further, suppose the mask matrix MH E (0,1)*x is defined such that there is a threshold T [QK e,112
Te, Then we can find M exactly with probability 1 - O(1/n?), and with probability at least 28, the outputs S, D of Algorithm satisfy Eq. (1). The runtime is Odd n t)))))).
()()) such that M. = if and only if (QK R, 2
The key idea behind the proof of Corollary 2 is to first sketch Q by ExpanderSketch T. which is efficient since T has a small number of rows. Then compute (T- Q)- K" which is again efficient since (T Q) has small number of rows. Thus, we never form the matrix Q K
3.1 CAUSAL MASKING
Language models commonly employ causal masking. The causal mask is a lower triangular binary square matrix denoted as MC where M = (122)) The causal attention mechanism is defined as:
Att. = DE (M @ A)V,
where A = exp (QKT) is defined as before and De is an x n diagonal matrix derived from the sum of rows of the masked attention Mc @ A, specifically [De]., = (MC. A..)) for E [m]. To approximate causal attention with a spectral guarantee, we require two components. First, we need spectral approximation for the diagonal matrix De. Second, we need to approximate the matrix product between D. (M @ A) and V. which can be achieved using the same sampling technique as described in Algorithm 3 and Lemma 2. The first component is more intricate, and we employ recursive method to address it. So we focus on how to efficiently approximate the diagonal De-
Our approach is based on a key observation, as depicted in Fig. 2. The masked attention M @.A can be decomposed into three non-zero matrices, each of which has half the size of the original attention matrix. The block A21, located entirely below the diagonal is unmasked attention. Consequently,

---

Published as a conference paper at ICLR 2024
CausalApproxD (recursive call)
M A11
M2 O A22
Approx
M O A
A21
Figure 2: Causal attention matrix can be divided into three equal-sized non-zero sections: M AA1 and M A22 are both causal attention matrices, and Ag1 is unmasked attention matrix.
Algorithm 4: CausalApproxD, recursive approximation of De for causal masking
1: input: matrices Q, K Roxd
2: Split Q and into equal sized sub-matrices: Q = qQI I and K [K] K.
3: De1 & CausalAAprooDD(Q),, Ki) and Dc22 + Causal CapsalApproxD (Q,, K₂)
4: Run the unmasked algorithm ApproxD (Algorithm 2) Q2, K1 to get D21
Deu 0 0,
5: return De
Dz1 + Dem
approximate its sums using Algorithm 2. The two diagonal blocks M O A11 and Mº 3 A22 shown in Fig. 2 are causal attentions with half the original size. To handle these, we apply recursive approach and further partition them into smaller blocks, and repeat this procedure. We present a pseudocode for this procedure in Algorithm 4.
4 EXPERIMENTS
In this section, we benchmark our algorithms by scaling up existing large language models to handle long-rang rangrrage sequences. All experiments are performed on a single A100 GPU with 40 GB memory and we use FlashAttention 2 (Dao, 2023) for the exact attention computation.
Implementation Detail. We implement HyperAttention based on sortLSH and uniform column sampling. Specifically, first apply sortLSH to all rows in Q, E Raxd Then, each of rows is partitioned into groups where is the block size as in Fig. 1. The i-th set of rows in Q is multiplied by the corresponding optimize Algorithm 2 for approximating D by sharing random indices {L, } with all rows in Q. This corresponds to uniformly sampling m rows in V. To further simplify, we reuse indices (l, 1m for the Approximate Matrix Multiplication (AMM) in Lemma 2. The required operations involve permuting rows, reshaping tensors, and small matrix multiplications. Since every batch, head and block the same configurations, the implementation can be parallelized using GPUs.
K, resulting in a block-diagonal approximation of Apg. Px- Next, we
4.1 MONKEY PATCHING SELF-ATTENTION
We first evaluate HyperAttention on two pre- pre-trained LLMs. We choose three models with different architectures that are widely used in practical applications: chatglm3-60 -b---kbbbkkkk (Du et al., 2022), and phi-1.5 (Li al., 2023b). We patch their final attention layers by replacing with HtperAttentions where é can vary from 0 to the number of all attention layers in each LLM. Note that attentions in both models requires causal masking and we make use of Algorithm 4 by recursively applying until the input sequence lengths are less than 4,096. We set both bucket size b and the number of sampled columns m to 256 for all sequence lengths. We evaluate the performance of such monkey patched models in terms of perplexity and speedup.
We use LongBench (Bai et al., 2023), a collection of long context benchmark datasets, which con- tains 6 different tasks ranging from single and multiple-document question answering, summariza-

---

Published as a conference paper at ICLR 2024
4.5
41
perplexity
4.4
2.0
speedup
perplexity
39
4.3
1.5
38
4.2
37
4.1
36
1.0
0
4
8
12
16
20
24
28
12 16
20
24
number of replaced layers
number of replaced layers
(a) chatg1m3-6b-325
(b) phi-1.5
Figure 3: Perplexity and speedup of chatg1m3-6b-322 (left) and phi-1.5 (right) monkey patched with HyperAttention. We vary the number of replaced attention layers in the final order.
Task
Number of Replaced Layers
single-qa
multi-qa
summarization
few-shot
synthetic
code
0 (exact)
96.07
88.19
60.29
203.39
103.00
107.16
7
91.82
90.02
60.84
202.77
104.00
107.42
14
87.98
89.64
58.84
200.46
101.00
103.19
21
84.61
77.51
61.55
201.13
97.00
101.12
28
53.30
73.98
60.45
200.01
71.00
104.99
Table I: Performance evaluation of chatglm3-6b-32k equipped with HyperAttentions on Long- Bench datasets (Bai et al., 2023). They contain 6 different tasks and we evaluate each of them with its own metric where higher value indicates better performance.
tion, few-shot learning, synthetic tasks, and code completion. We select a subset of dataset whose encoded sequence lengths are larger than 32,76% and trim them if the length is over 32,768 that all data have sequence lengths 32,768. Then, we compute the perplexity (i.e., loss on next tokens prediction) of each model. To highlight the scalability the long sequences, we calculate the total speedup on all attention layers whether performed by HyperAttention or FlashAttention.
The results are summarized in Fig. 3. Observe that chatglm3-68 1atpg11mm6b-32k- shows reasonable perplexity even after monkey patched by Hyper e.g., after replacing 20 layers the perplexity increases approximately by and it slowly goes up until 24 layers. But improves runtimes in attention layers about 50%. If all the layers are replaced then the perplexity goes to up 12 but it runs about 2.3x faster. For phi-1.5, similar happens but the perplexities are linearly increasing as the number of HyperAttention grows.
In addition, we evaluate the performances of monkey patched chatglm3-6b-32k LongBench datasets and compute task-specific evaluation scores on each task including single-documet ques- tion answering, multiple-document question answering, summarization, few-shot learning, synthetic tasks and code completion. Results are provided in Table 1. While replacing HyperAttention gen- erally leads performance degradation, we observe that role can vary depending on the task at hand. For example, summarization and code completion are more robust to other tasks. Notably, when half of all attention layers are patched (i.e., 14 layers), we verify that most of the tasks do not degrade more than 13%. In particular, the performance of the summarization task remained almost unchanged, suggesting that this task may more robust to partial alterations the attention mechanism. We recall that computations in attention layers can be 1.5x faster when = 32k.
4.2 SINGLE SELF ATTENTION LAYER
We further explore the speedup of HyperAttention with varying sequence lengths from 4,096 to 131,072. We measure wall-clock times of both forward and forward+backward operations when they computed with FlashAttention with and without causal masking. All inputs Q, K, V have the same length and their dimensions fixed 64 and the number of attention heads is set 12. We chose the same parameters in
accelerated by HyperAttention. We measure the times
8

---

Published as a conference paper at ICLR 2024
Forward
Forward
50
Forkard+Backward
Forward Backward
speedup
40
speed
30
3
20
2
10
16k
32k
65k
131k
4k
8k
16k
32k
65k
131k
sequence length
sequence length
(a) Without causal masking
(b) With causal masking
Figure 4: Speedup of the exact computation using FlashAttention (Dao, 2023) and HyperAttention (this work) in single self-attention layer during forward and backward operations. For n -131k, HyperAttention runs to 54x faster without causal masking and 5.4x with causal masking.
HyperAttention as described in the previous section. In Fig. 4, we observe that HyperAttention runs to 54x faster without causal masking and 5.4x when the causal masking applies. Although time complexities of both causal masking and non-masking are the same, a practical algorithm for causal masking (Algorithm 1) requires additional operations such partitioning Q. K, V, and merging attention outputs which result in an increase of practical runtime. However, those speedups will increase when the sequence length grows. We believe this opens the door scale self-attention only for inference but also for training or fine-tuning the LLMs to fit significantly long sequences.
4.3 EMPIRICAL VERIFICATION OF ASSUMPTION
In addition, we empirically verify our assumption in Theorem i.e., squared eo-norm of columns in D-1AI is upper bounded by =*( We investigate pretrained transformers with and without causal masking. For non-causal masking attention, we use the T2T-ViT model (Yuan al., 2021) and take Q,K, from its first attention layer on the ImageNet test data set as it is the main computational bottleneck. For each image, we compute
be the largest of the squared &g nn-nomms of the columns in D-' 3--A-ee l| and collect the value over 50k images. The sequence length of the model is given by 3,136 and the averaged values of a is observed be 8.1801. This much smaller than and can be possibly sublinear in n.
To further investigate the dependence on n, we utilize the chatglm3-66-322 and LongBench narrative-qa dataset, changing the sequence length from 1 to 9k. We trim pad the input context so that length is strictly n. Unlike the vision model, we notice that the first columns in D-1 A often contain heavy entries; hence we compute a the largest squared norm excluding the first 32 columns. We collect these values for all heads and layers and compute their average. Fig. 5 plots the value of a with various sequence length It is observed that the value of g decreases as grows, supporting the claim that our assumption a (1) holds in practice.
[D-IA oz
0.08
[D
0.06
(1/r)
0.04
1k 2k
3k
4k
5k
6k
7k
8k
9k
sequence length
Figure 5: Empirical values of a (i.e., the largest squared fn-nrs of the columns in D-1 Theorem assumes = (1) and the plot supports the claim that our assumption holds in practice.

---

Published as a conference paper at ICLR 2024
5 CONCLUSION
In this work, we propose a simple linear time attention approximation algorithm by simplifying the existing algorithm based on kernel density estimation (KDE). We introduce a more general parameterization for spectral approximation guarantee based on the condition number, which does require assumptions used prior work. Our algorithm makes use of sortLSH to find large entries and we adopt fast matrix multiplication via row norm sampling. We additionally study how our algorithm used for causal masking by recursive partitioning. Empirically, we illustrate that pre-trained LLMs using algorithm can enhance both inference and training speeds with only minimal performance degradation.
ACKNOWLEDGEMENT
D. Woodruff worked on this while at Google Research and also while visiting the Simons Institute for the Theory of Computing.
REFERENCES
Thomas D Ahle, Michael Kapralov, Jakob BT Knudsen, Rasmus Pagh, Ameya Velingker, David Woodruff, and Amir Zandieh. Oblivious sketching of high-degree polynomial kernels. In Sympo- sium on Discrete Algorithms (SODA), 2020.
Josh Alman and Zhao Song. arXiv:2302.13214. 2023.
Fast attention requires bounded entries.
arXiv preprint
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Neural Information Processing Systems (NeurIPS), 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. Proceedings of the Eu- ropean Conference Computer Vision(ECCV), 2020.
Moses Charikar, Kevin Chen, Martin Farach-Coltow Finding frequent items in data streams. In International Colloquium on Automata, Languages, and Programming (ICAMP), 2002.
Moses Charikar, Michael Kapralov, Navid Nouri, and Paris Siminelakis. Kernel density estimation through density constrained near neighbor search, In Foundations of Computer Science (FOCS), 2020.
Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models. In Interna- tional Conference on Learning Representations (ICLR), 2021a
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song. Atri Rudra, and Christopher Re. Scatterbrain: Uni- fying sparse and low-rank attention. Neural Information Processing Systems (NeurIPS), 2021b.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane. Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking Attention with Performers. In International Conference on Learning Representations (ICLR), 2021.
Tri Dao. FlashAttention-22: Faster Attention with Better Parallelism and Work Partitioning. 2023.
Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf-efficien atten- tion using asymmetric clustering. Neural Information Processing Systems (NeurIPS), 2020.
10

---

Published as a conference paper at ICLR 2024
Jyotikrishna Dass, Shang Wu, Huihong Shi, Chaojian Li, Zhifan Ye, Zhongfeng Wang, and Yingyan Lin. Vitality: Unifying low-rank and sparse approximation for vision transformer acceleration with a linear taylor attention. arXiv preprint arXiv:2211.05109 2022.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Conference of the North American Association for Computational Linguistics (NAACL), 2018.
Jiayu Ding, Shuming Ma, Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint vrxiv:2307.02466, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition Scale. In International Conference on Learning Representations (ICLR), 2021.
Petros Drineas and Ravi Kannan. Fast monte-carlo algorithms for approximate matrix multiplica- tion. In Proceedings 42nd IEEE Symposium on Foundations of Computer Science, pp. 452-459. IEEE, 2001.
Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, and Jie Tang. Glm: General language model pretraining with autoregressive blank infilling. 2022.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308.1613, 2023.
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. arXiv preprint a:Xiv:2310001655, 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and Francois Fleuret. Transformers
mns:: Fast autoregressive transformers with linear attention. In International Conference on Ma- chine Learning (JCML), 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The Efficient Transformer. In International Conference Learning Representations (ICLR), 2020.
Kasper Green Larsen, Jelani Nelson, Huy L. Nguyen, and Mikkel Thorup. Heavy hitters via cluster- 2016.
preserving clustering. CoRR,
François Le Gall. Faster algorithms for rectangular matrix multiplication. In 2012 IEEE 53rd annual symposium on foundations of computer science, pp. $14-523, IEEE, 2012.
Yi Li, Honghao Lin, Simin Liu, Ali Vakilian, and David P. Woodruff. Learning the positions in countsketch. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023a.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorn,, Suriya Gunasekar, and Yin Tat Lee. Textbooks Are All You Need II: phi-1.5 technical report. arXiv preprint arXiv:2309.054633 2023b.
Colin McDiarmid. Concentration. In Probabilistic methods for algorithmic discrete mathematics, pp. 195-248. Springer, 1998.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the Limits of Transfer Learning with a Unified Text-to- Text Transformer. Journal of Machine Learning Research (JMLR), 2020.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguis- tics (ACL), 2021.
Zhiqing Sun, Yiming Yang, and Shinjac Yoo. Sparse Attention with Learning to Hash. In Interna- tional Conference Learning Representations (ICLR), 2021.

---

Published as a conference paper at ICLR 2024
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Neural Information Processing Systems (NeurIPS), 2017.
David Woodruff and Amir Zandich. Leverage score sampling for tensor product matrices input sparsity time. In International Conference on Machine Learning, pp. 23933-23964, PMLR, 2022.
David P Woodruff and Amir Zandich. Near Input Sparsity Time Kernel Embeddings via Adaptive Sampling. In International Conference on Machine Learning (ICML), 2020.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xinet: Generalized autoregressive pretraining for language understanding. Neural Information Processing Systems (NeurIPS), 2019.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang. Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-toker vit: Training vision transformers from scratch on imagenet. In International Conference on Computer Vision (ICCV), 2021.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. In Neural Information Processing Systems (NeurIPS), 2020.
Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In International Conference on Machine Learning (ICML), 2023.
Haoyi Zhou, Shanghang Zhang. Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. In Conference on Artificial Intelligence (AAAI), 2021.
A OMITTED PROOFS
Here include the proofs that were omitted in the main body of the paper. First, present the proof Lemma
Proof of Lemma 1: First, we show that 7 calculated in line 3 of Algorithm is close to the maximum row sum the matrix (1, - M2) O A. is easy to check that & (1- M cxp ( (Pp()..)) S тR for all é E [n]] because the definition of x in the lemma statement. Furthermore, we define the set:
Sp { E [n] (1 M
> >].
then we can show that [So] is small. subset of rows, denoted by where [T] m. Hence, Pr[[Sp] 2 t] < (1 (1-10 t/n)m. Choosing
Recall that is the maximum of the row sums of a random
m T(kTa3= a²g- log n) and f = O Ta 22) gives that with probability at least poly(a)
Sol < 0 7 a 2...))))
(5)
Next, let define the upper-capped version of matrix A where entries of i-th row on positions where the mask M value equal to zero are capped value C, (line 6 of the algorithm) as:
min (Ai.ir C.)
if M -0 otherwise
Ā paxn À.,
for every i.j [n].
Aij
12

---

Published as a conference paper at ICLR 2024
We proceed by bounding the total mass of large entries of matrix D A lost through capping (i.e., e'm
entries of A that are larger than thresholds C.). If we define constant C := sen logn" we can write,
D (A- (A- 1))1,
Σ
(1 - M
(Aig min (Ai.) /r.
8
Ars min (Ai,). Ci) Di,i
Σ
Σ
1 1
(2' ---1)DD D.,
, 0,50 -11)
logg 2
2
[( E [n] M = 0& An, > C2+ - 1)
1=0
log,
x2
2 C
a o
( -
ed kb . . 12
(6)
(2'
The inequality in Eq. (6) follows because, for every i E [n]. the cardinality of the set {i.j E [n]: Mt = 0 & A., > > (2 - 1)āD } must be bounded by (अ- The proof of this is by contradiction because otherwise, there must be an 1 E [n] such that the cardinality of the set H4
{i E [m] M = 0 A Ar, C. + (2 - 1)(Da, } is at least [Hy > (2
This implies that DD-1A A el01 112 > Eien, (A../Da)² = Eien, (D... + 12'- >
[H log + (2 - 1)²) > n2' however, this contradicts with the precondition of the lemma
about a = n maxie[n]
|D- 1 A ee Now. if we defined the sets S2 S [n] as
S { i E [m]]
eDu
D# S2 = {i fie [m] | - Ar,
D
< A A
(7)
3
then it follows from Eq. (6) that the cardinalities of S₁ and S2 are bounded by
S S O (x a-..a-10 e2n), |S2l S < (N a-l. .."n).
(8)
Next note that di computed in line 7 of the algorithm is an estimator for i-th row norm of the capped and masked matrix (1,, M)*) @ A. Let define an estimator for the unmasked capped matrix A d, := d + (M
%L A..). By invoking Chernoff-Hoeffdhne inequality (see e.g., (McDiarmid, 1998)) along with union bound, because the lemma statement assumes that m = xx log n the following holds simultaneously for all i E [n] \ 82 with probability 1 poly(n)
A
A
< d. <
1 c/6
1 e/6
This inequality combined with definition of S, S2 in Eq. (7) implies that for any i E [m]
\ (S, U82).
(1 =/2 D S dे z
(1+ D Now we bound the operator norm of the error as follows:
| - D-1 A
(5-1- D
A
+
(D
(D-1-D)
D-1)
(
A
op
op
8,UBg
lop
3
|
1
A
| D
A
D
Dim (8)
A
2
op
8/2
lop
2
op
3
DIDE
1
A
x²| Dish
A
!l.p
+ KK
D
A
+
D Pin) A
2
up
op
lpp
3
2
||D
AA
+ x
D
A
+
D
A
+
D (m)
A
(10)
op
op
op
2
op
13

---

Published as a conference paper at ICLR 2024
where the second inequality above follows from the inequality in Eq. (9) and also because the defi- nition of 1 ensures that Da 2 (1 1133)- Da, for any E Sj. The third inequality above follows because the lower capping in line
the algorithm ensures that Dj,s 2 (M A₁t) + 7/K for any E S2 while Djs- S (M AJ.) + 7 for any j $ Sq by definition of set So and we know that Dj S (M). Az.)) for E So-
Finally we conclude the proof by bounding the terms | D for r E 1,, in Eq. (10). Fix some r E 1,2. Let e be the unit-normed vector that realizes the operator norm of DØ. A. assume that v is a non-negative vector. More
Since D& A is a non-negative matrix, w.l.o.g.,
precisely v == arg max xER ||
D A =l|2- One has that | A v =|12 || All.p We
define the sequence of binary matrices B B¹, B²... which have same shape as D 1 A as follows:
Bii = {2 "Va/u<(pr
for every integers / 2 0.
(11)
A S2 'Va/a) (r.
Note that because of the precondition of the lemma about a = n maxie[m] ||D-AA implies ID-1Als, S Va/n, we have the following inequality on each entry of the matrix Dz A:
A el0) e001 which
D A) S Va/n.
t=0
Since D A and u both have non-negative entries, the above inequality implies the following:
Dø. A A.A.vl ul1l22 va/n-
E B² v S va/n.
2 1 1B-.0ll2.
(12)
t=
t=0
Now to bound ||B' 0ll2 we first find bounds on the number of I's in rows and columns of B Using the definition of B¹ in Eq. (11) and the fact that row sums in matrix D A are equal to 1, we have:
Bi. S min(2 Vn/a,n).
(13)
Additionally, using the precondition of the lemma about n n- D- A e I| we have:
|B < m2n [S,1)-
(14)
Now bound the norm | 2|l2 for arbitrary integer t 2 0 as follows:
18,
|8
2 11
B 1o |B 2 001
(Cauchy- Schwarz inequality)
|&,1
< 21 Vn/a-
[ 0 = 2 Vn/a. Σ
B 03
JE(M]
+ Vn/a-
Σ B
SE(n]
S 21 Vn/a 2
+2
bØ 2 Vn/a- 2 2,,++2,18)1)
JEIn]
where the inequality in second line above follows from Eq. (13) and the inequality in the last line follows from Eq. (14). The last equality follows from the assumption that ||0|| 1. Therefore,
2 عبي (n/a)1/4
if 2 < V
||B" . -012 <
2 (n/
otherwise
14

---

Published as a conference paper at ICLR 2024
Now by plugging the above inequalities into Eq. (12) find that:
loga
80
| Allop S Va/n-
Σ
2
B
+
Σ
2 .
B²
1
1=loga
|
a+ [
1/1
3/2
if
S 12
if =-r11 1
n
if
where the last line above follows from the upper bound on the size of sets S, we obtained in Eq. (5) and Eq. (8). Finally, by plugging the above inequality into Eq. (10) and using the fact that D-1 is a row-stochastic matrix and thus || D
A op 2 1 the lemma follows.
Next, we prove the main theorem.
Proof of Theorem 1: The diagonal matrix in line 2 is computed by invoking Algorithm 2. By Lemma 1 with probability at least poly(n)* holds that ( A S e/2
op
|D- All Furthermore, in line 3 of the algorithm S is defined as the sampling matrix according
the row norms of V. To invoke Lemma we need have a bound on the stable rank of D-1 A.
First, from Lemma we know that 5
-
A
ž (1-e/2), ( || Allop 2 1/2, where the second
op
inequality follows because D-1
A is row-stochastic matrix. Therefore we have srank D 'A) S
D
A
All F Second, the lower capping in line of Algorithm 2 ensures that D- 1 S S no()) thus srank( (D'A)--A)) S nl)
2
D
S (1)
With this bound on the stable rank and since m d - o11) = 52 (2 2d srank( D-1A)).
Lemma implies that A- SV - D- S e/3 カ- D-1A A Vilpp S 8/2
| Al.p
VVllop with probability 0.99.
By combining these two inequalities, we obtain the spectral approximation guarantee in Eq. (1). The runtime of this algorithm is primarily determined by the time it takes invoke Algorithm 2 in line 2, which is dominated by the time to multiply two matrices of sizes n x and x m and the time to calculate attention matrix entries at the positions defined by M Using the matrix multiplication result from (Le Gall, 2012), the first computation be done in time O(dn +o()) and the latter can be done in nnz m (M*).
Proof of Corollary 1: Because r = log, n, we have Pr..) H(K,,,) S 1/n"- when- ever ((Q... K;,-) 2 x (1 - o(1)). As there are at most n² total pairs, the expected number of such 2
pairs that collide under H is at most n I=o(1)) and so by a Markov bound is at most n 11+((1) with failure probability (1)
Since we also assume there are at most 1++(11 pairs (1, with (QQ,,, K2,) (1 - o()), there can be at most n ++((1) additional pairs that collide.
Thus, in total we have n ++(1) collisions, and consequently the number of non-zero entries in M at most ++1( with failure probability 1/no(1), The proof now follows by the assumptions of the corollary statement as well as Theorem 1.
Proof of Corollary Let T be an ExpanderSketch matrix with O(7 log n)) rows. By the guarantees (Larsen et al., 2016) of T. we have that with probability least 1 n3, for any fixed vector R..
recover set S of indices i [r]] such that if x3 > || then i E S,
that from T x,
whereas if xf <
||>|| 27
then i S. Further, S can be found from Tr. in (1)) time.
15

---

Published as a conference paper at ICLR 2024
We compute T Q. followed by - Note that the time for this computation is O(rn log n). Next, for each column j of QK this allows us to construct a set Sy with the property that
(Q K T )cs 2 IIQ then E Sj. This holds simultaneously for all columns with probability at least 1 n² by union bound. The time for constructing all the sets Sj is n nl++ol)dd
Note that |Sy] < 2r for all j. and we can explicitly compute the exact value of (Q KT), )i. for all i E Sn and all J. in (nrd) time. By the assumptions the corollary, we have that Sş contains superset of the support of the j-th column of M and since we can compute the values exactly, we can exactly construct the mask M matrix that the corollary requires, and n ( d time. The proof now follows by the assumptions of the statement as well as Theorem 1.
16


Reviews:
[{"review_id": "2dEEyPmpgY", "reviewer": "Reviewer_cCwA", "summary": "The paper proposes a new approach of accelerating the computation of attention layer. The proposed method has near linear time (under plausible assumption on column norm) and demonstrates significant advantage in experiments.\n\nThe attention layer in Transformer requires $\\Omega(n^2)$ computation given input sequence of length $n$, this costs is prohibitive for long sequence and there are many recent work on reducing this computation cost. The paper proposes a new approach that uses near linear time, when the column norm of the attention matrix (i.e., $\\exp(KQ)$, where $K$ is the key matrix and $Q$ is the query matrix at attention layer) is small (or balance). This is a challenging task as one needs to take account into the softmax function. The algorithm first uses LSH to identify large entries, then apply fast approximate matrix multiplication by random sampling rows.\n\nThe paper also performs empirical study on real word dataset. When the sequence length is roungly 30k, the inference time decreases roughly 50% with a slight increase of perplexity (roughly 20% -- 50%, depend on the dataset and the number of layer replaced). The speedup is significant for longer sequence with one single transformer layer (up to 50x without causual masking and 5x with causual masking).\n\n--------------------------------------------------------------------\nI have read the author's response and I would keep my positive evaluation towards the paper. Personally I feel the paper could benefit from conducting more extensive experiments, but I feel the idea itself sounds interesting and it is worth publication.", "questions": ".", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The proposed approach has nearly linear runtime in theory (ableit with some assumptions) and has good performance in practice.", "weaknesses": "There is no major weakness.\n\nThe presentation is good overall, but the theory part could be improved.\nFor example, add in-line comments for algorithm description and give explanation on which steps the assumption are required.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rW2xAnaqRc", "reviewer": "Reviewer_NFiG", "summary": "Authors propose a new LSH-based approximation algorithm for attention layer, called HyperAttention. Similar to KDEformer, authors aim to bound the operator norm of the error, and this is done by two-steps: approximating the normalization matrix, and the attention matrix. Both approximations are primarily based on sortLSH, as with KDEformer, although authors generalize the analysis to be applicable to other sketching algorithms. Subquadratic runtime is proved. For empirical experiments, authors evaluate with LongBench tasks, perplexity on LongBench. Wall-clock time of HyperAttention on 131k sequence is also measured.", "questions": "Eq (10) in Section 1.2 probably mean Eq(1).\n\nIn equation (5) in Appendix A, which result are authors using to prove this?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Significance: This paper provides a new theoretically-grounded approximation algorithm for attention layers. The algorithm doesn't rely on kernel density estimation, and the analysis is more straightforward to associate the assumed properties of attention matrix with theoretical guarantees. Hence, future work will find it useful to build upon this work to adapt assumptions and derive new results. \n\nOriginality: Although the algorithm is based on the same primitive sortLSH, as prior work KDEformer did, the actual algorithm is new. Also, authors generalize to other sketching algorithms as well, which makes a stronger connection to sketching literature.\n\nClarity: The overall proof strategy is clearly described. High-level intuition and interpretation of proof parameters are explained such that readers shall follow the main logic without having to read the proof. However, intuitive explanation of algorithms is mostly reserved to appendix, which is unfortunate but understandable given page constraints.\n\nQuality: I wasn't able to closely verify proofs, but overall proof strategy of the work seems sound.", "weaknesses": "Empirical evaluations of this paper is not aligned well with prior work, which makes it difficult for readers to understand the practical usefulness of the proposed approach. For example, prior work KDEformer evaluates on BigGAN, ImageNet Classification, and Long Range Arena Benchmark, comparing against not only exact baseline but stronger baselines such as Reformer/Performer. While HyperAttention shared many similarity with KDEformer and claims to be more practical due to not using KDE, HyperAttention is not compared against KDEformer, making it difficult to see whether the difference between to actually make a practical difference. Also, end-to-end training experiment is lacking, and only forward/backward computation time was evaluated.\n\nAlthough HyperAttention shares the overall proof strategy with KDEformer, they differ significantly in assumptions they make. I believe it is potentially a strength, as HyperAttention shall make more realistic assumptions and strengthen guarantees. However, it wasn't clear how HyperAttention's assumptions relate to KDEformer. A more detailed discussion of comparison between KDEformer and HyperAttention will clarify the theoretical contribution of this paper despite sharing the overall proof strategy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "InJIAFIMGD", "reviewer": "Reviewer_xXnD", "summary": "The authors propose a method to reduce the computation complexity of transformers down from quadratic in squence length to near liner. The proposed idea proposes to perform sampling to approximate Attention. The authors claim that compared to prior works there proposed idea is practical and is able to outperform existing approaches.\n\nThe core idea is to perform sampling of the value matrix based on the norms and the \\sigma(QK^T) matrix to come up with important entries.", "questions": "See weakness section.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. I think this a resonably effective idea. The bounds provided are kind of standard and the paper is well structured.\n2. The idea of separating into multiple matrix when performing causal masking is very interesting. I think this is a very nice contribution.", "weaknesses": "Before highlighting the weaknesses I would like to point out where my criticism's are coming from. More often then not in Machine Learning approximation results, I have personally observed that the approximation often takes more time than the actual operation, often due to hardware constraints, which makes the approach moot and impractical. Stating this most my questions are coming from a point of view to discern how realistic will these speedups be in real world. I hope the authors are not too bothered by my naive questions.\n\n1. My first question is regarding evaluation. Are the speedups proposed end to end. As in the runtime reported in Figure 3(a) and 3(b) Figure 4 include the run time for generating the mask M_h . The reason I am a bit skeptical is that approximation using Algoirthm 1 does not look very GPU friendly to my naive understanding. And you have mentioned in your contributions (On Page 3) \"We assume these procedures are fast, and that after removing the heavy entries, two parameters in the resulting attention matrix are small: (1) the max column l1-norm, and (2) the ratio of row norms in the un-normalized attention matrix.\" This makes me wonder if you include the runtime. I would ideally love to see a breakdown of the time spent in different procedures as that will help new and unfamiliar readers to understand how with your novel approach bottlenecks have shifted and where the new bottlenecks are.\n\n2. For your proofs you assume the following - \"To be more precise, we assume that for any i ∈ [n] there exists some α = no(1) such that D−1A · e(i) 2 ≤ αn .\" Can you verify this assumption experimentally. I have generally found a proof to stand the test of time a bit more if the assumptions have some sort of experimental verification.\n\n3. Do the authors have some understanding if there is better way of choosing the number of layers to use for their proposed \"HyperAttention\" mechansim. It currently feels arbitary to choose these layers. It would be interesting to understanding why it works in certain cases and where it doesn't.\n\n4. The other thing I found missing is the \"m\" value (Algorithm 2). Can authors enumerate the \"m\" values used in their setup. Is it based on the bound or is it another hyper-parameter. Especially for long sequence results, since the accuracy can not be verified there it would be interesting to see the parameters which authors have used.\n\n\nIn general I found the paper easy to follow and appreciate the various figure the authors have included. I like the experiments which authors have performed and showed how on long sequence length their method is more useful. I am happy to bump up the score to a weak accept or accept based on authors reponses. Nice work !!.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)