Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers

Dear reviewer cHLB,

We sincerely appreciate your efforts in reviewing our manuscript. We respond to each comment in the following content.

---

[W1, Q1] Comparison to H2O

As also noted by Reviewer NqoB, the key novelty of our work lies in proposing the compress-gather-recompute pipeline, which enables effective and efficient, training-free context extension for large foundation models. Below, we clarify the main distinctions from H2O:

---

Key Difference 1: REFORM overcomes the context forgetting problem inherent in H2O

Recurrent approaches like H2O compress input segments incrementally, which often leads to context forgetting, particularly when handling longer inputs. This issue arises because information from earlier segments may be lost during successive compressions, harming performance on tasks that require long-range context.

To address this, REFORM introduces a compress-gather-recompute pipeline. The 'compress' stage encodes long input into token-level embeddings. Then, the 'gather' stage identifies important tokens for generating a response. Finally, the 'recompute' stage reconstructs the KV cache based on selected tokens. This design allows the model to access relevant information regardless of its position in the input, thus overcoming the forgetting issue.

Thanks to this fundamental difference in the method's capability, REFORM significantly outperforms H2O across diverse models and benchmarks.

Table 1. Performance comparison with H2O (Llama-3.1-8B-Instruct).

Table 2. Performance comparison with H2O (Mistral-Nemo-Instruct-2407).

Table 3. Performance comparison with H2O (Qwen2.5-Coder-1.5B/7B-Instruct).

Table 4. Performance comparison with H2O on MM-NIAH (Pixtral-12B-2409).

---

Key Difference 2: Different usage of recurrent KV cache compression

While H2O uses recurrent compression directly to generate new tokens from the compressed KV cache, REFORM leverages the compression differently. We use it to generate token-level embeddings, which are later employed to select relevant tokens for recomputation. The suggested approach is orthogonal to the specific compression mechanism (e.g., H2O, TOVA, StreamingLLM), and empirically performs well with compression mechanisms other than H2O, as shown in Table 6(a) of the main text.

Such a different usage enables early-exit optimization, as embeddings only require QKV states from selected layers (see L144–L149). H2O, by contrast, requires KV caches from all layers to be used for generation. Therefore, the early-exit optimization cannot be used for H2O.

Thanks to the early-exit optimizations, REFORM is more efficient than H2O, both in terms of computation and memory. Specifically, our efficiency analysis in Table 6(b) from the main text suggests that REFORM reduces the inference time by 34.1% and peak memory usage by 7.53% compared to H2O.

---

We believe these technical differences and empirical improvements clearly demonstrate that REFORM offers substantial advancements over H2O. Additionally, the proposed pipeline outperforms newer baselines like InfiniPot and InfLLM in both accuracy and efficiency, as demonstrated throughout the main text.

---

[L1] Deeper discussion on limitations.

Thanks for the constructive suggestions regarding more deep and insightful limitation discussions. Here, we provide some additional limitations of our work.

One limitation is that the effectiveness of REFORM depends in part on the gather stage’s ability to identify informative tokens. While our results across diverse benchmarks demonstrate that this step is generally reliable, edge cases may exist where the model's selection could be further optimized.

Another limitation is the redundant attention score computation in the current implementation. Our current implementation recomputes attention scores using matrix multiplication to compute the eviction scores, as Flash Attention does not output attention weights. Integrating eviction score computation into the Flash Attention kernel could further improve REFORM’s efficiency.

Dear reviewer fog1,

We sincerely appreciate your efforts in reviewing our manuscript. We respond to each comment in the following content.

---

[W1, Q1] Results with Larger Models

As requested, we have included additional results using the Qwen2.5-32B-Instruct model on the RULER and BABILong benchmarks. As shown in Table 1, REFORM consistently outperforms both H2O and InfiniPot across all evaluation settings, demonstrating the scalability and robustness of our method on larger models.

Table 1. Performance of Qwen2.5-32B-Instruct.

---

[Q2] Analysis of Eviction Strategies

We appreciate the request for further analysis of eviction strategies. Table 2 compares REFORM’s performance using alternative eviction methods (StreamingLLM and TOVA), in addition to our primary method (H2O). While H2O yields the highest accuracy, alternative methods still significantly outperform baseline approaches. These results underscore the flexibility of REFORM with respect to the eviction strategy. (Refer also to Table 6(a) and Lines 318–322 of the manuscript.)

Table 2. Performance with different eviction methods (Mistral-Nemo-Instruct-2407)

---

[Q3] Similarity Search and MNR Results for Llama Models

In Table 3, we showcase the similarity search results for Llama-3.1-8B-Instruct. As with Mistral-Nemo (Table 1 and Table 7 in the manuscript), we observe that cosine similarity using either Q, K, or V embeddings often yields the best performance, outperforming alternatives such as attention scores or cosine similarity using the hidden states.

Although we cannot share the MNR score distribution plots due to NeurIPS policy (image sharing restrictions), the observed trends are consistent with our previously reported results. Specifically, the top-performing heads in the pattern matching task tend to cluster at the early and the final layers, while top-performing heads for the multi-hop QA task tend to cluster in the middle layers.

These findings indicate that attention head significance patterns are consistent across architectures.

Table 3. Similarity search results for (top) pattern matching task and (bottom) multi-hop QA task, measured with Llama-3.1-8B-Instruct.

---

[Q4] Effect of Aggregation Window Size

In Table 4, we conducted a more detailed ablation study on the aggregation window size. Performance improves with increasing window size up to 129, after which it begins to saturate or slightly degrade. This suggests that our original selection of 129 is a balanced and effective kernel size.

Table 4. Impact of Aggregation Window Size (Mistral-Nemo-Instruct-2407)

Dear reviewer gpyK,

We sincerely appreciate your efforts in reviewing our manuscript. We respond to each comment in the following content.

---

[W1] Clarification on Research Objective

We would like to emphasize that our main research goal is to propose an efficient, training-free context window extension method for large foundation models (see L18–L19, L28, L87, L226–L231). We will revise the manuscript to more clearly state this objective.

Hence, the suggested baselines, i.e., sparse attention variants such as DuoAttention, MInference, and MoA address a different problem, as they aim to improve efficiency within the model's native context window (e.g., up to 128k tokens for LLaMA-3.1-8B-Instruct). Due to the reason, we did not compare them in the paper as they cannot be used for inputs that are longer than the model's context limit. Nevertheless, they are orthogonal to our approach and could potentially be combined with REFORM to further enhance efficiency.

While CEPE is capable of extending the context window, it requires training the bidirectional encoder and cross-attention layers, deviating from our training-free objective. Furthermore, REFORM will likely outperform CEPE as it is one of the early efforts on context extension, first released in February 2024. Notably, Figure 4 of the CEPE paper shows that it struggles on simple needle-in-a-haystack evaluations, whereas REFORM maintains 100% accuracy even at 1M tokens.

Finally, we remark that SnapKV, while originally designed for inference within the native context window, can be adapted to longer inputs by applying it iteratively over input chunks (as suggested by Reviewer NqoB). As shown in Table 1, REFORM consistently outperforms such a variant of SnapKV.

We believe that these discussions further clarify REFORM’s positioning in the long-context inference space, and will add them to the revised manuscript.

Table 1. Performance comparison with recurrent SnapKV variant (Llama-3.1-8B-Instruct).

---

[W1, Q2, Q3] Comparison with Baselines Within the Context Limit

Although our focus is on extending beyond the context window, we agree that evaluating REFORM within the context limit provides valuable insight, especially given that there are several open-source models that natively support 1M tokens. Table 2 compares REFORM with full attention, DuoAttention, and Minference using the Llama-3-8B-Instruct-Gradient-1048k model (1M-token support).

Somewhat surprisingly, REFORM consistently outperforms sparse baselines and often surpasses full attention, especially at longer context lengths. This is likely due to the gather-and-recompute stage, which filters out irrelevant information, making the task easier for the model to solve.

In Table 3, we further compared the performance of REFORM and full attention on Llama-3.1-8B-Instruct and Mistral-Nemo-Instruct-2407. We observed similar trends, where REFORM matches or outperforms full attention, and the gap becomes larger at longer contexts. While our main goal is in context extension beyond the model's context limits, these results suggest that REFORM can be useful even within the native context window by selectively handling only the relevant inforamtion.

Table 2. Performance within the context limit (Llama-3-8B-Instruct-Gradient-1048k). Default configurations from the official repositories were used for Duo Attention and MInference. RULER-Multi corresponds the average accuracy of the multikey, multivalue, and multiquery tasks.

Table 3. Performance comparison with full attention, using (Top) Llama-3.1-8B-Instruct and (Bottom) Mistral-Nemo-Instruct-2407. Full attention values are taken from the official leaderboards (BABILong performance for full-attention Nemo is not available).

---

[W2-1] Implementation Details and Flash Attention Compatibility

All efficiency experiments were done using Flash Attention 2 (see L539), and we will make this clearer in the captions of relevant tables. Since Flash Attention does not provide attention weights, we re-computed them using PyTorch matmul over the Q/K states. To minimize overhead, we limited these computations to Q vectors from the final 128 tokens of each chunk for H2O (see L230–L231), rather than using the full attention matrix.

The prefill latency comparisons in Table 4 (and Table 6(b) in the main text) highlights the empirical effect of the implementation. Specifically, StreamingLLM does not require the additional score computation, while InfiniPot and REFORM do. While InfiniPot shows some increased latency, REFORM shows reduced latency even with the additional operations, thanks to the computation savings from the early-exit strategy. Furthermore, the latency breakdown measurements show that the recomputation overhead for REFORM is minimal.

Table 4. Prefill time (seconds) on Mistral-Nemo-Instruct-2407, single H200 GPU.

---

[W2-2] Complexity analysis

In Table 6(b) in the main text and Table 4 in the rebuttal, we present empirical analysis on the latency and memory requirements for each method. Here, we discuss the time- and memory-complexity of different approaches.

Recurrent baselines such as StreamingLLM have a time complexity of O(L), where L is the input length, assuming a fixed (or bounded) chunk size and query length. Each recurrence step involves a constant amount of computation, and the total number of steps scales linearly with the input length. Their memory complexity is O(1), since only a fixed-size KV cache is maintained throughout.

In contrast, random-access baselines like InfLLM incur a time complexity of O(L^2), primarily due to periodic cache lookups across the full input length. These methods also require O(L) memory, as they store the entire KV cache in memory. This memory burden is significant, often requiring CPU offloading, which further increases latency.

REFORM maintains a time complexity of O(L) through its recurrent chunked processing. Although it has an O(L) memory cost due to storing token-level embeddings, these embeddings are very small in practice. Moreover, the early-exit strategy significantly reduces memory and compute requirements. Consequently, REFORM achieves faster speed and lower memory usage than both baselines (as demonstrated in Table 6(b) of the main text), highlighting the practical efficiency of our method.

---

[Q1] Clarifications of attention head selection.

We thank the reviewer for pointing this out and will revise the manuscript for better clarity. The head selection process is as follows:

Step 1: Data preparation. We construct two synthetic tasks for attention head selection: a pattern matching task and a multi-hop question answering task. For the pattern matching task, we insert sentences into the WikiText corpus with the format: “The value corresponding to the id {key} is {value},” where both the key and value are random 10-character ASCII strings. For the multi-hop QA task, we use passages from HotPotQA that contain the information required to answer a given question. The tokens corresponding to this key information are labeled as “golden tokens.”

Step 2: Context encoding. We apply recurrent chunked forwarding using H2O to each sample in the dataset, extracting token-level QKV embeddings from every layer and attention head.

Step 3: Significance score computation. For each QKV head in each layer, we compute a token-wise significance score. This is done by (1) calculating cosine similarity between context tokens and question tokens, (2) aggregating the similarity scores over the question tokens via max-pooling, and (3) smoothing the scores using mean pooling across a 20-token window.

Step 4: Performance measurement and head selection. We evaluate the effectiveness of each attention head using the Mean Normalized Rank (MNR) score, as described in Appendix C.1. Based on these scores, we select four heads: two that perform best on the pattern matching task, and two that perform best on the multi-hop QA task.

We will revise the manuscript to incorporate this explanation more clearly.

Dear reviewer NqoB,

We sincerely appreciate your efforts in reviewing our manuscript. We respond to each comment in the following content.

---

[W1] Evaluations on GSM-Infinite

Following your suggestion, we evaluated REFORM and RAG variants on the 'medium' split of GSM-Infinite. As shown in Table 1, REFORM consistently outperforms both sparse and dense RAG baselines, demonstrating its ability to handle more complex reasoning tasks beyond simple retrieval.

Table 1. GSM-Infinite ('medium' split) evaluation with Llama-3.1-8B-Instruct. RAG methods use an 8k retrieval budget and REFORM uses an 8k recomputation budget. Responses were generated without chain-of-thought reasoning for efficiency. Details follow Appendix C.3.

---

[W2] Comparison with Related Works

We compared REFORM to the suggested recurrent SnapKV baseline in Table 2. Notably, REFORM significantly outperforms SnapKV across multiple tasks, demonstrating its effectiveness.

Table 2. Comparison with recurrent SnapKV (Llama-3.1-8B-Instruct).

Before discussing other related works, we would like to emphasize that our main research goal is to propose an efficient, training-free context window extension method for large foundation models. Quest and MagicPIG tackle a different problem to improve efficiency within the model's native context window (e.g. <128k for Llama-3.1-8B-Instruct). We did not compare them in the paper as they cannot be used for inputs that are longer than the model's context limit.

Furthermore, Table 3 in the MagicPIG paper reports that Quest and MagicPIG scores 74.9 and 81.7 respectively on RULER 96k using Llama-3.1-8B-Instruct. In contrast, REFORM achieves 84.9 at 128k and 82.6 at 200k with the same model, suggesting that it outperforms these methods, particularly in longer contexts.

LM-Infinite employs a lambda-shaped mask to extend context, similar to StreamingLLM (a baseline in our paper). It optionally attends to prior tokens outside the lambda-mask for retrieval, but achieves only 70.3% accuracy on simple pass key retrieval tasks at 6k (Llama-2, 1.5x context limit), while REFORM reaches 100% at 1M (Llama-3.1, 8x context limit). Moreover, LM-Infinite’s full cache retention will lead to high memory requirements and offloading latency, similar to InfLLM.

We believe that these comparisons and discussions further clarify REFORM’s positioning in the long-context inference space, and will add them to the revised manuscript.

---

[W3, W4] More Detailed Latency Analysis

Here, we extend our latency analysis to three context lengths (previously only 256k), and separately present the pre-fill and decoding latency in Table 3 and Table 4.

In Table 3, InfiniPot and InfLLM shows slower processing speed compared to StreamingLLM due to the additional computation associated with dynamic cache eviction and retrieval. On the other hand, REFORM consistently shows faster prefill time compared to all baselines, thanks to the early-exit optimization. Further latency breakdown of the Compress+Gather stage and Recompute stage suggests that the recomputation overhead is minimal, taking under 1.52% of the total prefill time.

Table 4 shows a similar trend for the decoding time. InfiniPot and InfLLM shows slower decoding time compared to StreamingLLM due to the additional operations. On the other hand, REFORM also shows improved decoding speed as the generation is conditioned on KV cache created with a small recomputation budget (8k), unlike the baselines that condition on the full compressed cache (32k). (Note that these budgets are the hyperparameters used for the Mistral-Nemo experiments in our main text.)

Table 3. Time to first token (seconds) measurements (Mistral-Nemo-Instruct-2407, single H200, averaged over 20 runs).

Table 4. Time per output token (seconds) measurements (Mistral-Nemo-Instruct-2407, single H200, averaged over 20 runs and 200 tokens generated per measurement).

---

[W5-1] Discussion on Compression-Recomputation Trade-Off

While cache recomputation overcomes the forgetting problem, it introduces an additional recomputation overhead, leading to a trade-off. Latency breakdown in Table 3 suggests that the recomputation overhead is marginal, while our extensive evaluations in the main text demonstrate that the recomputation gives significant improvements in downstream performance.

---

[W5-2] Theoretical Analysis on Head Selection

Our attention head selection is motivated by findings from mechanistic interpretability literature, which show that different heads specialize in different functions [1,2]. Some heads are more useful than others for tasks like pattern recognition or reasoning. We leverage this diversity by selecting only a few high-performing heads, both for better identification performance and to keep the size of token-level embeddings tractable (including all heads would be memory-intensive). Our method thus aims for functional specialization as well as computational efficiency.

Although we use cosine similarity rather than attention weights to evaluate heads, our approach shares key insights with existing interpretability work. Prior studies show that lower-layer heads often detect syntactic or structural features, mid-layer heads handle semantic reasoning, and upper-layer heads guide output generation [3]. This aligns with our empirical findings: the most useful heads for pattern matching tend to appear in lower and upper layers, while those for multi-hop QA (i.e. semantic understanding) concentrate in the mid-layers. These trends are reflected in our MNR visualizations (Fig. 2 and Fig. 4). Also, Skean et. al. [4] suggests that the middle layer representations are more useful than the final layer representations for 32 MTEB (Massive Text Embedding Benchmark) tasks, similar to what we observe for the multi-hop QA task.

Furthermore, our choice of the head identification tasks is rooted in information retrieval literature, which shows that hybrid retrievers [5,6] combining sparse (structural) and dense (semantic) retrievers outperforms either alone. Inspired by this, we use two distinct head evaluation tasks: pattern matching and multi-hop QA. The former identifies heads that capture structural similarity, while the latter targets heads that encode semantic relevance. This dual-task strategy helps us select a complementary set of heads that are robust across diverse long-context scenarios.

Together, these design choices provide high token identification accuracy while maintaining efficiency.

[1] Voita et. al., Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned

[2] Ren et. al., Identifying semantic induction heads to understand in-context learning

[3] Zheng et. al., Attention Heads of Large Language Models: A Survey

[4] Skean et. al., Layer by Layer: Uncovering Hidden Representations in Language Models

[5] Bruch el. al., An analysis of fusion functions for hybrid retrieval.

[6] Chen el. al., Out-of-domain semantics to the rescue! zero-shot hybrid retrieval models.

Dear Reviewer NqoB,

Thank you very much for acknowledging our rebuttal. We are grateful for the time and effort you put in the review and reading through the rebuttals, which has greatly contributed to improving the clarity and soundness of our paper.

Through the rebuttal, we believe that we have clarified your concerns successfully. We would greatly appreciate any further comments or insights you may have after reading our response, especially if any of your earlier concerns remain. We would be glad to engage in further discussion during the remainder of the rebuttal period.

Many thanks,

Authors