ReAttention: Training-Free Infinite Context with Finite Attention Scope

Results on LongBench. All experiments are conducted in 32k context length, while LLaMA3-8B-8K is evaluated with Dynamic NTK.

Regarding W4, on theoretical analysis and visualization explanations, first, we thank the reviewer for acknowledging our experimental results. Regarding the latter, we have detailed descriptions in our paper, refer to subsection 4.2. Through experimental observations, we discover that token filtering without positional embeddings can more effectively locate key information in the previous context, enabling ReAttention to solve problems that full attention cannot. Regarding theoretical analysis, we are actively discussing this and will treat it as our future work.

[1] Xiao, Guangxuan, et al. "Efficient Streaming Language Models with Attention Sinks." The Twelfth International Conference on Learning Representations.

[2] Zhang, Zhenyu, et al. "H2O: Heavy-hitter oracle for efficient generative inference of large language models." Advances in Neural Information Processing Systems 36 (2023): 34661-34710.

[3] Xiao, Chaojun, et al. "Infllm: Training-free long-context extrapolation for llms with an efficient context memory." The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.

[4] Lu, Yi, et al. "LongHeads: Multi-Head Attention is Secretly a Long Context Processor." arXiv preprint arXiv:2402.10685 (2024).

[5] Oren, Matanel, et al. Transformers are multi-state RNNs. arXiv preprint arXiv:2401.06104 (2024).

[6] Li, Yuhong, et al. "SnapKV: LLM knows what you are looking for before generation." arXiv preprint arXiv:2404.14469 (2024).

[7] Cai, Zefan, et al. "PyramidKV: Dynamic kv cache compression based on pyramidal information funneling." arXiv preprint arXiv:2406.02069 (2024).

[8] Xiao, Guangxuan, et al. "DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads." arXiv preprint arXiv:2410.10819 (2024).

[9] Ge, Suyu, et al. "A little goes a long way: Efficient long context training and inference with partial contexts." arXiv preprint arXiv:2410.01485 (2024).

[10] Kim, Minsoo, et al. "InfiniPot: Infinite Context Processing on Memory-Constrained LLMs." Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing. 2024.

[11] Bai, Yushi, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508 (2023).

Dear Reviewer bDJC

We thank you for your valuable feedback and have addressed your concerns as follows.

Regarding W1, on the similarity with other works, we do have similarities with StreamingLLM [1] and H2O [2], in that we preserve the beginning and end tokens during attention and perform cache selection from the middle. However, ReAttention is fundamentally different from StreamingLLM and H2O. ReAttention is a work on extending context length, not a KV cache optimization method[1-2]. If preserving the beginning and end tokens and filtering middle tokens is concerned, there are many works including InfLLM (accepted by NeurIPS 2024) [3], LongHeads (accepted by EMNLP 2024) [4], TOVA [5], and additionally the works mentioned by the reviewer such as SnapKV (accepted by NeurIPS 2024) [6], PyramidKV [7], DuoAttention [8], which also preserve adjacent tokens and filter previous tokens.

Regarding W3, on comparisons with KV cache optimization works, first, we need to clarify that we are not a KV cache optimization work, but a work on extending model context length. Among the KV cache optimization works mentioned by the reviewer, SnapKV observes attention feature differences across different heads and clips historical tokens based on tokens at the prompt's end [6]. PyramidKV observes sparsity of attention patterns between different layers, and based on SnapKV, sets different numbers of KV cache storage in different layers to further improve computational and storage efficiency [7]. These works achieve efficient inference in a plug-and-play manner but cannot effectively achieve extrapolation. DuoAttention, building on SnapKV's approach, trains different attention heads' tendencies towards StreamingLLM and full attention, artificially classifying attention heads and applying different attention patterns during inference [8]. LongGen fine-tunes the model's intermediate layers to adapt to sparse attention patterns while maintaining full attention in the first and last layers [9]. Compared to these methods, our approach can effectively extrapolate without any training. InfiniPot is also a KV cache optimization work, defining importance metrics from future and past tokens as new token clipping criteria [10]; we similarly discussed cache selection in our paper for better context extrapolation.

Considering the reviewer's concern about ReAttention's downstream task performance, we choose the more influential SnapKV for downstream task comparisons with LongBench [11] and find that ReAttention still demonstrates competitive performance.

Dear Reviewer bDJC,

We hope we have adequately addressed your issues. We would be very grateful if you could provide feedback on our rebuttal since the deadline is approaching. If you require further clarification or have any additional concerns, please do not hesitate to contact us. We are more than willing to continue communicating with you.

Best wishes,

The Authors

Dear Reviewer xrHG

We thank you for your valuable feedback and have addressed your concerns as follows.

Regarding W1 on latency concerns, ReAttention and full attention have comparable latency, but with a critical difference, while full attention encounters Out-of-Memory (OOM) issues around 100K tokens on a single A100/800, ReAttention can continue running, as demonstrated in Figure 7. Despite additional algorithmic overhead, ReAttention proves to be more efficient practically.

Regarding W2, we compared our FlashAttention2-optimized ReAttention against standard FlashAttention2 [1] implementations, looking at the throughput during decoding on a single A800 GPU. These new results complement the TTFT (Time To First Token) metrics presented in our paper.

Results on decoding phase. The hyperparameters such as local size and global size used in our experiments, as well as the experimental hardware, remain consistent with the configurations we previously used in the paper. As in the paper, we compare against the official Hugging Face modeling of the Llama3 8B model with Flash Attention 2 enabled; to test decoding independently, we made minor adjustments to the code. The generation length is exactly 1024.

It's worth noting that we only utilized approximately 60% of a single A800's memory at 256k context length, without any offloading. The results show performance matching or exceeding FlashAttention, and we believe there's still untapped potential, particularly once inference engine(e.g., vLLM) optimizations are implemented.

Regarding W3 on storage and computational efficiency, contrary to the reviewer's statement that "The method requires sufficient memory resources where full attention is also feasible", our ReAttention uniquely enables operation beyond 100K tokens on a single A100, as shown in Figure 7. This is achieved through chunk-based reading to reduce peak memory utilization and a custom kernel for TopK Attention that minimizes intermediate computation memory read/write. These strategies are detailed in section 3.4.

Regarding W4 on generation process details, we appreciate the reviewer's concern. Indeed, our approach implements dynamic token filtering at each generation step, as detailed in Section 2.1. We've drafted detailed pseudocode showing how ReAttention works during both prefilling and decoding. We plan to include this in the appendix to make the algorithm more accessible to readers. For now, we've prepared a markdown version (due to website rendering problems for Latex) for your review – would love to get your thoughts on it.

Prefilling Phase

# Input: seq_len tokens
# Configuration:
#   - global_size tokens at start of cache
#   - local_size tokens at end of cache  
#   - chunk_size for processing

start = 0 
end = global_size + local_size
past_key_values = None

while start < seq_len:
    # Process next chunk
    input_chunk = input_ids[:, start:min(end, seq_len)]
    hidden_states = embedding(input_chunk)

    # Forward through each transformer layer
    for layer_i in range(num_layers):
        # Project to Q,K,V
        query = q_proj(hidden_states) 
        key = k_proj(hidden_states)
        value = v_proj(hidden_states)

        if layer_i == 0:
            # Apply position embeddings
            position_ids = get_position_ids(start, end)
            cos, sin = get_rope_embeddings(position_ids)
            query, key = apply_rope(query, key, cos, sin)

        # Update KV cache for this layer
        if past_key_values:
            key = concat(past_key_values[layer_i].key, key)
            value = concat(past_key_values[layer_i].value, value)
        past_key_values[layer_i] = (key, value)
        
        score = query @ (key[global_size:-local_size])^T
              idx_recall = score.topk(dim=0)
        
        key_recall = concat(past_key_values[layer_i].key[global_size:],
                           past_key_values[layer_i].key[idx_recall],
                           past_key_values[layer_i].key[:-local_size])
        value_recall = concat(past_key_values[layer_i].value[global_size:],
                           past_key_values[layer_i].value[idx_recall],
                           past_key_values[layer_i].value[:-local_size])
        
        # Compute attention
        attention_output = flash_attention(query, key_recall, value_recall)
        hidden_states = ffn(attention_output)

    # Move to next chunk    
    start = end
    end = end + chunk_size

return hidden_states, past_key_values

Decoding Phase

# Input: Last one token + KV cache from last iter
# - hidden_states of that token
# - past_key_values

for layer_i in range(num_layers):
    # Project queries just for last token
    query, key, value = q_proj(hidden_states), k_proj(hidden_states), v_proj(hidden_states)

    # First update KV cache
    key = concat(past_key_values[layer_i].key, key)
    value = concat(past_key_values[layer_i].value, value)
    past_key_values[layer_i] = (key, value)

    # Compute recall scores on middle section of cache (excluding global & local windows)
    score = query @ (key[global_size:-local_size])^T
    idx_recall = score.topk(dim=0)

    # Build attention context with global + recalled + local tokens
    key_recall = concat(
        past_key_values[layer_i].key[:global_size],          # Global window
        past_key_values[layer_i].key[idx_recall],            # Recalled tokens
        past_key_values[layer_i].key[-local_size:]           # Local window
    )
    value_recall = concat(
        past_key_values[layer_i].value[:global_size],
        past_key_values[layer_i].value[idx_recall], 
        past_key_values[layer_i].value[-local_size:]
    )

    # Compute attention only for new token
    attention_output = flash_attention(query, key_recall, value_recall)
    hidden_states = ffn(attention_output)

# Get next token prediction
logits = lm_head(hidden_states)
next_token = argmax(logits)

return next_token, past_key_values

Regarding Q1 on token eviction methods, we emphasize that we are not a KV cache optimization method. We do not perform KV cache eviction [2], but instead select critical tokens before self-attention computation, thereby achieving algorithmically infinite context. Considering the reviewer's interest in downstream task performance, we conduct a comparative analysis with H2O [2] on LongBench [3].

Results on LongBench. All experiments are conducted in 32k context length, while LLaMA3-8B-8K is evaluated with Dynamic NTK.

Best regards

Authors

[1] Dao, Tri. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning." The Twelfth International Conference on Learning Representations.

[2] Zhang, Zhenyu, et al. "H2O: Heavy-hitter oracle for efficient generative inference of large language models." Advances in Neural Information Processing Systems 36 (2023): 34661-34710.

[3] Bai, Yushi, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508 (2023).

Dear Reviewer xrHG,

We hope we have adequately addressed your issues. We would be very grateful if you could provide feedback on our rebuttal since the deadline is approaching. If you require further clarification or have any additional concerns, please do not hesitate to contact us. We are more than willing to continue communicating with you.

Best wishes,

The Authors

I appreciate the authors' detailed rebuttal and the additional insights provided regarding ReAttention's capabilities and limitations. While I acknowledge that ReAttention addresses the memory-bound problem effectively, allowing previously intractable sequence lengths to become tractable, I believe its approach shares conceptual similarities with token-dropping methods like H2O and FastGen. Specifically, selecting and attending to critical tokens is functionally equivalent to maintaining their KV cache for use during the decoding phase. This strategy effectively reduces memory overhead by discarding less important tokens, which aligns with the principles of token eviction techniques.

However, a significant drawback of ReAttention remains its latency. The two-round attention calculation introduces notable computational overhead. The additional throughput results in rebuttal does not address my concerns but somehow validate my concerns due to the trade-off between latency and throughput when increasing batch size. As shown by the throughput results, while ReAttention enables up to a 5x larger batch size compared to full attention, the throughput between the two methods is nearly equivalent. This indicates that the latency (of generating one token) of ReAttention is approximately 5x higher than that of full attention. Can authors fix the sequence length (e.g, 32k) and show the comparison to full attention about figure of latency (or throughput) v.s. batch size? The figure of latency v.s. batch size can show the efficiency of ReAttention and can be done given different fixed length (32k and 64k). Given that full attention already imposes substantial latency and computational complexity, ReAttention’s additional latency represents a considerable bottleneck that undermines its overall efficiency. Additionally, I encourage the authors to provide a comprehensive accuracy comparison between full attention to establish a more holistic understanding of the trade-offs between memory efficiency, latency, and predictive performance.

To make ReAttention more practical and widely applicable, I believe its computational complexity and latency need to be further optimized. These improvements would significantly enhance the applicability of ReAttention for long-sequence tasks.

Dear Reviewer xrHG,

We sincerely thank you for your continued engagement and thoughtful feedback.

Regarding the comparison to token dropping or eviction methods, there are differences in both approach and capability that need clarification. While token-dropping methods like H2O and FastGen permanently discard tokens, ReAttention preserves the complete context in the top-k attention. Token-dropping methods make irreversible decisions about which information to discard, creating an inherent shortage of losing potentially valuable context that cannot be recovered later. ReAttention, on the other hand, implements a dynamic selection mechanism during computation. All tokens remain available in the KV cache, and our approach determines which tokens to attend to during each specific computation step. This means the model retains access to the full context and can flexibly select different subsets of tokens in each attention head, layer, and inference step. ReAttention could see the token unseen in the previous inference step, but H2O fails to see the token dropped before. Therefore, ReAttention extends the context length to up to 1M tokens with perfect retrieval accuracy on NIAH tasks, where token-dropping methods like H2O and FastGen fail to do so.

Regarding the comparison between full attention and ReAttention in efficiency, the throughput and TTFT (Time To First Token) are two different measuring dimensions. TTFT measures performance under identical conditions (same batch size and context length). Our TTFT results demonstrate that ReAttention achieves latency on par with full attention equipped with FlashAttention2 under these controlled conditions.

Using a single A800 GPU, we test our ReAttention approach and the standard implementation of Hugging Face's **Llama 3 8B modeling with FlashAttention2. We keep all the settings, e.g. local size and global size, the same as in our paper. The batch size is exactly 1.

Throughput, on the other hand, evaluates maximum processing capacity when fully utilizing available GPU memory. ReAttention demonstrates higher efficiency by supporting larger batch sizes while maintaining comparable or better token processing rates. Besides, beyond context length 64K, full attention equipped with FlashAttention2 encounters memory limitations, while ReAttention continues working.

Using a single A800 GPU, we test our ReAttention approach and full attention with FlashAttention2 with LLaMA3-8B. We keep all the settings, e.g. local size and global size, the same as in our paper. The generation length is exactly 1024.

This performance pattern indicates not a trade-off where we sacrifice one capability for another. Rather, it represents extending context length while maintaining competitive throughput and enabling larger batch sizes.

Regarding W5 in a detailed description, we appreciate the reviewer's concern. Our approach uses dynamic token filtering at each generation step, which we covered in Section 2.1. We have prepared pseudocode for ReAttention's prefilling and decoding phases, which we plan to include in the appendix to avoid misunderstanding. Since the website's having some trouble rendering LaTeX right now, we've prepared a markdown version for you to look at.

Prefilling Phase

# Input: seq_len tokens
# Configuration:
#   - global_size tokens at start of cache
#   - local_size tokens at end of cache  
#   - chunk_size for processing

start = 0 
end = global_size + local_size
past_key_values = None

while start < seq_len:
    # Process next chunk
    input_chunk = input_ids[:, start:min(end, seq_len)]
    hidden_states = embedding(input_chunk)

    # Forward through each transformer layer
    for layer_i in range(num_layers):
        # Project to Q,K,V
        query = q_proj(hidden_states) 
        key = k_proj(hidden_states)
        value = v_proj(hidden_states)

        if layer_i == 0:
            # Apply position embeddings
            position_ids = get_position_ids(start, end)
            cos, sin = get_rope_embeddings(position_ids)
            query, key = apply_rope(query, key, cos, sin)

        # Update KV cache for this layer
        if past_key_values:
            key = concat(past_key_values[layer_i].key, key)
            value = concat(past_key_values[layer_i].value, value)
        past_key_values[layer_i] = (key, value)
        
        score = query @ (key[global_size:-local_size])^T
              idx_recall = score.topk(dim=0)
        
        key_recall = concat(past_key_values[layer_i].key[global_size:],
                           past_key_values[layer_i].key[idx_recall],
                           past_key_values[layer_i].key[:-local_size])
        value_recall = concat(past_key_values[layer_i].value[global_size:],
                           past_key_values[layer_i].value[idx_recall],
                           past_key_values[layer_i].value[:-local_size])
        
        # Compute attention
        attention_output = flash_attention(query, key_recall, value_recall)
        hidden_states = ffn(attention_output)

    # Move to next chunk    
    start = end
    end = end + chunk_size

return hidden_states, past_key_values

Decoding Phase

# Input: Last one token + KV cache from last iter
# - hidden_states of that token
# - past_key_values

for layer_i in range(num_layers):
    # Project queries just for last token
    query, key, value = q_proj(hidden_states), k_proj(hidden_states), v_proj(hidden_states)

    # First update KV cache
    key = concat(past_key_values[layer_i].key, key)
    value = concat(past_key_values[layer_i].value, value)
    past_key_values[layer_i] = (key, value)

    # Compute recall scores on middle section of cache (excluding global & local windows)
    score = query @ (key[global_size:-local_size])^T
    idx_recall = score.topk(dim=0)

    # Build attention context with global + recalled + local tokens
    key_recall = concat(
        past_key_values[layer_i].key[:global_size],          # Global window
        past_key_values[layer_i].key[idx_recall],            # Recalled tokens
        past_key_values[layer_i].key[-local_size:]           # Local window
    )
    value_recall = concat(
        past_key_values[layer_i].value[:global_size],
        past_key_values[layer_i].value[idx_recall], 
        past_key_values[layer_i].value[-local_size:]
    )

    # Compute attention only for new token
    attention_output = flash_attention(query, key_recall, value_recall)
    hidden_states = ffn(attention_output)

# Get next token prediction
logits = lm_head(hidden_states)
next_token = argmax(logits)

return next_token, past_key_values

Regarding W2 and Q1 on the NIAH comparison, we sincerely thank you for your suggestion. We have performed experiments comparing ReAttention with baseline methods and InfLLM [3] in LLaMA3-8B-8K. Our results conclusively show that ReAttention significantly outperforms FullAttention, FullAttention with DynamicNTK [5], and InfLLM on the NIAH task [6]. More results will be added to the paper later.

LLaMA3-8B-8K with DynamicNTK, scaling_factor=4

LLaMA3-8B-8K with StreamingLLM, local_size=4096, global_size=32, chunk_size=512

LLaMA3-8B-8K with InfLLM, local_size=4096, global_size=128, seg_size=128 (recommended by InfLLM), num_seg=31, chunk_size=512

LLaMA3-8B-8K with ReAttention, local_size=4096, global_size=32, seg_size=32, num_seg=127=k', chunk_size=512

Dear Reviewer fRQF

We appreciate your valuable feedback and have addressed your concerns as follows.

Regarding W1 on missing long-context evaluation results, thank you for highlighting the need for comprehensive evaluations. We have conducted additional experiments comparing ReAttention with StreamingLLM [1] on Infinitebench [2] and with InfLLM [3] on LongBench [4]. Our results demonstrate that ReAttention achieves the best performance on average.

Results on InfiniteBench

Results on LongBench

Regarding W6 on clarification on the "infinite" context, we appreciate your guidance on terminology. While our context length remains constrained by memory conditions, we have indeed achieved a training-free algorithm for infinite context length, representing a significant advancement over existing approaches [3,5].

Regarding Q3 on differences with InfLLM, thank you for encouraging a comprehensive comparison. We have already addressed distinctions in our method and related work (subsection 2.1 line 129, subsection 2.2 line 157, section 5 line 522). The key differences between our ReAttention with InfLLM can be summarized as follows, and we will include this part in the appendix.

1. First, on KV cache filtering, InfLLM uses block-based representation, pre-dividing KV caches into blocks [3]. However, Existing research shows this approach can negatively impact downstream tasks by breaking context coherence [9]. ReAttention directly selects the K cache, avoiding these issues.
2. Second, on positional embedding, InfLLM applies identical position indices for areas beyond adjacent parts, which significantly deviates from pre-training position encoding patterns, leading to cumulative errors in the KV cache, and potentially degrading downstream task performance. We've demonstrated this in subsection 3.1's Table 1 and our previous response.
3. Finally, regarding efficiency tradeoff, InfLLM emphasizes offloading to reduce memory consumption [3], while our approach focuses on efficient KV cache selection via the Triton kernel [8]. Inspired by FlashAttention [7], we extract selected intermediate KV caches without storing intermediate results and achieve computation time close to the original FullAttention with lower peak memory usage.

Regarding Q4 on figure readability, we will adjust the text size in figures to ensure clarity in the next version of our paper.

We sincerely appreciate your detailed and constructive feedback.

Best regards

Authors

[1] Xiao, Guangxuan, et al. "Efficient Streaming Language Models with Attention Sinks." The Twelfth International Conference on Learning Representations.

[2] Zhang, Xinrong, et al. "∞ Bench: Extending long context evaluation beyond 100k tokens." Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.

[3] Xiao, Chaojun, et al. "Infllm: Training-free long-context extrapolation for llms with an efficient context memory." The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.

[4] Bai, Yushi, et al. "Longbench: A bilingual, multitask benchmark for long context understanding." arXiv preprint arXiv:2308.14508 (2023).

[5] Dynamically scaled rope further increases performance of long context llama with zero fine-tuning, July 2023a. https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/.

[6] Needle in a haystack - pressure testing llms. https://github.com/gkamradt/LLMTest_NeedleInAHaystack, 2023

[7] Dao, Tri. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning." The Twelfth International Conference on Learning Representations.

[8] Tillet, Philippe, Hsiang-Tsung Kung, and David Cox. "Triton: an intermediate language and compiler for tiled neural network computations." Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages. 2019.

[9] Luo, Kun, et al. "BGE Landmark Embedding: A Chunking-Free Embedding Method For Retrieval Augmented Long-Context Large Language Models." arXiv preprint arXiv:2402.11573 (2024).

Regarding Q2 on Figure 7, we confirm that we use full attention with FlashAttention2 [7] as our baseline in the referenced figure.

Regarding W3, we've now run a new set of experiments, looking at throughput during decoding on a single A800 GPU, a complement to the TTFT (Time To First Token) metrics presented in our paper. Using a single A800 GPU, we test our ReAttention approach and the standard implementation of Hugging Face's Llama 3 8B modeling with FlashAttention2. We keep all the settings, e.g. local size and global size, exactly the same as in our paper. The generation length is exactly 1024. Here's what we have found:

Not only can we handle much longer sequences, but we're also faster at shorter lengths where both methods work. We're seeing about 7-20% better throughput at 32k and 64k contexts, while also managing larger batch sizes. Beyond that, the baseline runs out of memory, but our method keeps going strong while only using about 60% of the GPU memory(49.3GB at 256k).

Regarding W4 on memory overhead tradeoff, we appreciate your insightful feedback on memory considerations. We are committed to continuous optimization of our Triton operator [8]. As a training-free context extrapolation method primarily, ReAttention focuses on demonstrating significant advantages in long-context scenarios. Specifically, our approach requires less memory compared to full attention in contexts exceeding 100k tokens and enables extrapolation of LLaMA3.1-8B to 400k tokens on a single GPU.

Dear Reviewer fRQF,

We hope we have adequately addressed your issues. We would be very grateful if you could provide feedback on our rebuttal since the deadline is approaching. If you require further clarification or have any additional concerns, please do not hesitate to contact us. We are more than willing to continue communicating with you.

Best wishes,

The Authors

Thanks to the authors for the comprehensive response, it has addressed many of my concerns (missing benchmarks and decoding throughput comparisons). Please include these results in the final version of the paper. I have raised my score to a 6.

Comments:

• From the LongBench table, it seems like the base LLaMA model achieves better performance than ReAttention. Why do you think this is the case?
• I still don't agree with the use of the word "infinite" in the title. In practice, the method cannot support infinite sequence lengths due to memory constraints, so in what sense has "infinite context length" been achieved?

Dear Reviewer iiPf

We appreciate your valuable feedback and recognition of our work. We have addressed your concerns as follows.

Regarding W1 on writing issues and typos, we sincerely appreciate your meticulous review of our manuscript. We will carefully address and correct the spelling errors in the subsequent version of our paper.

Regarding W2 on more related work, we have comprehensively expanded our discussion on KV cache optimization, especially on token eviction methods like H2O [1] and SnapKV [2]. The evolution of token eviction techniques begins with DCP [3], Scissorhands [4], and H2O [1]. These approaches demonstrate that retaining only a small subset of tokens with high attention scores could maintain near-equivalent performance while significantly reducing memory and computational overhead. Then, StreamingLLM [5] advances this concept by identifying that critical tokens always occur in the beginning and end part of the attention scope. By preserving only corresponding tokens, StreamingLLM enables stable model outputs after infinite input lengths, dramatically reducing computational and memory costs. However, this approach substantially compromised long-context performance [2,6]. Subsequent researches explore mitigation strategies, such as head-adaptive methods like FastGen [7] and DuoAttention [6]. Specifically, SnapKV [2] introduced a more nuanced approach to pruning previous KV caches by querying the above context with later input segments, substantially enhancing long-context performance for token eviction techniques. Building upon these foundations, researchers such as PyramidKV [8], PyramidInfer [9], Gemfilter [10], and LongGen [11] take layer-wise optimization, maintaining inference effectiveness while improving memory efficiency further. Besides, works like VATP [12] and InfiniPot [13] expand token pruning metrics beyond traditional QK dot product scoring. Our ReAttention draws inspiration from these token eviction methods, specifically applying the principle of selecting KV cache to context extrapolation, enabling comprehensive context perception throughout the inference process while computationally addressing each attention calculation with a finite context, thereby achieving algorithmically infinite context length.

Regarding Q1 on performance on position-sensitive tasks, given the scarcity of code completion and mathematical reasoning tasks in long-context scenarios, we are actively seeking and adapting appropriate evaluations such as RepoQA [14] and Infinitebench.MathFind [15]. For the existing evaluation results, LongBench [16] includes code completion tasks like LCC [17] and RepoBench-p [18]. We have compared the performance of ReAttention and full attention across different models on LongBench's code tasks. The results are shown as follows. Our analysis reveals that ReAttention achieves comparable overall performance with full attention and outperforms StreamingLLM [5]. We are committed to further exploring and validating our method's effectiveness on position-sensitive tasks.

Results on code task in LongBench. All models are evaluated in 32k context length, with LLaMA3-8B-8K evaluated with Dynamic NTK.

Thank you for your valuable feedback and suggestions.

Best regards

Authors

[1] Zhang, Zhenyu, et al. "H2o: Heavy-hitter oracle for efficient generative inference of large language models." Advances in Neural Information Processing Systems 36 (2023): 34661-34710.

[2] Li, Yuhong, et al. "Snapkv: Llm knows what you are looking for before generation." arXiv preprint arXiv:2404.14469 (2024).

[3] Anagnostidis, Sotiris, et al. Dynamic context pruning for efficient and interpretable autoregressive transformers. Advances in Neural Information Processing Systems 36 (2024).

[4] Liu, Zichang, et al. "Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time." Advances in Neural Information Processing Systems 36 (2024).

[5] Xiao, Guangxuan, et al. "Efficient Streaming Language Models with Attention Sinks." The Twelfth International Conference on Learning Representations.

[6] Xiao, Guangxuan, et al. "DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads." arXiv preprint arXiv:2410.10819 (2024).

[7] Ge, Suyu, et al. "Model tells you what to discard: Adaptive kv cache compression for llms." arXiv preprint arXiv:2310.01801 (2023).

[8] Cai, Zefan, et al. "Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling." arXiv preprint arXiv:2406.02069 (2024).

[9] Yang, Dongjie, et al. "PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference." arXiv preprint arXiv:2405.12532 (2024).

[10] Shi, Zhenmei, et al. "Discovering the gems in early layers: Accelerating long-context llms with 1000x input token reduction." arXiv preprint arXiv:2409.17422 (2024).

[11] Ge, Suyu, et al. "A little goes a long way: Efficient long context training and inference with partial contexts." arXiv preprint arXiv:2410.01485 (2024).

[12] Guo, Zhiyu, Hidetaka Kamigaito, and Taro Watanabe. "Attention Score is not All You Need for Token Importance Indicator in KV Cache Reduction: Value Also Matters." arXiv preprint arXiv:2406.12335 (2024).

[13] Kim, Minsoo, et al. "InfiniPot: Infinite Context Processing on Memory-Constrained LLMs." arXiv preprint arXiv:2410.01518 (2024).

[14] Liu, Jiawei, et al. RepoQA: Evaluating Long Context Code Understanding. arXiv preprint arXiv:2406.06025 (2024).

[15] Zhang, Xinrong, et al. ∞ Bench: Extending long context evaluation beyond 100k tokens. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.

[16] Bai, Yushi, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508 (2023).

[17] Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. 2023. Longcoder: A long-range pretrained language model for code completion. arXiv preprint arXiv:2306.14893.

[18] Tianyang Liu, Canwen Xu, and Julian McAuley. 2023b. Repobench: Benchmarking repository-level code auto-completion systems. arXiv preprint arXiv:2306.03091.