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

Authors: Yuxiang Huang, Binhang Yuan, Xu Han, Chaojun Xiao, Zhiyuan Liu
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: Long-context Inference, Memory Efficient Inference, Large Language Models

Abstract:
Large language models (LLMs) have shown remarkable advances in supporting long-context comprehension and processing tasks. However, scaling the generation inference of LLMs to such long contexts incurs significant additional computation load, and demands a substantial GPU memory footprint to maintain the key-value (KV) cache of transformer-based LLMs. Existing KV cache compression methods, such as quantization, face memory bottlenecks as context length increases, while static-sized caches, such as selective eviction, suffer from inefficient policies. These limitations restrict deployment on consumer-grade devices like a single Nvidia 4090 GPU. To overcome this, we propose Locret, an efficient framework for long-context LLM inference that introduces retaining heads to evaluate the causal importance of KV cache units, allowing for more accurate eviction within a fixed cache size. Locret is fine-tuned on top of the frozen backbone LLM using a minimal amount of data from standard long-context SFT datasets. During inference, we evict low-importance cache units along with a chunked prefill pattern, significantly reducing peak GPU memory usage. We conduct an extensive empirical study to evaluate Locret, where the experimental results show that Locret outperforms the recent popular and competitive approaches, including InfLLM, Quantization, SirLLM, and MInference, in terms of memory efficiency and the quality of generated contents --- Locret achieves over a $20\times$ and $8\times$ KV cache compression ratio compared to the full KV cache for Phi-3-mini-128K and Llama-3.1-8B-instruct. Additionally, Locret can be combined with other efficient inference methods, such as quantization and token merging. To the best of our knowledge, Locret is the first framework capable of deploying Llama-3.1-8B or similar models on a single Nvidia 4090 GPU, enabling 128K long-context inference without compromising generation quality, and requiring little additional system optimizations.

Decision: reject
Reviewer scores: 8, 3, 5, 5, 8

Reviews and discussion:
<thread>
creation date:
1729833470164

thread_type:
official_review

<thread_object>
creator: reviewer_wkxd

<content>
summary:
This paper presents LOCRET, a framework that enhances long-context large language model (LLM) inference on consumer GPUs. LOCRET introduces "retaining heads," lightweight components added to a frozen LLM backbone to estimate the importance of each key-value (KV) cache unit. LOCRET optimizes cache eviction and reduces GPU memory usage during inference by predicting which cache units are crucial. Combined with chunked prefill, it outperforms methods like InfLLM and SirLLM in memory efficiency and generation quality, enabling models like Llama-3.1-8B to run 128K context inference on a single Nvidia 4090 GPU without performance loss.

strengths:
Novelty: The introduction of retaining heads for estimating causal importance is a novel approach to KV cache management.
Practical Impact: Enables deployment of large LLMs on consumer-grade GPUs without significant performance loss.
Comprehensive Evaluation: Extensive experiments across multiple datasets and models validate the effectiveness of LOCRET.
Compatibility: LOCRET can be integrated with other efficient inference methods like quantization and token merging.
Lightweight Training: Requires minimal additional training time and resources

weaknesses:
Clarity of Presentation: The paper contains grammatical errors and unclear notations, hindering understanding.
Theoretical Depth: The theoretical underpinnings, particularly regarding the causal importance score and its properties, could be more thoroughly developed.
Hyperparameter Analysis: Limited discussion on the impact of key hyperparameters (e.g., cache budget, chunk size) on performance.
Limited Discussion of Limitations: The paper does not sufficiently explore potential drawbacks or scenarios where LOCRET may underperform.
Reproducibility: Some essential details for reproducing results are located in the appendix rather than the main text.

questions:
1. 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?
2. 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?
3. 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?
4. 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?
5. Limitations: Are there specific tasks or contexts where attention pool-based methods might outperform LOCRET? How does LOCRET handle scenarios with severe context discontinuity?
6. 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?
7. Combination: You mention the possibility of combining your approach with other efficient inference methods. Could you expand on this with results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730329134250

thread_type:
official_review

<thread_object>
creator: reviewer_kmjz

<content>
summary:
This paper proposes LOCRET, a novel framework for long-context LLM inference aimed at reducing GPU memory usage through trained retaining heads. Unlike existing static cache eviction methods, LOCRET uses lightweight training to estimate the causal importance of KV units, achieving more accurate cache eviction. The experimental results demonstrate memory efficiency and competitive generation quality with models like Llama-3.1-8B to perform 128K token inference on a single RTX 4090 GPU.

strengths:
The motivation of the paper is well-articulated, and the experiments are thoughtfully designed. Specifically:

- The claims are strongly supported by comprehensive experimental results. The framework addresses the core issue of KV cache growth through the use of retaining heads, with detailed benchmarks comparing LOCRET against several existing methods.
- The selective eviction strategy, guided by the use of CIS, is convincingly motivated. The experiments are well-structured, thoroughly exploring various datasets, models, and baselines, providing strong evidence of LOCRET’s effectiveness.
- The empirical evaluations comprehensively assess memory usage, inference speed, and performance across a diverse set of tasks. The results are consistently underpinned by sound theoretical analysis. Additionally, LOCRET facilitates long-context inference on GPUs like the Nvidia 4090, significantly enhancing the accessibility of advanced LLMs on consumer-grade hardware.

weaknesses:
The core idea of this paper is to develop an effective eviction policy through training retaining heads. However, several weaknesses need to be addressed:

- SirLLM is not an appropriate baseline for evaluating token eviction strategies. SirLLM is designed primarily for multi-turn conversations and is not tested on benchmarks like InfiniteBench or L-Eval. A more suitable baseline for eviction-based methods would be SnapKV [1]. 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.
- The benchmark suite lacks depth, particularly for information retrieval tasks. The retrieval task within InfiniteBench is overly simplistic, comprising repeated sentences that can be trivially discarded. I recommend that the authors incorporate experiments on RULER [2], following the MInference settings, to provide a more meaningful evaluation of retrieval performance.
- Token eviction based methods may struggle in multi-turn conversation scenarios. For example, in key-value retrieval tasks, if the user queries a different key-value pair during a subsequent turn, the model’s accuracy could degrade significantly due to missing context or prematurely evicted tokens.

If the authors address these concerns, I would consider raising my score.

[1] SnapKV: LLM Knows What You are Looking for Before Generation

[2] RULER: What's the Real Context Size of Your Long-Context Language Models?

questions:
- Could you clarify why there is a significant difference in performance between SirLLM and LOCRET in Table 3? If both methods operate under the same KV budget, the latency bottleneck should primarily stem from the attention operation. What factors contribute to LOCRET’s superior performance despite this similarity?
- Why is it necessary to keep the last $n_s$ caches? Could the retaining head detect and manage these recent tokens effectively? Does this indicate that the retaining head’s predictions are not sufficiently accurate for recent tokens, and if so, what improvements could address this limitation?
- How does LOCRET handle noisy datasets, such as conversational data with inconsistent or off-topic turns? Are there cases where retaining incorrect KV pairs causes irreparable errors during generation, and if so, how does the method mitigate such risks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730514145381

thread_type:
official_review

<thread_object>
creator: reviewer_j7cb

<content>
summary:
This paper proposes a training-based KV cache compression framework LOCRET for long-context LLM inference. The framework introduces retaining heads to evaluate the causal importance of KV cache units, allowing for more accurate eviction within a fixed cache size. The proposed framework is evaluated with two LLMs on Nvidia 4090 GPU.

strengths:
1.	This paper proposes a training-based KV cache compression framework LOCRET for selective KV cache eviction for long-context LLM inference. The proposed framework on two LLMs outperforms related methods on two LLMs and two benchmarks.
2.	The paper is easy to follow.

weaknesses:
1.	The paper claimed “LOCRET is also applicable to all transformer-based LLMs and various hardware”. However, 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).
2.	The proposed framework is validated with ∞Bench and L-Eval. How is the performance on other long-context benchmarks, such as longBench, et al. ?

questions:
Please refer to weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730547546609

thread_type:
official_review

<thread_object>
creator: reviewer_xs3n

<content>
summary:
To address the substantial overhead of KV cache in long-context reasoning with large language models, this paper introduces a novel method named LOCERT for KV cache pruning. LOCERT utilizes a more precise pruning metric called the causal importance score (CIS) to preserve the most significant KV cache entries.

strengths:
- The method proposes a lightweight training-based selective key-value cache eviction paradigm for long-context language model inference, with an offline training cost of less than 1 GPU hour.
- Extensive validation on various datasets confirms the superiority of our proposed method over the baselines discussed in the paper.
- An efficient inference system implementation is provided, integrating a retaining head mechanism into a segmented pre-filling inference framework. It maintains a fixed-size cache set by evicting cache units with low predicted importance, thereby controlling GPU memory usage.
- The paper discusses the inadequacies of existing methods such as KV quantization, which fail to address the overhead caused by linear growth in KV size. Our selection-based KV cache eviction method utilizes a static-sized KV cache and outperforms previous strategies in preserving important KV cache entries.

weaknesses:
- The proposed method requires additional training, and although the authors claim it only needs one hour, it also utilizes an eight-card A800 server, which is still resource-intensive.
- The novelty of the proposed method is modest. It is unclear why the training of heads to perform KV cache eviction, predicting each KV's importance, and using the causal importance score (CIS) for pruning, is superior to existing methods like H2O.
- The paper lacks a detailed analysis of the causal importance score (CIS) and needs a deeper discussion to explain why this metric effectively reflects the importance of KV cache.

questions:
- Regarding the use of a static-sized KV cache in selection-based KV cache eviction methods, can you explain why "the weakening correlation between local and global importance as sequences grow exacerbates this issue"?
- During training, the first loss term merely learns the maximum value of each column in the attention score. How effective would it be to directly use the maximum value of each column as a metric during inference?
- The paper mentions that methods like H2O cannot be effectively combined with KV quantization approaches. What are the actual performances of these methods?
- There are many papers similar to H2O that use attention score statistics for pruning, such as SnapKV and PyramidKV [2]. How does the method proposed in this paper compare with these approaches?
- Is the Stabilizer used only for selecting recent tokens?
- Is the performance improvement in this paper due to the SFT? What would be the effect if SFT were directly applied to the model?
- Should the number of heads in a retaining head be the same as in Query, or should it match Key/Value? If it matches Query, in structures like Grouped-Query Attention where each head's Key/Value corresponds to multiple heads' Query, how did you train this setup?

[1] LLM Knows What You are Looking for Before Generation
[2] Dynamic KV Cache Compression based on Pyramidal Information Funneling
</content>
</thread_object>

</thread>
<thread>
creation date:
1730579687365

thread_type:
official_review

<thread_object>
creator: reviewer_v3dc

<content>
summary:
The paper proposes LOCRET, an framework designed to enhance memory efficiency in long-context large language model (LLM) inference by using retaining heads to score and selectively retain key-value (KV) cache units. The primary challenge addressed is the high computational and memory demands posed by long-context LLM inference, which often limits deployment on consumer-grade devices. LOCRET introduces a trained retaining head mechanism that evaluates and prioritizes cache units based on their causal importance, offering a scalable and efficient approach that maintains inference quality on devices such as Nvidia 4090 GPUs. The paper conducts a comprehensive evaluation, comparing LOCRET with various memory-efficient inference baselines, demonstrating notable improvements in memory compression and inference quality without sacrificing speed.

strengths:
1. The paper presents a framework combining trained retaining heads with chunked prefill, contributing a distinctive approach to KV cache management in long-context inference. Unlike previous methods, LOCRET’s retaining heads learn a heuristic for cache importance, adapting to specific model architectures and sequence types, which provides greater flexibility across transformer-based LLMs.
2. The empirical evaluation is rigorous, with comparisons across a diverse set of baselines, including INFLLM, Quantization, SIRLLM, and MINFERENCE. The experiments cover both long and shorter context scenarios, supporting the paper’s claims of LOCRET’s superiority in maintaining performance while reducing memory usage.
3. LOCRET offers a good solution for deploying long-context LLM inference on consumer-grade hardware by significantly reducing the KV cache size without compromising quality. This contribution is valuable given the rising importance of long-context LLM applications in various fields.
4. The paper is well-organized, providing a clear explanation of LOCRET's architecture, training process, and the underlying intuition behind retaining heads. Diagrams effectively illustrate the framework and its mechanisms, enhancing reader understanding of the complex process of cache unit scoring and selective eviction.

weaknesses:
1. While the use of retaining heads to score and retain cache units is a valuable idea, the approach may benefit from further differentiation from existing token-dropping and quantization-based methods. Some parts of the scoring approach appear to overlap with traditional token importance estimation techniques (e.g., heavy-hitter approaches). A more comprehensive analysis highlighting LOCRET’s distinctions from similar heuristics in cache management would strengthen the contribution.
2. The results indicate promising efficiency gains but lack granular performance data on how LOCRET’s accuracy scales with different cache budgets across various architectures. 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.
3. Although LOCRET is tested across two LLM architectures, the applicability of this approach to a broader set of LLMs with diverse attention mechanisms (e.g., sparse attention) is not explored in depth. Discussing potential limitations or adjustments required for alternative models would enhance the generalizability of the method.

questions:
1. 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?
2. While SIRLLM performs poorly on memory-demanding tasks, it performs well on comprehension tasks. Could the authors comment on potential reasons LOCRET outperforms SIRLLM in these scenarios, particularly when both approaches manage memory through cache eviction?
3. Could the authors provide further insights or examples where the heuristic scoring might diverge significantly from the true causal importance? This would clarify the potential trade-offs in LOCRET's eviction policy.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732030401987

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revised Manuscript updated by Authors

comment:
We are very grateful to the reviewers for their detailed and thoughtful comments! We have responded to each comment and conducted additional experiments and clarifications. Based on the reviewers' feedback, we have updated and uploaded the revised manuscript. 

**We use blue color for all newly added content.** We will turn them back to black in the final version. Next, we will provide a one-on-one response to each reviewer, stating the specific locations in the paper where modifications have been made.

Due to the addition of new figures and tables, the numbering in the revised manuscript differs from the initial version. We provide a table comparing the old and new figure/table numbers for the reviewers' convenience when reading our response. In the response, we use the old numbering.

### **Figures**

| Old numbering | New numbering |
| :-: | :-: |
| Figure 1 | Figure 2 |
| Figure 2 | Figure 3 |
| Figure 3 | Figure 4 |
| Figure 4 | Figure 5 |
| Figure 5 | Figure 7 |
| Figure 6 | Figure 8 |
| Figure 8 | Figure 9 |

### **Tables**


| Old numbering | New numbering |
| :-: | :-: |
| Table 1 | Table 1 |
| Table 2 | Table 3 |
| Table 3 | Table 4 | 
| Table 4 | Table 9 |
| Table 5 | Table 10 |
| Table 6 | Table 2 |
| Table 7 | Table 6 |
| Table 8 | Table 12 | 
| Table 9 | Table 13 |
| Table 10 | Table 14 | 
| Table 11 | Table 15 |
</content>
</thread_object>

</thread>
<thread>
creation date:
1732389677664

thread_type:
official_comment

<thread_object>
creator: area_chair_m9ij

<content>
title:
Engage in Discussions Before Nov 26 (AoE)

comment:
Dear Reviewers,

First, let me thank you for your invaluable contributions to the ICLR review process. Your constructive feedback plays a key role in enhancing the quality of submissions.

---

As we approach the final days of the discussion phase (ending **Nov 26, 2024, AoE**), I kindly remind you to:

- Please take a moment to review the authors' responses to your comments. This is an opportunity to clarify any remaining questions, acknowledge misunderstandings, and refine your evaluation.

- If you need further clarification, don't hesitate to post your comments as soon as possible.

- If the authors' responses address your concerns or provide new insights, please consider updating your score to reflect this.

---

Your thoughtful participation during this phase is especially valuable for borderline papers, where additional input can be critical to ensuring a fair decision-making process.

I understand how busy this time of year can be and truly appreciate the time and care you dedicate to this important role. Your efforts make a tangible impact on the success of ICLR.

Thank you once again for your dedication.

Best regards,

Area Chair, ICLR 2025
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces LOCRET, a framework for memory-efficient long-context inference in LLMs. LOCRET employs lightweight retaining heads trained to predict the causal importance of KV cache units, enabling selective cache eviction. Using this mechanism with `chunked prefill`, LOCRET achieves efficient inference on consumer-grade GPUS while maintaining strong performance on long-context tasks.

## Strengths

- LOCRET introduces a query-independent mechanism for cache eviction, which addresses the bottleneck of KV cache growth in long-context inference (Reviewer WKXD). The paper enables long-context inference on consumer-grade hardware. 
- The paper includes detailed comparisons across multiple benchmarks and baselines (e.g., SnapKV, SirLLM, and H2O) and explores various hyperparameters like cache budget, stabilizer length, and chunk size (Reviewer WKXD, Reviewer V3dC).
- LOCRET consistently demonstrates strong memory efficiency and performance on a variety of tasks (Reviewer V3dC).
- The framework requires minimal training resources (less than 1 GPU hour) and makes only minor modifications to existing LLM architectures, ensuring diverse applicability (Reviewer xs3N, Reviewer WKXD).

## Weaknesses

- LOCRET has been evaluated primarily on decoder-only architectures and does not generalize to encoder-decoder models or alternative attention mechanisms, such as sparse attention (Reviewer WKXD, Reviewer j7cb).
- Poor performance on complex query-driven tasks like RULER indicates that LOCRET struggles without query-awareness, limiting its applicability in certain retrieval scenarios (Reviewer KmJZ, Reviewer xs3N).
- While the experiments are extensive, some reviewers noted a lack of focus on challenging benchmarks or real-world datasets (Reviewer KmJZ). For example, InfiniteBench’s simplicity and the lack of off-topic turns in multi-turn conversation datasets reduce the impact of the results.
- SnapKV and similar baselines were tested under chunked prefill scenarios, which are not compatible with their design, potentially inflating LOCRET’s relative performance (Reviewer KmJZ).
- LOCRET exhibits performance degradation on complex retrieval tasks like multi-key RULER subtasks, raising concerns about its robustness and generalizability (Reviewer KmJZ).

Based on the reviews and the rebuttal, I recommend **Reject** for this paper. The paper tackles an important problem, however, the paper in its current form has few issues that seem not to be properly addressed in the rebuttal:

- LOCRET fails to generalize beyond decoder-only architectures, struggling with encoder-decoder models and alternative attention mechanisms (Reviewer WKXD, Reviewer j7cb).
- Its poor performance on query-driven benchmarks like RULER, especially on complex multi-key tasks, raises concerns about its applicability in real-world scenarios (Reviewer KmJZ). The results suggest that LOCRET lacks the generalizability needed to handle tasks where query-awareness is critical.
- While the authors run extensive evaluations, some of the chosen benchmarks (e.g., InfiniteBench) are overly simplistic and fail to challenge the proposed method sufficiently (Reviewer KmJZ).
- The comparisons with baselines like SnapKV under chunked prefill scenarios are unfair, as these methods are not designed for such setups (Reviewer KmJZ). 
- The proposed method, while practical, does not demonstrate significant contribution over existing approaches like SnapKV and H2O. The core idea of training retaining heads to predict cache importance lacks sufficient theoretical development and differentiation from prior work (Reviewer xs3N, Reviewer WKXD).
- The paper does not convincingly address challenges in query-driven or multi-turn conversation tasks. LOCRET struggles with off-topic or inconsistent input contexts, and the experiments lack robust evaluation in such scenarios (Reviewer KmJZ).