Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads

Q3: Hyperparameter Sensitivity: Have the authors conducted ablation studies on the cache budget b and chunk size B? How do these parameters impact performance and memory usage?

Please refer to our response to W3, where we have introduced additional experiments to illustrate the hyperparameter sensitivity of cache budget and chunk size.

Q4: Generalization: It might be out-of-scope, but how well does LOCRET generalize to other transformer architectures, such as encoder-decoder models or those with different attention mechanisms?

Thank you for this valuable question. Locret is compatible with any transformer-based LLM that utilizes a KV cache. While encoder-decoder models can benefit from Locret by reducing the KV cache burden in the decoder, the encoder part, which does not use a KV cache due to non-causal inference, is not compatible with Locret.

Additionally, other attention mechanisms, such as linear attention and multi-latent attention, are compatible with Locret. Linear attention modifies only the calculation of attention while maintaining the KV cache structure, and multi-latent attention features a single-head KV cache-like structure that is also compatible with Locret.

Q5: Limitations: Are there specific tasks or contexts where attention pool-based methods might outperform LOCRET? How does LOCRET handle scenarios with severe context discontinuity?

When the budget is extremely limited, such as 128 or 256 tokens, attention pool-based methods can outperform cache eviction methods. This is because eviction-based methods, including Locret, degrade to a StreamingLLM pattern by retaining only the initial and local tokens. As illustrated in Figures 5 and 6, Locret also exhibits a StreamingLLM pattern in these scenarios. Since LoCoCo surpasses StreamingLLM, it can also outperform eviction-based methods under such budget constraints.

However, Table 5 demonstrates the compatibility between LoCoCo and Locret. In scenarios with strict budget constraints, LoCoCo can enhance performance, and when the budget is larger, Locret can further boost standalone LoCoCo's performance. Replacing H2O in LoCoCo with Locret is an effective strategy for improving performance across all scenarios.

Locret assigns continuous position information to the evicted, discontinuous cache to mitigate performance issues related to severe context discontinuity. We are currently exploring training techniques for the LLM backbone to better process discontinuous context, which we plan to address in future work.

Q6 & Q7: Quantization Methods: You mention KV cache quantization techniques, mentioning the computation overhead as their limitation. Could you compare these techniques, e.g., KVQuant, with sparse attention methods such as FastGen? Combination: You mention the possibility of combining your approach with other efficient inference methods. Could you expand on this with results?

Thank you for this advice. Locret is orthogonal to KV cache quantization, making it compatible with KVQuant. Additionally, Locret is orthogonal to FastGen, which employs a mixture of eviction policies across different heads. Locret can be applied as a specific policy to selected heads to achieve a higher compression ratio.

We are also interested in exploring how Locret can be combined with other efficient inference methods, such as LLM backbone quantization and speculative decoding. Existing works have demonstrated the potential of such combinations; for instance, TriForce [6] integrates H2O-based KV cache compression with speculative decoding, achieving improved decoding throughput. Since Locret can function as a cache importance scoring mechanism, it could be utilized in scenarios where H2O is currently applied. These explorations and combinations involving Locret and other efficient inference methods will be part of our future work.

[1] H2O: Heavy-hitter oracle for efficient generative inference of large language models

[2] SnapKV: SnapKV: LLM knows what you are looking for before generation

[3] SirLLM: Streaming infinite retentive LLM

[4] Longbench: A bilingual, multitask benchmark for long context understanding

[5] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory

[6] Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding

We sincerely appreciate your detailed review and valuable feedback.

W1: Clarity of Presentation: The paper contains grammatical errors and unclear notations, hindering understanding.

Thank you for your comments. We will conduct a thorough grammar review and ensure that all notations are clearly defined before use in our next revision soon. We are committed to improving the overall quality of our writing.

W2: Theoretical Depth: The theoretical underpinnings, particularly regarding the causal importance score and its properties, could be more thoroughly developed.

Thank you for pointing this out. I would like to provide the motivation behind designing the CIS, which may help to intuitively understand our proposed methods.

The goal of this paper is to enable long-context inference on consumer-grade devices, particularly in memory-constrained scenarios where GPU memory usage must be strictly controlled. To manage GPU memory consumption, we employ a chunked prefill pattern, which is crucial for reducing memory requirements. However, existing cache importance scoring functions cannot be seamlessly integrated with chunked prefill, as their importance estimation for a token (or cache) depends on subsequent tokens. Locret addresses this limitation by using a causal importance score, meaning the importance score of a cache unit does not depend on subsequent cache units.

A key property of CIS is its causality, which ensures compatibility with chunked prefill. To further explore the incompatibility of existing scoring functions with chunked prefill, we conducted an experiment measuring the consistency of the top 10% most important cache positions in prefixes of varying lengths compared to the full context (truncated to 6K for reference). The results are as follows:

This experiment demonstrates that scoring functions relying on subsequent information, such as H2O and SnapKV, show significant discrepancies in predicted cache importance when future cache units are not considered. SirLLM, while being a causal importance scoring function, suffers from inaccuracy that leads to substantial performance degradation, as shown in Table 2 and Table 3. Locret, however, avoids these discrepancies and delivers strong performance.

W3: Hyperparameter Analysis: Limited discussion on the impact of key hyperparameters (e.g., cache budget, chunk size) on performance.

Thank you for your comment. We have included an ablation study on the length of stabilizers (Figure 3) and will incorporate additional ablation studies on cache budget and chunk size in our next revision.

Regarding the impact of cache budget, we conducted experiments with different cache budgets on subsets of LongBench [4], using Phi-3-mini-128K with retaining heads trained on LongAlign. The results indicate that a larger cache budget generally leads to improved performance, although certain tasks are more sensitive to cache budget variations than others.

For chunk size, we conducted an experiment evaluating Locret trained with Phi-3-mini-128K LongAlign on the L-Eval's Natural Questions (NQ) dataset using different chunk sizes. The results are as follows:

This experiment demonstrates the hyperparameter stability of chunk size. A smaller chunk size can be adopted for memory-limited scenarios, such as end-side devices, while a larger chunk size can be used to boost inference speed. Adjustments to chunk size have only a minor impact on performance.

W4: Limited Discussion of Limitations: The paper does not sufficiently explore potential drawbacks or scenarios where LOCRET may underperform.

Thank you for your advice. In our paper, we evaluated Locret on two LLMs (Phi-3-mini-128K and Llama-3.1-8B-instruct) and two types of hardware platforms (A800/H800 and 4090). Locret is compatible with various scenarios due to minimal modifications to the model architecture and low hardware requirements.

However, we recognize some potential scenario limitations. For example, the eviction action requires GPU support for the "gather" kernel, making Locret incompatible with some NPU platforms that lack this capability. Additionally, Locret is designed for decoder-only transformer architectures and is not suitable for encoder models or models based on other architectures (e.g., RNNs). We will expand our limitations section to include these and other potential drawbacks in our next revision.

W5: Reproducibility: Some essential details for reproducing results are located in the appendix rather than the main text.

To enhance readability and provide clearer instructions for reproducing our experiments, we will move key hyperparameters and training/evaluation details from the appendix to the main text. Specifically, we will transfer the essential hyperparameters outlined in Appendix A.1 and A.2 to Section 4.1 (Experimental Setup). Additionally, we will include more detailed instructions on training the retaining heads and evaluating the trained model.

Q1: Stabilizer Length: Could the authors provide more insight into how the stabilizer length ns affects performance across different models and datasets? Is there an optimal range for ns?

There is a tradeoff between stabilizer length and effective cache budget. Stabilizers occupy space in the retained cache, so a larger stabilizer length reduces the space available for other retained cache units, potentially resulting in the loss of important information and leading to performance degradation. Conversely, a shorter stabilizer length can increase the instability of CIS predictions, leading to more errors during eviction.

From our observations, context retrieval tasks require a larger $n_s$​ due to their need for more accurate CIS prediction and eviction, rather than ample space for important caches. In Figure 3(a), when $n_s$​ is small, there is significant performance degradation in retrieving the correct number. On the other hand, natural language understanding tasks, such as summarization, benefit from shorter $n_s$​ values, as maximizing space for important caches is crucial for better performance.

We conducted an additional experiment on the QMSum task from L-Eval with various stabilizer lengths, keeping the cache budget fixed at 6000. The results demonstrate that overly large $n_s$​ values occupy too much space in the cache, causing performance degradation. The results are as follows:

We recommend using ~2500 as the optimal value for $n_s$, as shown in Table 6. This value strikes a balance across different types of tasks. We will include the above discussion and this recommendation in our next revision.

Q2: Theoretical Justification: Can the authors elaborate on the causal importance score's theoretical properties and explain how it ensures minimal approximation error during cache eviction?

The intuition behind the design of training-based CIS prediction is explained in W2, and we have added further experiments to support this discussion.

Defining a definitive "golden label" for the KV cache eviction problem is difficult. Estimating the impact of evicting a group of KV cache units is complex, and to our knowledge, no prior work has specifically focused on this topic. In fact, the most commonly used approaches and the baselines included in this study—such as H2O, SnapKV, SirLLM, and InfLLM [5]—rely on heuristics to evaluate eviction impact or estimate importance based on statistical metrics like attention scores. Consistent with these methods, we identify important cache units using higher attention scores to generate training labels, as a higher attention score generally correlates with greater importance in the attention mechanism.

Manuscript updates are as follows. Note: we use the new figure/table numbering here.

W2: Figure 1 and line 076-084.

W3, Q3: Figure 6(a) and Figure 6(c).

W4, Q5, Q6, Q7: line 530-539.

W5: line 349-350, line 355-356, and Table 2.

Q1: Figure 6(b).

Manuscript updates are as follows. Note: we use the new figure/table numbering here.

W1: Figure 6(a), Table 5 and Table 16.

W3, Q3: Table 17.

Q1: line 858-859.

Q2: line 530-539.

We sincerely appreciate your detailed review and valuable feedback.

W1.1: SirLLM is not an appropriate baseline. A more suitable baseline for eviction-based methods would be SnapKV.

Thank you for your suggestion. We have conducted additional experiments using SnapKV [1] with chunked prefill, and our results show that Locret outperforms SnapKV in this scenario. The necessity of eviction with chunked prefill is further explained in W1.2.

Here, we evaluate Locret against SnapKV (as well as H2O [2], a well-established scoring function for cache importance) on various subsets of InfiniteBench, utilizing Phi-3-mini-128K. The experiment follows the same setup as outlined in Table 6, with the exception of varying the scoring function. The results are summarized below.

Additionally, we test Locret and SnapKV in the chunked prefill scenario on LongBench with varying cache budgets. We exclude the Chinese subtasks from LongBench, as the model we are using, Phi-3-mini-128K, is not specifically trained on a Chinese corpus. The retaining heads are trained on the SFT dataset, LongAlign, for 3000 steps. Hyperparameters are consistent with those outlined in Table 6 of our paper.

From the experiments above, Locret demonstrates superior performance compared to SnapKV. SnapKV experiences significant performance degradation across all tested subtasks, particularly on context retrieval tasks, such as R.Number in InfiniteBench, highlighting the incompatibility of SnapKV with chunked prefill. In contrast, Locret proves to be a more accurate scoring function under these conditions.

W1.2: Although chunk prefilling may not align perfectly with SnapKV, the authors could still avoid OOM errors and reduce GPU peak memory usage by employing layer-by-layer token dropping during prefilling.

Thank you for your insightful comment. As demonstrated in the additional experiments in W1.1, SnapKV does not align well with chunked prefill. We believe that eviction with chunked prefill is more efficient than layer-by-layer token dropping, as the latter requires storing a full KV cache for each layer at some point. This behavior can still be resource-intensive when processing longer sequences. For instance, when handling context with a length of 10 million tokens, the full cache for a single layer occupies 120GB of GPU memory, which is difficult to accommodate even with a single A100/H100 GPU (with only 80GB of memory). In contrast, Locret is able to process such large contexts while keeping GPU memory usage within reasonable limits.

Additionally, we would like to include the 10M pass key experiment suggested by Reviewer V3dC. We conducted the experiment using the same hyperparameters as in Table 6, with a budget size of only 6000. The maximum GPU memory used in this setup was less than 18GB.

The experiment demonstrates that Locret can successfully handle pass key retrieval with a context length of 10 million tokens, and it can be conducted on a single GPU. In contrast, SnapKV with layer-by-layer token dropping requires at least 120GB of GPU memory to store the KV cache for a single layer, consuming significantly more hardware resources. Moreover, Locret is capable of processing streaming context input, where the total context length is infinite. Layer-by-layer token dropping, however, is unsuitable for such scenarios.

We will include the additional experimental results from W1.1 and W1.2 in our next revision and will emphasize the discussion of eviction with chunked prefill versus layer-by-layer token dropping, as it is crucial for clarifying our research objectives.

Thank you for your response. I believe that 1.5GB peak memory consumption is acceptable even for the 4090, which is entirely satisfactory.

If you are not inclined to measure the 128K RULER, you can opt for shorter ones, such as the 32K or 64K RULER. Alternatively, you can choose GQA models. The sequence length is not the key point; the key point is that you need to measure complex tasks like RULER fairly.

If the proposed method cannot outperform SnapKV, which is a widely used baseline, I think it is not good enough.

Thank you to the authors for providing additional experiments and detailed explanations. I appreciate your efforts. I have a few further suggestions and comments:

• Unfair Evaluation on SnapKV: I recommend that the authors report the performance of SnapKV on InfiniteBench or RULER (<128K) without chunk prefill, as this scenario would not cause out-of-memory issues with a single A100/H100 GPU (with 80GB of memory). Conducting chunk prefill would not be a fair evaluation for SnapKV, as it is not designed for such usage.
• 10M Pass Key Experiment: The 10M pass key experiment appears to be of limited significance because the task is overly simplistic—the “haystack” and “needle” are too distinct, making the task easy to hack through training.
• Poor Generalization: The performance of the proposed method on RULER for complex tasks suffers from accuracy drops, even when trained on datasets specifically curated for RULER. This indicates that the method lacks generalizability.
• Unsuitable Multi-Turn Task: The current evaluation dataset for multi-turn tasks does not seem ideal. As shown, StreamingLLM’s performance is not particularly poor (over 30 points as reported in SirLLM’s paper), and your method’s improvement is marginal. Since StreamingLLM is widely regarded as a weak baseline (e.g., its inability to perform needle retrieval), a better evaluation dataset is needed. A simple alternative could be RULER's multi-key task, where the model retrieves different keys in different turns. However, as the results suggest, Locret appears to struggle with such tasks.
• Lack of Off-Topic Turns Experiments: It would be beneficial to include an experiment that evaluates off-topic turns. For instance, in RULER's QA task with multiple documents, you could ask different off-topic QA questions over multiple rounds before asking the main question. This would help assess whether accuracy drops in such settings.

I hope these suggestions provide useful insights to strengthen your paper further. Thank you again for your effort for conducting additional experiments.

We sincerely appreciate your detailed review and valuable feedback.

W1: The additional training in the proposed method is still resource-intensive.

We addressed the cost of additional training for Locret in lines 344–345, where we specify that the training for both models in our benchmark requires less than 1 GPU hour on a single A800 GPU.

Although Appendix A.2 (System Environment) mentions that our experiments were conducted on an 8*A800/H800 GPU cluster, we want to clarify that only a single GPU was utilized for training. We acknowledge the potential confusion and appreciate your observation. We will make this clearer in our next revision.

W2: Novelty concern. It is unclear why to train the heads to obtain CIS instead of using existing methods like H2O.

We appreciate your thoughtful feedback. Below, we explain our rationale for employing a training-based approach for predicting CIS and performing eviction, instead of directly using existing cache importance scoring functions such as H2O.

Our primary objective is to enable long-context inference on consumer-grade devices, particularly under memory-constrained conditions where GPU memory usage must be strictly controlled. To achieve this, we utilize a chunked prefill pattern that is essential for reducing memory consumption. However, existing cache importance scoring methods cannot be adapted to chunked prefill due to their reliance on subsequent tokens for estimating the importance of a token or cache.

Locret addresses this limitation by employing a causal importance scoring mechanism, where the importance score of a cache unit does not depend on future units. This feature allows seamless integration with chunked prefill, solving the issue posed by non-causal methods.

To further illustrate the incompatibility of existing scoring functions with chunked prefill, we conducted an experiment measuring the consistency of the top 10% most important cache positions in prefixes of varying lengths compared to a 6K full context. The results are shown below:

This experiment highlights that scoring functions requiring future information, such as H2O and SnapKV, suffer from significant discrepancies when subsequent cache units are not considered. On the other hand, SirLLM, while also causal, shows notable inaccuracies, leading to performance degradation as demonstrated in Table 2 and Table 3 of our paper.

We also evaluated the end-to-end performance using H2O and SnapKV with chunked prefill on a subset of InfiniteBench:

The results demonstrate that discrepancies between local and global importance scores in H2O and SnapKV lead to severe performance drops, particularly in R.Number. Locret, however, avoids such inconsistencies and achieves superior performance.

We appreciate this insightful comment and will include this analysis and visualized results in our next revision.

W3: This paper lacks a detailed analysis of the CIS.

We have added further clarification regarding the purpose of designing CIS in our response to W2. Additionally, we will incorporate the discussion and experimental results outlined in W2 into the main text in our next revision.

Q1: Explaination of "the weakening correlation between local and global importance as sequences grow exacerbates this issue".

Existing cache importance scoring functions, such as H2O and SnapKV, are designed to identify important cache units only after the entire input sequence is prefilled. H2O relies on the complete attention scores to determine heavy hitters, while SnapKV’s voting mechanism requires the attention scores of the local window at the end of the input sequence, which also mandates full sequence prefill before eviction. When using chunked prefill, subsequent cache units that have not yet been processed are inaccessible, leading to significant discrepancies when applying H2O or SnapKV to prefilled cache units. This discrepancy arises because the predicted importance based on partial input diverges from the actual importance computed with the full sequence.

We demonstrated this effect through an additional experiment in W2, which highlights the inconsistency in H2O and SnapKV. We hope this experiment clarifies the limitations of these methods in scenarios involving chunked prefill.

Q2: How effective would it be to directly use the maximum value of each column as a metric during inference?

We appreciate this question. However, it is not feasible to use the maximum value of each column as a metric during inference. As discussed in W2, a key objective of this paper is to integrate an eviction policy with chunked prefill. In chunked prefill, the subsequent cache units (tokens or hidden states) are not accessible, preventing us from calculating the attention score of a token in relation to all subsequent tokens. The maximum value of each column represents the highest attention score of a token’s query to the keys of all subsequent tokens, which cannot be determined during chunked prefill. Therefore, using the maximum value of each column as a metric during inference is not possible.

Q3: What are the actual performances of H2O combined with quantization?

Quantization combined with H2O leads to an attention shift, resulting in inaccurate cache importance estimation, as demonstrated in the Q-Hitters paper[4]. In Section 4.2, the authors report that the overlap ratio of identified heavy-hitters drops below 50% when quantization is applied. Additionally, Figure 7 of the Q-Hitters paper illustrates significant performance degradation when H2O is used with standard quantization techniques. We hope this addresses your question and clarifies the limitations of using H2O with quantization.

Q4: Comparison between SnapKV and PyramidKV.

Thank you for pointing this out. We will include H2O and SnapKV as additional baselines and have conducted comparative experiments with Locret on specific subsets of InfiniteBench. Since PyramidKV primarily manages budget allocation, it operates orthogonally to the eviction function and can be combined with Locret. We also provide results for the combination of Locret and PyramidKV. Due to the time-consuming nature of running the full benchmark, we were unable to generate complete results for InfiniteBench at this stage, but these will be included in our final revision.

Additionally, the slow inference speed of H2O and SnapKV stems from their incompatibility with flash-attention, as both methods require access to the attention scores, which the current implementation of flash-attention does not support.

In our experiments, we modified only the scoring function, keeping all other hyperparameters consistent with Appendix A.2. When integrating with PyramidKV, we used maximum pooling among the CIS (following PyramidKV's setting) and set $\beta=2$.

The results indicate that Locret outperforms H2O and SnapKV in chunked prefill scenarios for long-context inference. H2O and SnapKV show limitations in accurately predicting context retrieval tasks, such as R.Number. Additionally, incorporating PyramidKV for budget allocation management further enhances overall performance, demonstrating the compatibility between Locret and PyramidKV.

Q5: Is the Stabilizer used only for selecting recent tokens?

No, the stabilizers refer to the last $n_s$ tokens in each chunk during the chunked prefill process. These tokens are retained without eviction to maintain a local and continuous context, thereby minimizing errors (as stated in line 292).

Q6: Is the performance improvement in this paper due to the SFT? What would be the effect if SFT were directly applied to the model?

The performance improvement is not attributed to SFT. We use a minimal amount (3,000 entries) of long-context SFT data solely to train the retaining heads. Importantly, there is no SFT loss involved in Equation 3, and the LLM backbone remains frozen throughout. In other words, the only learnable component in our framework is the scoring function responsible for identifying which cache units are more important.

Q7: Should the number of heads in a retaining head be the same as in Query, or should it match Key/Value?

The number of retaining heads must match the number of Key/Value heads. To train the retaining heads, we select the maximum attention score (before softmax) across different query heads within the same group (as described in line 236).

[1] H2O: Heavy-hitter oracle for efficient generative inference of large language models

[2] SnapKV: LLM knows what you are looking for before generation

[3] SirLLM: Streaming infinite retentive LLM

[4] Q-Hitter: A Better Token Oracle for Efficient LLM Inference via Sparse-Quantized KV Cache

Manuscript updates are as follows. Note: we use the new figure/table numbering here.

W1: line 841-842.

W2, W3: Figure 1, Table 5.

Q1: line 076-084.

Q4: Table 5, Table 11.

Dear Reviewer xs3N,

We would like to appreciate our gratefulness again for your valuable reviews. We have already uploaded the revised version based on your suggestions. Since the discussion period will be ending soon, we are looking forward to discussing the newly added content and our response with you. Should our responses have addressed your concerns, we would be grateful for an improved score. Thanks again for your time and effort.

Submission 5591's Authors

We deeply thank you for your thorough review.

W1: The proposed method is only evaluated with two LLMs (Phi-3-mini-128K and Llam-3.1-8B-instruct) and one hardware platform (Nvidia 4090 GPU).

Thank you for your valuable comments. The proposed method, Locret, has been evaluated on two LLMs with MHA and GQA architectures. Locret is compatible with any decoder-only LLM that utilizes a KV cache, as the modifications to the model architecture are minimal. Decoder-only Transformers with MHA or GQA attention mechanisms are widely adopted in both academia and industry, making them the most popular LLM architectures at present. We believe our experiments demonstrate that Locret performs effectively with such model architectures.

Regarding the hardware platform, we evaluated Locret on A800/H800 GPUs (Table 2) and the 4090 GPU (Table 3). Locret has minimal hardware requirements, as the algorithm is implemented entirely in PyTorch and does not require any CUDA-level modifications. Therefore, we believe it can be executed on any GPU platform while achieving comparable performance.

W2: How is the performance on other long-context benchmarks, such as longBench, et al. ?

Thank you for your valuable advice. We conducted additional experiments to evaluate Locret on LongBench [1], comparing it with baselines such as Full Attention, MInference [2], InfLLM [3], and SirLLM [4]. For this evaluation, we used Phi-3-mini-128K with a retained head trained on LongAlign. To ensure a fair comparison, we excluded all Chinese subtasks from LongBench and focused solely on the English subtasks, as Phi-3-mini-128K was not specifically trained on Chinese corpora. The results are presented below.

We also report the maximum memory usage, including GPU memory, CPU memory, and total maximum memory, alongside the average score on LongBench. For FullAttn, we exclude the maximum memory usage, aligning with Figure 4.

From the experiments above, Locret demonstrates the best overall performance and excels in the majority of subtasks. It outperforms all the baselines without any noticeable performance degradation while consuming less memory. Although MInference also avoids performance drops, it requires more GPU memory compared to Locret. SirLLM achieves comparable memory usage but shows some performance decline compared to FullAttn and Locret. InfLLM exhibits the most significant performance drop, and its offloading mechanism results in the highest CPU memory usage, making it the method with the largest total memory consumption. These results highlight Locret as an outstanding approach for evaluation on LongBench.

We plan to integrate the experiments mentioned above into the experimental section in our next revision. The comparison of LongBench scores and memory consumption, along with the accompanying discussion, will also be included.

Thank you for your valuable insights on improving our paper by incorporating additional experiments. We hope this response thoroughly addresses your questions. If there are any remaining concerns, we are happy to provide further clarifications or conduct additional experiments. We kindly request a raise of the overall rating if all issues have been addressed.

[1] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding

[2] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

[3] SirLLM: Streaming infinite retentive LLM

[4] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory

Manuscript updates are as follows. Note: we use the new figure/table numbering here.

W1: line 858-859.

W2: Table 7, Table 8.

Dear Reviewer j7cb,

We would like to appreciate our gratefulness again for your valuable reviews. We have already uploaded the revised version based on your suggestions. Since the discussion period will be ending soon, we are looking forward to discussing the newly added content and our response with you. Should our responses have addressed your concerns, we would be grateful for an improved score. Thanks again for your time and effort.

Submission 5591's Authors

Thanks for you advice on testing more models. We conduct experiments using the LongBench dataset on Llama-3.2-1B-instruct, and compare Locret with InfLLM and SirLLM, following the LongBench experiments with Phi-3-mini-128K above. We skip MInference, as MInference does not officially support Llama-3.2-1B-instruct (they have not provided the official configuration of heads).

The results are presented below. (The highest score among Locret, SirLLM and InfLLM is marked in bold font, and FullAttn is provided as a reference.)

From the results above, Locret also shows compatibility on Llama-3.2-1B-instruct. It only exhibits neglegible performance drop, while SirLLM and InfLLM shows large performance degredation. Locret also uses the less overall memory compared with the tested baselines. We will integrate the additional results to the next revision (maybe the camera ready version, as we cannot upload pdf now) of our paper.

We hope such additional results are able to show the generalizability of the proposed algorithm, Locret. It can be applied to the majority of decoder-only LLMs. We are also glad to test more settings if there are further needs.

W2.2: Additionally, while the framework shows reduced memory requirements, further evidence on latency and computation trade-offs associated with retaining heads would be beneficial for practitioners evaluating deployment feasibility.

We have evaluated the model's forward throughput under varying context lengths, both with and without retaining heads. The results are summarized below. "R" represents the retaining heads, and the throughput is reported in tokens per second (tok/s) in the format "Ave. / Std."

From the results, no significant latency increase is observed when using retaining heads. The numerical differences are attributed to systematic variations rather than additional overhead introduced by retaining heads during inference.

We have highlighted in our paper that the retaining heads trained for Phi-3-mini-128K and Llama-3.1-8B-instruct account for only 8% and 2.5% of the original model size, respectively (line 343). This minimal size overhead introduces negligible difficulty for deployment.

W3: Discussing potential limitations or adjustments required for alternative models would enhance the generalizability of the method.

Thanks for pointing this out. We will revise our limitation section to clearly state that Locret has been tested on decoder-only MHA and GQA architectures. Exploring the compatibility of Locret with other model architectures, such as MLA, remains part of our future work. Additionally, we aim to investigate the integration of Locret with other models, such as encoder-decoder architectures, in the future. This discussion will be incorporated into Section 5 (Discussion) in our next revision.

Q1: Could the authors clarify how LOCRET’s retaining heads would handle extremely high-context lengths (e.g., 10 million tokens)? Would additional constraints or modifications be required to manage the scoring of cache units in such contexts?

There are no additional constraints or modifications required for Locret to handle longer contexts.

We conducted the following experiment to demonstrate Locret's capability in extremely long-context scenarios. We created a dataset similar to InfiniteBench's R.Number, with an average length of 10 million tokens. Each data point contains a 10-digit number string inserted into an irrelevant context, and the task is to retrieve the inserted number. The dataset consists of 50 examples, with the number strings uniformly distributed throughout the context. We used the hyperparameters from Table 6, with the exception of setting the chunk size to 10240 to speed up inference. The results, presented below, show that Locret can efficiently process extremely long contexts. In this experiment, the cache budget is set to 6000, and the compression ratio is 1747.6×.

Q2: What are the potential reasons of Locret outperforming SirLLM on memory-demanding tasks?

There are two reasons why Locret outperforms SirLLM on memory-demanding tasks:

First, SirLLM uses token-entropy to estimate token importance, where higher token-entropy indicates a token is harder to predict based on its context, deeming it more important. However, as highlighted in the Limitation Section of SirLLM's paper, significant discrepancies between user data and the model's training data can lead to poor eviction decisions. Memory-demanding tasks, such as context retrieval, often fall into such scenarios. For instance, tasks like R.PassKey and R.Number involve data patterns that are rare in natural contexts. As these patterns are not well-represented in the model’s training data, SirLLM fails to provide accurate token-entropy predictions, resulting in suboptimal performance.

Second, SirLLM’s token-entropy is a token-level metric, while Locret's Causal Importance Score (CIS) operates at the cache unit level. This distinction allows Locret to assign different eviction policies across heads, enabling more flexible and effective cache management. In contrast, SirLLM lacks this flexibility, as its eviction strategy is uniformly applied. Previous studies [5, 6] have demonstrated that eviction patterns often vary among heads, and Locret’s ability to accommodate such variations provides a significant advantage in memory-demanding tasks.

We sincerely appreciate your detailed review and valuable feedback.

W1: A more comprehensive analysis highlighting LOCRET’s distinctions from similar heuristics in cache management would strengthen the contribution.

Thanks for pointing this out. We have conducted the following experiment to highlight the differences between Locret and existing cache importance scoring functions, e.g. H2O and SnapKV. Cache importance scoring functions can generally be categorized into two types: causal and non-causal.

• Non-causal scoring functions: Examples include H2O and SnapKV. These methods require information from subsequent cache units to determine the importance score of a cache unit, making them dependent on prefilling the entire sequence.
• Causal scoring functions: Examples include SirLLM and our proposed method, Locret. These methods predict cache importance without relying on subsequent information.

Non-causal scoring functions are incompatible with chunked prefill because they cannot calculate scores without access to the full sequence. If such functions are integrated with chunked prefill, they often face a significant discrepancy between the local importance score (without considering subsequent information) and the global importance score (with full context).

To investigate this discrepancy, we measured the consistency of the top 10% most important cache positions identified in prefixes of various lengths compared to the full context. For reference, the full context is truncated to 6K tokens. The results are as follows:

This experiment highlights that scoring functions requiring future information, such as H2O and SnapKV, suffer from significant discrepancies when subsequent cache units are not considered. On the other hand, SirLLM, while also causal, shows notable inaccuracies, leading to performance degradation as demonstrated in Table 2 and Table 3 of our paper.

We also evaluated the end-to-end performance using H2O and SnapKV with chunked prefill on a subset of InfiniteBench:

The results demonstrate that discrepancies between local and global importance scores in H2O and SnapKV lead to severe performance drops, particularly in R.Number. It is this discrepancy that leads to the failure of H2O and SnapKV in accurately retrieving information from the context. Specifically, the model is unable to identify the importance of certain cache units at the time they are first encountered. Locret, however, avoids such inconsistencies and achieves superior performance.

W2.1: The results indicate promising efficiency gains but lack granular performance data on how LOCRET’s accuracy scales with different cache budgets across various architectures.

In order to investigate the impact of cache budget, we conducted experiments with different cache budgets on subsets of LongBench [4], using Phi-3-mini-128K with retaining heads trained on LongAlign. The results indicate that a larger cache budget generally leads to better performance, although certain tasks are more sensitive to cache budget variations than others.

We have tested Locret on two LLMs—Phi-3-mini-128K with an MHA architecture and Llama-3.1-8B-128K with a GQA architecture—demonstrating its compatibility with some of the most widely adopted model architectures. Exploring its applicability to other architectures, such as Encoder-Decoder models or MLA models, would be an interesting direction for future work. We will highlight this as part of the limitations in our next revision.

Q3: Could the authors provide further insights or examples where the heuristic scoring might diverge significantly from the true causal importance?

Since Locret is a training-based method, it may perform poorly on contexts that have low probability in the training data. For instance, punctuation is often treated as unimportant in the training dataset. If the query focuses on retrieving specific punctuation marks from a long context, Locret may provide a biased CIS and incorrectly evict important cache units.

However, this issue can be mitigated by designing a more representative training dataset. Moreover, such problems are inherently challenging even for humans to solve, as they require understanding the context where insignificant elements like punctuation become important. We will address this limitation in our next revision.

[1] H2O: Heavy-hitter oracle for efficient generative inference of large language models

[2] SnapKV: SnapKV: LLM knows what you are looking for before generation

[3] SirLLM: Streaming infinite retentive LLM

[4] Longbench: A bilingual, multitask benchmark for long context understanding

[5] Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs

[6] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

Manuscript updates are as follows. Note: we use the new figure/table numbering here.

W1: Figure 1, Table 5.

W2: Figure 6(a), Table 18.

W3, Q3: line 076-084.

Q1: Table 16.