Paper: HyperAttention: Long-context Attention in Near-Linear Time (ICLR-2024)

Reviewer comment (Evaluation):
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.

Author response:
- 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.

Reviewer comment (Experiments):
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.

Author response:
- 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  |

Reviewer comment (Experiments):
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.

Author response:
- 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.

Reviewer comment (Reproducibility):
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.

Author response:
- In our experiments, we consider the value “$m$” to be a hyper-parameter, which corresponds to the number of columns in $\mathbf{D}^{-1} \mathbf{A}$ selected uniformly at random. Motivating our analysis that m can be sublinear in “$n$”, we fix this value to be 256 for all the experiments. We observe that a larger value of m shows a better performance on LLMs (e.g., perplexity) but it makes our algorithm slower.