PaperlayerView exact revision on arXiv

Source-linked AI summary

RestoreKV: Recovering Full-Cache Behavior Under Aggressive Query-Agnostic KV Cache Eviction

Changwoo Baek, Seungjun Shin, Kyeongbo Kong

arXiv:2608.01247v1cs.CLcs.LG

TL;DR

Query-agnostic KV eviction retains only selected original KV pairs, but aggressive compression can sharply degrade performance. RestoreKV adds a learned, context-conditioned restore cache within the same budget and consistently narrows the gap to full-cache performance across models, benchmarks, and eviction methods.

  • Problem

    Existing query-agnostic eviction methods primarily ask which original KV pairs to retain, leaving degradation under aggressive compression insufficiently addressed.

  • Method

    RestoreKV generates a compact, context-conditioned restore cache before eviction using a shared LoRA-adapted restoration mechanism under the same total KV budget.

  • Results

    RestoreKV consistently narrows the gap to full-cache performance across four backbones, four benchmarks, and five eviction methods; at 5% budget, KVzip rises from 38.2 to 73.2 on RULER-4K.

  • Takeaways & Limitations

    RestoreKV is a broadly compatible plug-in that improves existing eviction pipelines, with larger gains typically observed under tighter cache budgets.

  • Takeaways & Limitations

    RestoreKV generates its restore cache without observing the query, preserving query-agnostic cache construction rather than adapting construction to query-specific information.

Abstract

from arXiv · show

Query-agnostic KV cache eviction compresses a context once and reuses the resulting cache for arbitrary future queries, but performance can collapse under tight budgets. Existing methods primarily improve which original KV pairs are retained. We introduce RestoreKV, which complements this selection-based formulation with learned restoration under the same total KV budget. Our key insight is that, although the information lost through eviction is context-specific, the mechanism for generating its compact complement can be shared across contexts. After context prefill, a few restore tokens attend to the full KV cache in a single LoRA-adapted pass, generating a compact, context-conditioned restore cache. The base importance scorer and eviction rule remain unchanged, and the adapters are disabled for all subsequent queries and decoding. RestoreKV is trained through parameter-efficient self-distillation from the frozen full-cache model, optimizing only $0.4\%$ of the parameters and requiring no task-specific tuning. Across four backbones and four long-context benchmarks, RestoreKV substantially reduces compression-induced degradation. On Qwen3-4B, it improves 59 of 60 paired, budget-matched settings across five base eviction methods; at a $5\%$ budget, it raises KVzip from $38.2$ to $73.2$ on RULER-4K. Applied to KVzip+, RestoreKV reaches $86.4$ RULER accuracy at $16\times$ compression on the KVPress Benchmark, while adding less than $0.5\%$ one-time cache-construction overhead in a 32K-context evaluation. Our project page is available at https://paper.pnu-cvsp.com/RestoreKV/

1 Introduction

RestoreKV addresses the limits of query-agnostic KV eviction by generating a compact, context-conditioned restore cache alongside retained original states under the same KV budget. It narrows compression-induced degradation across models and benchmarks while adding minimal one-time construction overhead.

  • Motivation: Query-agnostic eviction compresses a context before future queries are known, leaving the central problem of selecting which KV pairs survive a fixed budget.The resulting cache contains only retained original states.
  • Results: 86.4 RULER accuracy is reached at 16× compression with KVzip+ on the KVPress Benchmark, while the 32K-context restore pass adds less than 0.5% of total compression time.The pass adds 0.03–0.04 seconds and 84 MB (0.4%) of peak memory while preserving query-time KV budget and decoding cost.
  • Method: RestoreKV complements retention with a shared mechanism that generates a compact, context-conditioned restore cache before eviction without changing the base scorer or eviction rule.The restore cache occupies part of the fixed budget, while retained context states fill the remaining slots.
  • Method: RestoreKV uses a single LoRA-adapted restore pass after prefill, then disables its adapters for subsequent query processing and decoding.Eight restore tokens attend to the full KV cache by default, generating restore states that combine with retained context states.

2 Related Work

Prior work estimates query-agnostic KV importance, learns lightweight estimators, or expands compressed-cache representations. RestoreKV instead preserves the base eviction pipeline while allocating part of the fixed budget to a shared mechanism for context-specific restoration.

  • Query-agnostic eviction: Query-dependent eviction methods cannot construct a compressed cache before future queries are observed, motivating query-agnostic importance estimation.Query-agnostic methods estimate individual KV-pair contributions through context reconstruction, while KVzip+ and ContrastKV refine importance using output-norm weighting and contrastive objectives.
  • Learned importance estimation: Lightweight trainable modules amortize or improve importance estimation by predicting importance scores, learning sink-attention gates, or using selectively activated LoRA modules.KVzap predicts KVzip+ importance from hidden states; Fast KVzip learns sink-attention gates through context reconstruction; LookaheadKV predicts response-induced importance in the query-aware setting.
  • Expanded cache representations: Other approaches expand the compressed cache representation by rewriting selected-token states or optimizing context-specific parameterized KV caches.KV-Distill jointly learns importance scoring and parameter-efficient adapters using full-cache predictive distributions for distillation, whereas Cartridges optimizes a context-specific parameterized cache through self-study.
  • Concurrent directions: Concurrent methods construct compact cache states through attention matching, value reconstruction, or context compilation with disposable LoRA modules.Attention Matching matches per-head attention outputs and mass; VECTOR reconstructs values with an offline-calibrated linear map; Latent Context Compilation uses buffer-token KV states.
  • RestoreKV’s design point: RestoreKV fixes the base importance scorer and eviction rule, preserves most selected KV pairs, and reserves part of the same budget for a compact, context-conditioned restore cache.Its premise is that eviction losses are context-specific while the mechanism generating their compact complement can be shared across contexts.

3 Method

RestoreKV augments query-agnostic KV eviction with a compact, context-conditioned restore cache generated from the full cache before eviction, while preserving the base scorer, allocation rule, and total KV budget. Shared restore-token embeddings and LoRA parameters are trained by self-distillation, but LoRA is used only during one-time cache construction, keeping subsequent queries query-agnostic.

  • Core design: RestoreKV combines a generated context-conditioned restore cache with retained context KV pairs under the same total budget, without changing the base importance scorer or eviction rule.The method reserves part of the fixed budget for restoration and lets the base evictor fill the remainder.
  • Restore-cache generation: A single LoRA-adapted restore pass lets shared restore tokens attend causally to the full KV cache and generate a context-specific restore cache.The restore-token embeddings and LoRA parameters are shared across contexts, whereas the generated restore cache depends on the current context.
  • Budget matching: Each of n restore tokens contributes one KV pair per layer and KV head, yielding |Cres| = nLH and allowing restoration to incorporate information from KV pairs later removed.RestoreKV reserves nLH of the B cache slots for Cres and gives the base evictor the remaining B − nLH slots.
  • Training objective: Self-distillation trains only the restore-token embeddings and LoRA parameters against the frozen full-cache model, while the backbone and base importance scorer remain frozen.The teacher generates an answer from the full cache, and teacher and student evaluate that same answer using the full and budget-matched caches.
  • Inference: The query and teacher answer define only the offline distillation target; restore-cache construction itself does not observe the query, preserving query-agnostic reuse across arbitrary subsequent queries.At inference, LoRA is enabled only for the n-token restore pass and disabled for all later query processing and decoding.

4 Experiments

Experiments evaluate RestoreKV across four backbones, four long-context benchmarks, and five query-agnostic eviction methods under matched KV budgets. RestoreKV consistently mitigates compression loss, with gains explained by learned attention-side restoration and achieved at negligible one-time overhead.

  • Evaluation setup: Four backbones are tested, while only restore-token embeddings and LoRA adapters train; the frozen backbone and base eviction rules remain unchanged.The models are Qwen3-0.6B/4B/8B and Llama-3.1-8B-Instruct, and the restore pass is reused across downstream tasks without task-specific tuning.
  • Evaluation setup: Four benchmarks span synthetic retrieval, tracing, and aggregation; scientific-paper QA; and additional long-context query types.The evaluation uses RULER at fixed 4K context, QASPER at approximately 5K tokens, and two further benchmarks described in the experiments.
  • Main results: 73.2 versus 38.2: at r=0.05 on Qwen3-4B RULER-4K, RestoreKV improves KVzip under the same total KV budget.RestoreKV+ likewise improves KVzip+ from 51.6 to 70.7 under the same budget.
  • Generalization across eviction methods: Five eviction methods improve with RestoreKV, with larger gains typically at tighter budgets, demonstrating compatibility across different importance criteria and eviction procedures.The methods are KVzip, KVzip+, ContrastKV, and query-agnostic SnapKV and H2O variants.
  • Ablations: 96% of RestoreKV’s full improvement is recovered by fixed embeddings with LoRA, reaching 71.9 versus 73.2 at r=0.05; q/k/v-only LoRA reaches 72.4.These results identify attention-side adaptation as the primary driver of the context-conditioned restore cache.
  • Efficiency: <0.5% additional end-to-end compression time: at 32K contexts, restore-cache construction adds 0.03–0.04 seconds and 84 MB peak prefill memory, with no measurable inference-time cost.The cache size is exactly matched between methods, decreasing from 4.20 to 0.20 GB.

5 Conclusion … A.2 Effect of the Query-Position Offset

RestoreKV combines retained KV pairs with a compact, context-conditioned restore cache generated before eviction, preserving the base eviction pipeline without query-time memory or decoding cost. The appendix reports compatibility with learned scorers and controls for query-position offsets, alongside additional benchmark and reproducibility results.

  • 5 Conclusion: RestoreKV generates a compact, context-conditioned restore cache before eviction and combines it with retained KV pairs under the same total budget.It uses a shared LoRA-adapted restoration mechanism while preserving the base importance scorer and eviction rule.
  • 5 Conclusion: Adapters are disabled after cache construction, adding negligible one-time overhead and no query-time KV-memory or decoding cost.This keeps restoration confined to the prefill and cache-construction stage.
  • Appendix Overview: The appendix covers additional experimental results, control experiments, implementation details, benchmark results, and reproducibility details.Its listed sections include learned-scoring applicability, query-position offsets, evictor mismatch, training ratios, LongBench, SCBench, training setup, baseline adaptations, environment, seeds, and metrics.
  • A Additional Experimental Results: The additional experimental results include applicability to a learning-based scorer, query-position offsets, training–inference evictor mismatch, and training ratio ranges.These topics are enumerated under Section A.
  • A.1 Applicability to a Learning-Based Scorer: 46.7 to 79.3: attaching RestoreKV raises Fast KVzip’s RULER-4K accuracy at 16× compression, a 32.6-point gain exceeding standard KVzip’s 59.0.Fast KVzip uses lightweight per-layer gates to predict importance from hidden states in a single forward pass, and retains single-pass scoring.
  • A.2 Effect of the Query-Position Offset: RestoreKV appends 8 restore tokens, shifting subsequent queries 8 positions later and changing their relative RoPE offsets to retained context KV pairs.The control variant removes restore tokens and restore cache while applying the same +8 query-position shift to isolate this effect.

A.3 Effect of Training–Inference Evictor Mismatch

RestoreKV is evaluated under training–inference evictor mismatch by fixing inference to KVzip and comparing checkpoints trained with KVzip or SnapKV. Matching becomes increasingly important under aggressive eviction, although mismatched restoration still provides a small benefit at the tightest budget.

  • A.3 Effect of Training–Inference Evictor Mismatch: At 16× compression, RestoreKV raises Fast KVzip from 46.7 to 79.3, exceeding plain KVzip at 59.0.This composition result uses RULER-4K with Qwen3-8B.
  • A.3 Effect of Training–Inference Evictor Mismatch: Inference fixes KVzip while RestoreKV checkpoints are trained with either KVzip or SnapKV to test training–inference evictor mismatch.The experiment directly examines whether a checkpoint trained with one base evictor remains effective with another.
  • A.3 Effect of Training–Inference Evictor Mismatch: At r=0.10, SnapKV-trained RestoreKV reaches 81.6 versus 88.8 for KVzip-trained RestoreKV; at r=0.05, it reaches 42.0 versus 73.2.Matched and mismatched checkpoints perform similarly at r=0.20, but the gap widens as eviction becomes more aggressive.
  • A.3 Effect of Training–Inference Evictor Mismatch: At r=0.05, mismatched RestoreKV still reaches 42.0 versus 38.2 for the no-restoration KVzip baseline.This indicates a limited benefit from restoration even when the training evictor does not match inference.

A.4 Effect of the Training Ratio Range · B Additional Benchmark Results

Training-range choices have little effect at a mild RULER-4K budget but matter increasingly under aggressive compression. The default U(0.025, 0.25) range performs best, reaching 88.8 at r=0.1 and 73.2 at r=0.05.

  • A.4 Effect of the Training Ratio Range: The study varies the sampled KV budget ratio as r ∼U(rmin, rmax) on RULER-4K.
  • A.4 Effect of the Training Ratio Range: At the mild evaluation budget r=0.4, all training-range configurations achieve comparable performance.
  • A.4 Effect of the Training Ratio Range: As the evaluation budget tightens, performance differences between training ranges become increasingly pronounced.
  • A.4 Effect of the Training Ratio Range: The default training range U(0.025, 0.25) achieves the strongest overall performance.
  • A.4 Effect of the Training Ratio Range: 88.8 is achieved by the default range at the evaluation budget r=0.1.
  • A.4 Effect of the Training Ratio Range: 73.2 is achieved by the default range at the tighter evaluation budget r=0.05.
  • A.4 Effect of the Training Ratio Range: Increasing either the lower or upper bound of the training range reduces performance under aggressive evaluation budgets.

B.1 LongBench Results

RestoreKV was evaluated on all 16 LongBench tasks under the KVPress protocol, covering 3,750 examples across approximately 5K–15K-token contexts. It improved average performance across both backbones and cache budgets, with larger gains at r=0.0625.

  • B.1 LongBench Results: RestoreKV was evaluated on all 16 LongBench tasks using the KVPress protocol.The evaluation comprised 3,750 examples with context lengths ranging from approximately 5K to 15K tokens.
  • B.1 LongBench Results: RestoreKV improved the overall average across both backbones and both cache budgets.Results were aggregated by task category in Table D.
  • B.1 LongBench Results: At r=0.0625, the average score increased from 33.5 to 37.7 on Llama-3.1-8B-Instr.The passage reports larger gains at this tighter cache budget.

B.2 SCBench Results · C Implementation and Reproducibility Details · C.1 Training Data Construction

RestoreKV generalizes to approximately 104K-token SCBench contexts, improving KVzip averages across all tested budgets despite training only up to 15K tokens. The paper also documents its training-data construction and experimental environment for reproducibility.

  • B.2 SCBench Results: 29.7 vs. 29.4: RestoreKV raises KVzip’s SCBench average at r=0.4.SCBench covers nine tasks averaging approximately 104K tokens.
  • B.2 SCBench Results: 33.1 vs. 31.6: RestoreKV raises KVzip’s SCBench average at r=0.2.These gains generalize roughly an order of magnitude beyond the training lengths.
  • B.2 SCBench Results: 32.4 vs. 30.0: RestoreKV raises KVzip’s SCBench average at r=0.1.RestoreKV is trained only on contexts up to 15K tokens.
  • B.2 SCBench Results: 25.6 vs. 24.1: RestoreKV raises KVzip’s SCBench average at r=0.05.The evaluation follows KVzip’s protocol across retrieval and global-context understanding tasks.
  • C.1 Training Data Construction: 2,488 examples: LongAlpaca self-study documents yield filtered question-answer examples for RestoreKV training.The teacher generates five queries per document spanning factual, summarization, multihop, method, and comparison questions.
  • C.1 Training Data Construction: 6,248 total samples span contexts from 512–15,237 tokens, according to RestoreKV’s training-data statistics.The training set combines three sources, including LongAlpaca and PG-19 self-study.
  • C Implementation and Reproducibility Details: Experiments use an NVIDIA RTX PRO 6000 Blackwell with 96 GB, an Intel Xeon Gold 6530, and 503 GiB memory.The reported software environment is Ubuntu 22.04.5 LTS with PyTorch 2.8.0 and CUDA 12.8.

C.2 Training Setup Details · C.3 Baseline Adaptations (SnapKV and H2O). · C.4 Experimental Environment

The paper trains a small restore-token and LoRA component while freezing the base model, evaluates benchmarks with specified task metrics, adapts SnapKV and H2O query-agnostically, and uses the official KVPress environment for fair comparisons.

  • C.2 Training Setup Details: 0.4% of a 4B model is trained: n=8 restore-token embeddings and LoRA adapters, while the base model remains frozen.The adapters use rLoRA=8, α=16, dropout 0, and target attention and MLP projections.
  • C.2 Training Setup Details: 5,000 steps use AdamW with learning rate 2 × 10−4, cosine scheduling, 50 warmup steps, gradient clipping at norm 1.0, and one context per batch.AdamW uses β=(0.9, 0.999) and weight decay 0.01.
  • C.2 Training Setup Details: RULER-4K uses official string-matching, QASPER uses maximum token-F1, and QuALITY and LongHealth use answer-letter accuracy.LongBench uses its official per-task metric, while SCBench uses its official task-specific evaluator.
  • C.3 Baseline Adaptations (SnapKV and H2O).: Both SnapKV and H2O compress the context cache before future queries are observed and evaluate importance at the individual KV-pair level under the same budget framework.These are the query-agnostic, prefill-based adaptations used in the KVzip evaluation setup.
  • C.3 Baseline Adaptations (SnapKV and H2O).: SnapKV scores KV pairs from a final w=32 context-position observation window, averaging attention and applying max-pool smoothing with kernel size 7.The trailing observation window is always retained, making the resulting cache query-agnostic because the window comes entirely from context.
  • C.3 Baseline Adaptations (SnapKV and H2O).: H2O defines KV-pair importance as maximum attention over all causal queries during context prefill, rather than mean attention.This prefill-based adaptation follows KVzip’s finding that maximum attention provides better compression performance.
  • C.4 Experimental Environment: The KVPress Benchmark comparison implements RestoreKV in the official KVPress codebase and runs evaluations in its environment for fair comparison with published baselines.The comparison covers RULER and LongBench entries.

C.5 Sensitivity to Training Seeds · C.6 Evaluation Metrics

RestoreKV inference is deterministic, while training-seed variation remains small across cache ratios on RULER-4K. Evaluation follows each benchmark’s official metric, with specified protocols for KVPress-based and task-specific assessments.

  • C.5 Sensitivity to Training Seeds: Inference is deterministic because importance scoring uses no sampling and decoding is greedy.Residual stochasticity arises during training rather than inference.
  • C.5 Sensitivity to Training Seeds: All main results use the seed-0 training run.Additional runs with two different seeds assess robustness.
  • C.5 Sensitivity to Training Seeds: Two additional training seeds were evaluated alongside seed 0 on RULER-4K.The study varies random initialization and retention-ratio sampling during training.
  • C.5 Sensitivity to Training Seeds: Below 0.57 standard deviation was observed across all cache ratios.This result summarizes the three evaluated training runs on RULER-4K.
  • C.6 Evaluation Metrics: RULER and LongBench follow the KVPress evaluation protocol.This preserves consistency with standard evaluation and comparability with prior work.
  • C.6 Evaluation Metrics: SCBench uses official task-specific evaluators, while QASPER uses maximum token-level F1 over reference annotations.These metrics follow each benchmark’s established evaluation procedure.
  • C.6 Evaluation Metrics: QuALITY and LongHealth use answer-letter accuracy.Table H summarizes the benchmark metrics.
Loading 2608.01247v1…