HyperAttention: Long-context Attention in Near-Linear Time

Official Comment by Authors

> 1. My first question is regarding evaluation. ... and where the new bottlenecks are. - All runtimes in our results are end-to-end. Although in the right-most figure in Figure 1 the sparsified attention by ${\mathbf {M}}^{\mathcal{H}}$ seems to be computationally inefficient and not parallelizable, we follow the operations in the middle figure in practice, which does not compute the mask ${\mathbf M}^{\mathcal{H}}$ explicitly. More specifically, rows in $\mathbf{Q}$ and $\mathbf{K}$ are sorted using Hamming LSH, then they are partitioned into “$b$” (i.e., the block size in Figure 1) groups. Finally, the i-th set of rows in $\mathbf{Q}$ is multiplied by the corresponding set in $\mathbf{K}$, resulting in a block-diagonal approximation of $\mathbf{A}_{P_Q, P_K}$ . The required operations involve permuting $n$ rows, reshaping tensors, and small matrix multiplications. Since every batch, head and block has the same configurations, the implementation can be parallelized using GPUs. We have added more details in the updated submission. - The two parameters $\alpha$ and $\kappa$ are used to analyze the runtime of our algorithm but they are not required to compute in our practical algorithm.

Official Comment by Authors

> 2. For your proofs you assume the following ... the assumptions have some sort of experimental verification. - We additionally conducted experiments to verify our assumption. For the attention with non-causal masking, we use the T2T-ViT model and take query, key and values from its first attention layer on the ImageNet test data set. Then, we compute $\alpha$ to be the maximum of the squared $\ell_2$-norms of the columns in $\mathbf{D}^{-1} \mathbf{A}$. The sequence length of T2T-ViT is $n=3136$ and the average value of $\alpha$ is observed to be 8.1801, which can be sublinear in $n$. - To further investigate the dependence on $n$, we choose the ChatGLM and longbench narrative-qa dataset used in the submission and change the sequence length n from 1k to 9k. As a result, we observe that the value of $\alpha$ increases sub-linearly in $n$, as below. This supports the claim that our assumption holds in practice. We have added these results in our updated submission. | $n$ |$\alpha / n$| | :------------- |:-------------:| | 1024 | 0.09170 | | 2048 | 0.06379 | | 3072 | 0.05395 | | 4096 | 0.04817 | | 5120 | 0.04559 | | 6144 | 0.04128 | | 7168 | 0.03918 | | 8192 | 0.03694 | | 9216 | 0.03601 |

Official Comment by Authors

> 3. Do the authors have some understanding ... It would be interesting to understanding why it works in certain cases and where it doesn't. - In our experiments, we choose to replace the final L-layers where L changes from 1 to the total number of layers in the given pretrained Transformers, e.g., [1]. We motivate this setting from fine-tuning methods where the last few layers are updated. We believe that finding better configurations of layers to approximate is a very interesting future problem. [1] Chen, Boqi, Fandi Yi, and Dániel Varró. "Prompting or Fine-tuning? A Comparative Study of Large Language Models for Taxonomy Construction." https://arxiv.org/pdf/2309.01715.pdf.

Discussion ends in 2 days

Dear reviewer, Since you indicated that you were > happy to bump up the score to a weak accept or accept based on authors reponses. we wanted to inquire whether you were satisfied with our responses. If not, please let us know soon so that we could provide more details or address your other questions/concerns before the discussion period ends soon. Thank you again for your time!

Official Comment by Authors

> Empirical evaluations of this paper is not aligned well with prior work, .... A more detailed discussion of comparison between KDEformer and HyperAttention will clarify the theoretical contribution of this paper despite sharing the overall proof strategy. - We thank the reviewer for their insightful comments. As the reviewer points out, HyperAttention is much simpler and more streamlined than the KDEformer. Particularly, the KDEformer requires a kernel density estimation (KDE) algorithm which is impractical and fairly hard to implement in general, in particular, with GPUs. We streamline such a complicated KDE to a simple approximation via sortLSH and uniform sampling. The assumptions we use are for guaranteeing that our uniform sampling approach can well-approximate the kernel density. We empirically observe that these assumptions hold in practice (please refer to our responses to reviewer **xXnD**). - Additionally, KDEformer does not support **causal masking** and this might be the reason why it mostly works on vision transformers, which does not require causal masking. In this work, we propose a novel method (HyperAttention) to support causal masking, and it allows us to apply it to causal language models. Furthermore, HyperAttention can be implemented in a modular way, thereby allowing us to utilize other accelerated self-attention algorithms such as FlashAttention. This makes the gap in runtimes between HyperAttention and KDEformer even larger. - To see the practical differences between HyperAttention and KDEformer, we conduct ImageNet classification experiments with the T2T-ViT model. We apply HyperAttention and KDEformer to the first attention layer because it is the main computational bottleneck. The sequence length is 3136 and we set all hyperparameters in both methods so that they have the same memory and time complexity. The per-image runtime of KDEformer in the attention layer was 13.175 ms while HyperAttention shows 8.634 ms runtime (x1.53 faster). The test accuracies of KDEformer and HyperAttention were 81.11% and 81.14%, respectively, which are the same within a standard deviation. > Eq (10) in Section 1.2 probably mean Eq(1). - We appreciate your bringing attention to this typo. Equation (1) is indeed correct and we have corrected it in our updated submission. > In equation (5) in Appendix A, which result are authors using to prove this? - Equation (5) shows an upper bound on the number of rows for which the row sum is greater than $\tau$. To derive this bound, we make use of the fact that tau is the maximum of the row sums of a random subset of rows, denoted by T. Recall that $|T|=m$. Now note that the probability that the size of the set $S_0$ is greater than $t$ for some $t$, equals $\Pr[|S_0| \ge t] \le (1-t/n)^{m}$. Plugging $m=\Omega(\kappa^7 \alpha^2 \varepsilon^{-6} \log n)$ and $t=O(\kappa^{-7} \alpha^{-2} \varepsilon^6 n)$ gives Equation (5) with probability $1-1/\text{poly}(n)$. We have updated this in our submission file.

Official Comment by Authors

We thank the reviewer for their detailed and insightful comments. We have modified the submission to include in-line comments describing the purpose of each quantity in algorithm 2, as well as an explanation before the algorithm of where the assumptions of Lemma 1 are used; in particular, the assumptions are used in the proofs to show that the variance of the estimator is small, therefore allowing for small sample complexity.

Thanks

Thanks for sharing this work. Here we give a comparison to the paper [1] - That paper requires not only estimating the row norms of $\mathbf{V}$, but also estimating the column norms of $\mathbf{A}$, where $\mathbf{A} = \exp(\mathbf{QK}^T)$. One of our key observations is that for the approximate matrix product guarantee, i.e., for preserving $\mathbf{D}^{-1} \mathbf{A SS^T V}$ up to the same desired additive error as in [1], one needs only to sample from the row norms of $\mathbf{V}$ and *does not need column norm information* of $\mathbf{D^{-1} A}$ (we note that the failure probability dependence is worse, but this is not a crucial parameter since arbitrarily large constant failure probability is already achievable just from sampling according to the row norms of $\mathbf{V}$). Critically, we achieve the same error bound using only row norm sampling from $\mathbf{V}$. This observation was missed in [1] and is essential since computing the column norms of $\mathbf{A}$ is expensive. [1] instead tries to estimate each of the column norms of $\mathbf{A}$ by uniform sampling, which in general is not possible, even under the assumptions we make in our paper, as it is easy to miss large entries. We note that our assumptions are used instead only to estimate the row norms of $\mathbf{A}$, as we describe next. - Another main issue with [1] is that they need to estimate the row norms of $\mathbf{A}$ in order to compute $\mathbf{D^{-1}}$ to normalize the rows. [1] gives a heuristic for doing this based on interpolating the average value of the unread entries in a row based on the corresponding entries in the row in the columns of $\mathbf{A}$ that they sample, but there are examples where this heuristic fails, since heavy entries may be missed by this interpolation. In short, [1] provides no end-to-end theorems unlike our Theorem 1 and Corollaries 1 and 2, which we consider to be the main theoretical contributions of our work. - We would like to stress that the difference between our work and [1] is not only in the theoretical justification, but the algorithms themselves are also different in fundamental ways: (1) for finding the subset of columns of $\mathbf{D^{-1} A}$ to use, we use only row-norm sampling from $\mathbf{V}$ whereas [1] also uses column norm sampling from $\mathbf{A}$, and (2) although our work and [1] use uniform sampling of entries in $\mathbf{A}$, [1] uses this for the purpose of finding the subset of columns of $\mathbf{A}$ to use (and has no end-to-end guarantees for this since uniform sampling of columns is not guaranteed to estimate the column norms of A without strong assumptions), whereas we use uniform sampling only for the purpose of estimating the *row norms* of $\mathbf{A}$, so that we can create the normalization matrix $\mathbf{D}^{-1}$ (and [1] only provides a heuristic for estimating entries of $\mathbf{D}^{-1}$ and does not argue that they accurately estimate such entries).