Radar: Fast Long-Context Decoding for Any Transformer

In my mind, the paper would be stronger if there was a clearer link between Theorem 2 and the experiments. Namely, unless I've missed something, the authors don't actually measure whether their algorithm retrieves the correct top-k index set in their experiments. It would be useful to do so, and to comment on whether by those results the bound in Theorem 2 is practically useful or not. It would also be useful to know something about the distribution of attention scores, which would affect what value of $k$ is reasonable. More precise (and in some sense more mechanistic) quantifications could help the reader get a better sense of why your method works.

We appreciate the reviewer's insightful question. Theorem 2 is a worst-case analysis that provides a theoretical framework for understanding dot product approximation via random projection. The bound offers insights into factors impacting attention approximation, which could guide future work in this area to derive tighter bounds for specific practical settings. Empirically, we found that the approximation works well in many practical cases (34% hit@1 and 63% hit@3), which performed better than StreamingLLM (19% / 47%) and random (10% / 30%) on the first 100 tokens (after the first one) of the PG-19 dataset using the Llama 3.1 model. Visualizations of how the approximation works also confirm this. We have added this new analysis to Appendix D.

Do you have any intuition for why the H20 and SnapKV methods perform so poorly when applied to the Mistral PG19 model (Appendix C)? The difference is quite striking, so it would be useful to have at least a hypothesis for why these methods do not generalize across architectures. This could also provide useful context for why the proposed method successfully generalizes.

Thanks for asking this. We do have a hypothesis of why Radar generalizes better than H2O and SnapKV. Specifically, both H2O and SnapKV heuristically use the accumulated attention scores as the indicator of the importance. However, the heuristics may lead to evicting the current unimportant tokens which will be very useful later for new tokens, especially when each head is responsible for multiple query spaces (i.e., in group-query attention used in Llama 3 and Mistral models). Indeed, we observed that both methods tend to perform worse on these models. By contrast, our Radar is a proven approximation of the vanilla attention method. The error bound is well-established in our Theorem 2 regardless of the model architectures and weights, making it more versatile in applications. We have added this discussion in Appendix C.

There appears to be a typo on Line 837, as the square brackets enclose the equality.

Yes, you are correct. Thank you for identifying this typo. We have fixed it in the PDF.

It might be useful to state Lemma 1 in terms of concentration to the expectation at large $n$, as that could give a bound on error in a practical setting.

Yes, we agree that stating it in terms of concentration will be more informative. However, since building the concentration inequality for a single approximation does not directly justify our method, we only stated its expectation here to make it easier to understand. Nevertheless, we expect that one can easily build such an inequality following the same procedure as our Lemma 8 based on the provided properties in our Lemma 6.

Given that the claimed time-complexity improvement is in the exponent of a power law, it would be useful to use logarithmic scales for time in the figures, so that one can see how tight the bound is. It could be semantically useful to use a different color scheme for Figure 4, to distinguish the fact that here different dimensions are being used rather than different algorithms.

We appreciate these detailed suggestions on writing. We have tried both the log-scale time (this anonymous link) and other color schemes (this anonymous link). Since there are some overlapping curves in both figures that might hinder readability, we temporarily maintained our original figures. However, we are grateful for this suggestion and will explore other designs in the next revision.

Overall, we would like to thank the reviewer for recognizing our contributions and outlining the highly detailed comments. Based on the suggestion, we have improved our paper by visualizing the theory in Theorem 2. We believe that our revised paper further resolves concerns. We are more than happy to discuss further if the reviewer has further suggestions.

The paper lacks a comparison with several relevant works focused on attention acceleration. Notably, the following papers explore sparse, low-rank, and combinations of these approaches with near-linear runtimes for long context, and it remains unclear how the proposed method measures up against them. Additionally, this method does not tackle the memory concerns regarding storing key-value embeddings in the KV cache; it requires retaining all key and value embeddings in the cache during the generation phase. How does your method compare against the aforementioned works in experimental evaluations?

Thank you for suggesting the papers. In the initial submission, we mainly followed the baselines and experimental settings in recent papers StreamingLLM (at ICLR 2024) and SnapKV (at NeurIPS 2024). Based on the comment, we now have included all of the suggested papers in the related work section. However, we would like to point out that many of them are not directly comparable as they use different settings. Specifically, our method is a post-training acceleration method without using any data. Among 5 suggested papers, there are 3 methods (Scatterbrain, KDEformer, and Routing Transformers) that require training. As discussed in lines 504-512, we cannot directly use these models as baselines because they have different architectures trained with their unique data. Our method is specifically proposed to tackle the inefficiency of the vanilla Transformer attention for inference after training.

The remaining 2 methods (HyperAttention and the recent SubGen) are more related. However, their experiments were conducted on ChatGLM and LongChat, whereas our method and all our baselines use Llama and Mistral models. This creates difficulty in conducting fair comparisons given the limited amount of time during the discussion. Nevertheless, we are able to replicate SubGen (the newer method from the same group as HyperAttention) on the already included LongBench dataset using the Llama models. The preliminary results are as follows

Llama 2 with maximum 1024+32 tokens used in each step

The results show that our method performs consistently better than the baseline (including SubGen). We believe that the results strongly suggest the capability of our method.

However, we acknowledge that these results should not discard the values of previous work as their main target is memory efficiency, whereas Radar focuses on time efficiency and accuracy. We have added this discussion in Appendix E and will be committed to updating the discussion when we have complete results of the suggested baselines.

Thank you for your detailed and positive evaluation of our work. We greatly appreciate your recognition of the strengths of our approach! Encouraged by the positive comments, we have updated the PDF which further enhances the strength of this paper. Please feel free to let us know should you have any further questions.

The analysis of time complexity needs more detail, particularly regarding lines 8 to 11 of Algorithm 1. This dynamic restructuring step may take more time than $O(t)$. Based on my understanding, this section could have a total time complexity of $\sum_{c=1}^{\sqrt{t}} c^2 = O(t^{3/2})$. Thus, the time required for calculating the importance score of a single segment is not constant.

In line 11, the time complexity is $O(t)$ because we have $c$ segments with each containing $c$ tokens. Since each segment representation $\bar\phi$ takes $O(c)$ (according to Eqn. 5) to construct, line 11 takes at most $\sum_{i=1}^c O(c) = O(c^2)$ time. This is equivalent to $O(t)$ because $c=\sqrt{t}$ in line 10.

The overall restructuring steps (considering the outer loop containing lines 8-11) is $O(t^{3/2})$ time, as we will restructure the hierarchical approximation by at most $O(\sqrt{t})$ times (because of the condition in line 8). When amortized to all $t$ steps, each step will take $O(\sqrt{t})$ on average, matching the querying complexity. We analyzed this in lines 205-207.

We hope this could better illustrate the time complexity analyses.

We sincerely appreciate your efforts in verifying our calculation and capturing the technical details. By double-checking the relevant parts, we are confident in our analyses of the time complexities. We hope our response answers your concerns. We are also happy to answer any further questions.

Thank you for your reply. I understand your perspective that each segment representation takes $O(c)$. However, I believe Line 11 calculates $c$ segments for each $c$, with each segment containing $c$ tokens. Therefore, the computation cost for Line 11 at a specific $c$ is $O(c) \times c = O(c^2)$.

Let’s break this down clearly:

For a given $c$, there are $c$ segments to compute. Each segment contains $c$ tokens, requiring $O(c)$ to process. Hence, the total cost for that iteration is $c \times O(c) = O(c^2)$.

When summing over all $c$ in the range $[1, \sqrt{t}]$, the total time complexity is: $\sum_{c=1}^{\sqrt{t}} c^2 = O(t^{3/2}),$ as the sum of squares formula confirms.

If you disagree with $O(c^2)$ for Line 11, could you clarify where you think this step is optimized or reduced? I’m open to revisiting the logic if there’s something I missed.

Thank you for your response. That said, I believe the paper still faces some challenges in terms of clarity and organization, which may impact its overall readability and impact. For instance, the $O(t)$ complexity mentioned in line 11 is somewhat unclear and could benefit from further elaboration.
Nevertheless, I acknowledge the significant effort put into addressing the concerns raised, and I have adjusted my score accordingly to reflect this progress.
Lastly, I have one final question: how do we select $k$ indices into $S$ from $c$ segments? Are we randomly selecting one index from the top $k$ segments with the highest average attention scores? How to ensure that the selected index in the segment related to high attention score with query?

There seems to be an assumption that segments should be made of contiguous tokens. This might be a safe assumption, but should be discussed. I can imagine an 'adversarial' setting where the ideal attention matrix would have been centered on a subset of tokens spaced far apart (i.e., each in a different segment). The theoretical analysis could perhaps have focused on this as well (how likely it is for the top-k segments to include the top-k tokens).

Thank you for the insightful comment. We have added a discussion in Appendix E. Essentially, our reasoning for this choice is as follows. First, as you mentioned, this is naturally a safe assumption because the important tokens are sparse in a sequence (as shown in our new Figures and previous work like StreamingLLM). In addition, by selecting the surrounding tokens around the important ones, we could provide the context for the model to correctly identify the semantics for these tokens. Our design could naturally achieve this goal, while other methods like SnapKV need to explicitly apply smoothing techniques to encourage the “surrounding selection”. Nevertheless, we consider that there could exist better segmenting mechanisms for Radar by developing an adaptive method. We hope this paper serves as the foundation and leaves the exploration of this direction to future work.

It would have been nice to include some examples of the true and approximated attention maps (as heatmap figures, and perhaps also as text examples). We only see performance numbers at the final task, but not directly how well attention is being approximated.

Thank you for the suggestion. We have added the headmap figures in Appendix D based on the suggestion. We believe this further enhances the strength of our paper.

In addition to speaking to the weaknesses listed above, it would be good if the authors could discuss whether they think this kind of hierarchical approximation to attention might work even better if it were included during training, not just as a post-training step.

Thank you for the question. Due to the time constraint, we could not complete relevant experiments. However, we have discussed this in Appendix E, stating “We expect training with this kind of hierarchical approximation could yield better results.” and “We leave this to future exploration.”

Thank you for reviewing our paper. We highly appreciate the comments for acknowledging our contributions and novelty. We have also included the related discussions in our paper. We believe these discussions further strengthen our submission. Please do not hesitate to let us know for more questions.