PaperlayerView exact revision on arXiv

Source-linked AI summary

ReFreeKV: Towards Threshold-Free KV Cache Compression

Xuanfan Ni, Liyan Xu, Chenyang Lyu, Longyue Wang, Mo Yu, Lemao Liu, Fandong Meng, Jie Zhou, Piji Li

arXiv:2502.16886v4cs.CLcs.AIcs.LG

TL;DR

KV cache pruning often requires input-specific budget thresholds, limiting its use across diverse open-domain inputs. ReFreeKV introduces threshold-free pruning and achieves near-lossless compression robustly across 13 datasets without input-specific thresholds.

  • Problem

    Prior KV cache pruning methods typically require data-dependent budget thresholds, creating an unresolved challenge for robust compression across diverse inputs.

  • Method

    ReFreeKV uses a two-stage pruning process with an input-insensitive Uni-Metric stopping criterion that dynamically allocates KV cache budgets.

  • Results

    63.7% average KV budget slightly exceeded full-cache performance across 13 datasets for Llama3-8B, without an input-specific threshold.

  • Takeaways & Limitations

    ReFreeKV demonstrates near-lossless, input-insensitive KV compression across varied tasks, context lengths, and LLM models.

  • Takeaways & Limitations

    ReFreeKV lacks formal guarantees on performance degradation and incurred a 1.5% degradation on Mistral-7B.

Abstract

from arXiv · show

To reduce memory consumption during LLM inference, a handful of methods have been proposed for KV cache pruning. While these techniques can accomplish lossless memory reduction on many datasets, they often hinge on an under-emphasized condition: an input/domain-specific threshold for KV cache budget needs to be pre-determined to achieve the optimal performance. However, such input-sensitive design may be considerably limited in real-world scenarios, as open-domain inputs span diverse domains, lengths and difficulty levels, without clear boundaries for threshold selection. As a result, the dependence of such input-sensitive threshold can be a fundamental limitation that causes large degradation on arbitrary inputs. In this work, we propose a new objective that lifts the threshold constraints for robust KV compression, advocating for "threshold-free" methods that adaptively adjust budget allocation while preserving full-cache performance. We then propose a novel method, ReFreeKV, serving as the first instantiation of this objective. Extensive experiments across 13 datasets with diverse context lengths, task types, and model sizes demonstrate its efficacy and efficiency. Our code is publicly released at https://github.com/Patrick-Ni/ReFreeKV.

1 Introduction

KV cache reduction methods commonly require input-specific budget thresholds, which are difficult to preset across diverse real-world inputs and can cause performance degradation. ReFreeKV addresses this limitation with threshold-free budget adaptation designed to preserve full-cache performance across varied tasks and models.

  • Motivation: Llama3-70B can require up to 50GB of memory for 20K tokens, making efficient KV cache management crucial as model size and input length grow.The KV cache stores intermediate states retrieved for attention during autoregressive decoding.
  • Problem: Prior KV reduction methods rely on data-dependent budget thresholds, with D2O needing 20% on LongBench but experiments indicating 80% on GSM8K for full performance.These thresholds tune how many cache positions are retained after prefilling.
  • Problem: Because real-world inputs mix domains, lengths, and difficulty levels, optimal thresholds cannot be predetermined, making systems less robust and prone to significant performance degradation.A threshold that reaches full performance on one dataset may not transfer well to others, and performance varies substantially as the budget ratio changes.
  • Objective: The proposed objective is a universal input-insensitive threshold and performance consistently comparable to full-cache inference, with dynamic KV budget adjustment prioritized over maximal compression.This reframes KV pruning around robust arbitrary-input handling rather than memory savings on a fixed benchmark.
  • Method: ReFreeKV uses a two-stage process that ranks KV positions, progressively retains them until a stopping criterion is met, and implements the procedure with parallel operators to minimize overhead.Its Uni-Metric is designed to remain consistent and insensitive across input domains and sequence lengths; threshold-free means no pre-determined threshold calibration is needed.
  • Evaluation: Across 13 datasets spanning context lengths, tasks, and model sizes, ReFreeKV achieves highly comparable inference and can surpass full-cache performance without an input-specific threshold.The evaluation includes mathematical and commonsense reasoning, reading comprehension, and coding with multiple LLM sizes.

2 Related Work

Prior KV cache pruning methods reduce inference memory by selectively discarding cache positions using attention, recency, historical importance, or per-head retention strategies. However, recent approaches may still depend on input-sensitive hyperparameters, limiting stability across tasks.

  • KV Cache Pruning: KV cache pruning reduces LLM inference memory by selectively discarding less important cache positions.Scissorhands and SnapKV retain tokens based on high attention scores.
  • KV Cache Pruning: Attention-based methods such as Scissorhands and SnapKV retain tokens with high attention scores.These methods represent a prominent line of work in KV cache pruning.
  • KV Cache Pruning: Other methods use token recency and historical importance, while FastGen adapts retention strategies per attention head.H2O and StreamingLLM use recency or historical importance; FastGen refines retention on a per-head basis.
  • Input-Sensitive Design in Recent Works: Recent approaches replace fixed budgets with heuristic halting conditions but often remain constrained by input-sensitive hyperparameters.The passage identifies sparse and recent ratio as an example of such dependence in Lethe, with limited emphasis on cross-task stability.

3 Methodology

ReFreeKV targets threshold-free KV compression by removing input-specific budget selection while preserving near-full-cache performance across arbitrary inputs. Its two-stage design ranks positions, then uses an input-insensitive Uni-Metric stopping rule with efficient parallel implementation and practical safeguards.

  • Motivation: Prior KV compression methods depend on pre-selected thresholds whose optimal values can vary across inputs, limiting stable performance on arbitrary open-domain instructions.This motivates removing input-specific threshold constraints rather than prioritizing compression ratio first.
  • Objective: ReFreeKV’s objective is threshold-free compression that maintains performance comparable to full-cache regardless of inputs while eliminating optimal-threshold tuning.The method is presented as the first effective solution claimed to fulfill this objective.
  • Two-Stage Method: ReFreeKV first ranks KV positions per layer and attention head, then sequentially retains vectors until an input-insensitive Uni-Metric stopping condition is reached.The remaining KV cache positions are discarded after the stopping point.
  • Initial Ranking: The initial ranking takes the first m positions and reverses the remaining n − m positions, with m chosen to work well regardless of the input sequence.This position-based ranking is reported as effective and computationally advantageous compared with other ranking strategies.
  • Uni-Metric Eviction: The Uni-Metric uses attention-matrix Frobenius-norm differences to signal pruning degradation, selecting T = 1% as a universal threshold that balances minimal degradation with maximal cache eviction.Preliminary studies indicate that T < 1% remains robustly comparable to full-cache performance.
  • Efficient Implementation: Parallel PyTorch operators identify stopping positions without explicit loops, while ReFreeKV always retains the full KV cache of the first two layers for simplicity and robustness.The method is applied once after input prefilling, targeting reduced inference memory consumption and improved throughput.

4 Experiments

Across 13 diverse datasets and multiple LLM families, ReFreeKV is evaluated for near-lossless, threshold-free compression, robustness across budgets and model sizes, and inference efficiency. It dynamically adapts KV-cache budgets to preserve full-cache performance while improving compression and generation throughput.

  • Experimental Setup: Experiments cover Llama3-Instruct, Mistral-7B-Instruct-V0.3, and Qwen2.5-Instruct across model sizes from 7B to 72B and 13 datasets spanning diverse tasks and context lengths.The datasets include mathematics, science, commonsense reasoning, document comprehension, summarization, and coding.
  • Robustness to Budget Selection: At 90%, 50%, and 20% fixed budgets, prior pruning methods can suffer severe dataset-specific degradation, whereas ReFreeKV automatically adjusts pruning to prioritize full-cache performance.ReFreeKV also outperforms three 90%-budget baselines while using less than 90% budget.
  • Main Results: ReFreeKV achieves near-lossless dynamic compression across models, input lengths, and task types, surpassing full-cache performance by 0.12% on Llama3-8B and 2.63% on Qwen2.5-7B.The corresponding average KV-cache utilizations are 63.68% and 76.02%; Mistral achieves near 15% compression with a 1.5% performance reduction.
  • Dynamic Budget Allocation: ReFreeKV allocates over 90% budget on Math&Science datasets and as little as 15% on QA or Summarization datasets, reflecting task difficulty.Hard, concise tasks require more context and precise calculation, while easier or longer-context tasks permit greater compression.
  • Comparison with Twilight: ReFreeKV achieves performance comparable to Twilight on GSM8K, NarrativeQA, 2WikiMQA, and Musique, with both methods frequently surpassing the full-cache baseline.The comparison uses Twilight’s reported optimal p values: 0.95 for Llama3-8B and 0.85 for Mistral.
  • Efficiency: ReFreeKV attains the best overall generation time in 8 out of 12 comparisons, improves throughput over naive generation by 10–20%, and adds minimal pruning latency.Its lightweight two-stage design limits overhead, while the reduced KV cache lowers attention computation costs; throughput gains persist as batch size increases.
  • Ablations: Using a 1% threshold provides a robust trade-off, while 10% causes substantial degradation and attention-score ranking fails to maintain robust final performance.Reducing the threshold to 0.1% increases average budget with negligible performance improvement; on NarrativeQA, 10%n scores 21.44 using 48.7% budget versus 9.74 using 32% for the alternative.
  • Generalization: On Llama3-70B and Qwen2.5-32B/72B, ReFreeKV maintains near full-cache performance without threshold tuning, while compression approaches 50% on longer-context datasets.The same configuration used in Table 2 generalizes across model sizes and datasets.

5 Conclusion

The study introduces a threshold-independent KV cache compression objective for input-insensitive pruning and robust inference. ReFreeKV implements this objective through a two-stage pruning process and achieves near-lossless performance across diverse experimental settings.

  • The proposed KV cache compression objective lifts threshold dependency to enable input-insensitive pruning for robust inference performance.
  • ReFreeKV implements the objective with a straightforward yet effective two-stage KV cache pruning process.
  • Comprehensive experiments across diverse datasets, tasks, context lengths, and LLM models demonstrate that ReFreeKV achieves near-lossless performance.

Limitations · A Full Algorithm of ReFreeKV

ReFreeKV’s limitations include a gap between achieved and optimal compression budgets and a lack of formal performance guarantees. Its full algorithm takes a prompt and threshold, processes each Transformer layer, and returns compressed key-value caches.

  • Limitations: 84.3% budget retention on QMSum with Mistral-7B exceeds the 50% budget that remains viable without performance degradation.The authors identify this gap between achieved compression and the true optimal budget as a primary limitation.
  • Limitations: More aggressive KV-cache compression under dynamic budgets remains a future direction for satisfying the full-cache objective.This direction is proposed to address the gap between current and optimal budgets.
  • Limitations: 1.5% degradation occurs on Mistral-7B, while ReFreeKV surpasses full-cache performance for Llama3-8B and Qwen2.5-7B.Although compression is empirically near-lossless, the method lacks formal guarantees on performance degradation.
  • A Full Algorithm of ReFreeKV: The algorithm accepts a prompt and threshold t, outputs a compressed KV cache, and initializes empty lists Kc and Vc.These are the stated input, output, and initialization steps.
  • A Full Algorithm of ReFreeKV: For each Transformer layer Li in the LLM, the algorithm processes the layer before continuing to cache construction.The provided pseudocode explicitly iterates over Transformer layers.
  • A Full Algorithm of ReFreeKV: The procedure ranks Ai and uses the resulting selection to build the compressed cache lists Kc and Vc.The excerpt specifies ranking Ai and adding selected content to Kc and Vc.
  • A Full Algorithm of ReFreeKV: After processing all layers, the algorithm returns Kc and Vc as the compressed key-value cache.The pseudocode ends by returning the two constructed cache lists.

B Retain Bottom LLM Layers

ReFreeKV begins pruning from the third LLM layer, retaining the full KV cache in the first two layers. Without freezing those layers, examples on GSM8K and CoQA become incorrect and incoherent despite overall budgets exceeding 90%.

  • Method choice: ReFreeKV’s KV-cache pruning begins at the third LLM layer, retaining the first two layers.The method applies pruning from layer 3 onward.
  • Case study: At a 1% threshold with no frozen layers, GSM8K and CoQA outputs are incorrect and lack logical coherence despite overall budgets exceeding 90%.Layers 1 and 2 receive relatively low budgets, while layers 3–31 receive uniformly high budgets and the final layer drops again.
  • Case study: Layer 1: 67.71 and Layer 2: 82.29 contrast with Layer 3∼31: 95.83, while Layer 32: 37.50 and Avg.: 93.90.The per-layer budget pattern accompanies the observed failures when the first two layers are not frozen.

CoQA

On CoQA, ReFreeKV reports high average retention across most layers but exhibits a generation failure in the case study. Freezing the first two layers is identified as a balance between performance and budget, while retaining layer 31 adds negligible gains at higher cost.

  • CoQA: The case study’s output repeats “What is the question:” instead of answering the CoQA question about Cotton’s color.The displayed output is associated with a generation failure in the case study.
  • CoQA: 95.03 Avg. is reported, with 99.19 for Layers 3∼31 and lower values for Layers 1, 2, and 32.The reported values are Layer 1: 42.14, Layer 2: 42.54, Layer 3∼31: 99.19, and Layer 32: 79.64.
  • CoQA: Freezing the first two layers balances model performance and budget, whereas retaining layer 31 provides negligible gains at higher budget.The authors evaluate different layer-freezing strategies and summarize partial results in Table 8.

C Datasets Used in Experiments

The experiments cover mathematical and scientific problem solving, commonsense reasoning, document question answering, summarization, few-shot learning, and code generation across diverse datasets. Metrics include Accuracy, BLEU, F1 Score, Rouge-L, and Edit Sim.

  • Math & Science: Math & Science uses GSM8K and GPQA to evaluate mathematical and scientific problem solving with Accuracy.GSM8K contains 8,000 elementary-level math word problems, while GPQA tests physics understanding and problem solving.
  • Commonsense Reasoning (CR): Commonsense Reasoning uses TruthfulQA and CoQA to assess factuality and conversational question answering with BLEU.TruthfulQA tests distinguishing fact from fiction, while CoQA focuses on responding to questions in conversational contexts.
  • Single Document QA (Single-Doc QA): Single-Doc QA uses NarrativeQA and Qasper to evaluate comprehension and question answering over extended narrative and academic documents.NarrativeQA focuses on stories and underlying themes, whereas Qasper evaluates information extraction from academic papers.
  • Multi-Document QA (Multi-Doc QA): Multi-Doc QA uses 2WikiMultiHopQA and MuSiQue to test multi-hop reasoning across documents with F1 Score.Both datasets require integrating information from multiple sources, including Wikipedia articles, to answer complex questions.
  • Summarization: Summarization uses QMSum and Multi-News to assess concise and coherent summaries of meeting transcripts and multiple news articles with Rouge-L.QMSum emphasizes key points from multi-party discussions, while Multi-News requires synthesizing articles on the same topic.
  • Few-Shot Learning (FSL) and Code: Few-Shot Learning uses TriviaQA with F1 Score, while Code uses LCC with Edit Sim to evaluate code completion and generation.TriviaQA emphasizes retrieval and factual recall from large trivia collections; LCC evaluates understanding and generation over extended code contexts.

D Ablation

The ablation study examines how different k settings for reducing the full attention matrix affect model performance and pruning time. It identifies k = 1 as an efficient design choice that achieves better performance with a reduced budget.

  • D Ablation: The study evaluates various k values for reducing the full attention matrix and extends the results reported in Table 5 with Table 10.These experiments assess how parameter settings affect model performance and inform parameter selection.
  • D Ablation: The ablations provide a basis for understanding how parameter settings impact model performance and for optimizing parameter selection.
  • D Ablation: k = 1 conserves pruning time while achieving better model performance with a reduced budget, making it an ideal design choice.

E Comparison with StreamingLLM

Experiments on six datasets with Llama3-8B compare ReFreeKV against StreamingLLM while freezing the first two layers. StreamingLLM cannot consistently preserve optimal performance across datasets under varying KV-cache budgets, indicating that position bias alone is insufficient.

  • Comparison setup: Experiments evaluate StreamingLLM with Llama3-8B on six datasets, freezing the first two layers to assess the proposed two-stage process.This setup matches ReFreeKV’s layer-freezing strategy for the ablation.
  • Performance comparison: StreamingLLM fails to consistently maintain optimal performance across datasets such as GSM8K and GPQA under KV-cache budgets of 0.9, 0.5, and 0.2.The comparison uses the first two layers frozen for both methods.
  • Performance comparison: Relying solely on position bias, as in StreamingLLM, is insufficient to preserve full performance.This conclusion follows from StreamingLLM’s inconsistent performance across datasets and budgets despite adopting the same layer-freezing strategy as ReFreeKV.
Loading 2502.16886v4…