Summary (Overview)

• RestoreKV introduces a novel complementary restoration perspective to query-agnostic KV cache eviction: instead of only learning which original KV pairs to retain, it learns a shared mechanism that generates a compact, context-conditioned restore cache to complement the retained context cache under the same total KV budget.
• The method uses 8 learnable restore tokens that attend to the full KV cache in a single LoRA-adapted pass after context prefill, generating context-conditioned KV pairs that are combined with the base evictor's retained states.
• RestoreKV is trained via self-distillation from a frozen full-cache teacher, optimizing only 0.4% of parameters (restore-token embeddings + LoRA adapters) with no task-specific tuning.
• Across 4 backbones, 4 benchmarks, and 5 base eviction methods, RestoreKV substantially reduces compression-induced degradation. On Qwen3-4B at a 5% budget, it improves KVzip from 38.2 to 73.2 on RULER-4K; applied to KVzip+, it reaches 86.4 RULER accuracy at 16× compression on the KVPress Benchmark.
• The method adds less than 0.5% one-time cache-construction overhead (0.03–0.04 s and 84 MB) in a 32K-context evaluation, with no query-time KV-memory or decoding cost.

---

Introduction and Theoretical Foundation

Background

Large language models increasingly rely on long contexts, but their key–value (KV) caches grow linearly with sequence length. KV cache eviction reduces this cost by retaining only a subset of cached KV pairs. Query-agnostic eviction methods compress a context once before future queries are known and reuse the resulting cache across arbitrary subsequent requests.

The Selection-Based Formulation

Recent query-agnostic methods (KVzip, KVzip+, ContrastKV) largely follow the same formulation: select a better subset of the original KV pairs. The resulting cache remains composed only of retained original states. The central question has been which KV pairs should survive under a fixed budget.

Key Insight

Although the information lost through eviction is context-specific, the mechanism for generating its compact complement can be shared across contexts.

This motivates a complementary question: rather than relying only on a better subset, can we generate a small context-conditioned complement before eviction to preserve behavior that would otherwise be lost?

Theoretical Foundation

Consider a decoder-only Transformer $f_\theta$ with L layers and H KV heads. Given a context $x_{1:T}$, the prefill stage produces a full KV cache $\mathcal{C}$ containing $TLH$ KV pairs. At KV budget ratio $r$, the total KV budget is:

$$B = \lfloor r T L H \rfloor.\tag{1}$$

The retained context cache is:

$$\mathcal{C}' = \operatorname{Evict}(\mathcal{C}, \mathbf{s}, B),\tag{2}$$

where $|\mathcal{C}'| = B$ and Evict follows the layer- and head-wise allocation rule of the underlying method.

---

Methodology

3.1 Restore-Cache Generation with LoRA

RestoreKV introduces $n$ learnable restore-token embeddings:

$$E = [e_1, \ldots, e_n] \in \mathbb{R}^{n \times d},\tag{3}$$

where $d$ is the hidden dimension. After context prefill, the restore tokens are processed at positions $T+1, \ldots, T+n$ with causal access to the full KV cache $\mathcal{C}$. A single LoRA-adapted restore pass generates the context-conditioned restore cache:

$$\mathcal{C}_{\mathrm{res}} = \operatorname{Restore}_{\theta, \phi}(E \mid \mathcal{C}),\tag{4}$$

where $\theta$ denotes the frozen backbone parameters and $\phi$ denotes the LoRA parameters.

For each adapted linear projection W, LoRA applies:

$$\Delta W = \frac{\alpha}{r_{\mathrm{LoRA}}} B_\phi A_\phi, \qquad W_{\mathrm{res}} = W + \Delta W,\tag{5}$$

where $r_{\mathrm{LoRA}}$ is the LoRA rank, $\alpha$ is the scaling factor, and $A_\phi$, $B_\phi$ are trainable low-rank matrices.

3.2 Budget Matching

To preserve the query-time KV budget, RestoreKV reserves $nLH$ of the $B$ cache slots for $\mathcal{C}_{\mathrm{res}}$ and lets the base evictor fill the remaining $B - nLH$ slots:

$$\widetilde{\mathcal{C}} = \operatorname{Concat}\left(\operatorname{Evict}(\mathcal{C}, \mathbf{s}, B - nLH), \mathcal{C}_{\mathrm{res}}\right),$$ $$|\widetilde{\mathcal{C}}| = (B - nLH) + nLH = B = |\mathcal{C}'|.\tag{6}$$

Retained context KV pairs preserve their original RoPE phases, while restore tokens occupy positions $T+1, \ldots, T+n$ and future queries begin at position $T+n+1$.

3.3 Self-Distillation Training Objective

Given a training context $x_{1:T}$ and query $q = (q_1, \ldots, q_Q)$, the full-cache teacher first generates an answer $y = (y_1, \ldots, y_M)$. The teacher and student then evaluate the same answer:

$$p_i^{\text{full}}(\cdot) = p_\theta(\cdot \mid q, y_{<i}, \mathcal{C}),$$ $$p_i^{\text{res}}(\cdot) = p_\theta(\cdot \mid q, y_{<i}, \widetilde{\mathcal{C}}).\tag{7}$$

The loss is the token-averaged symmetric KL divergence:

$$\mathcal{L}_{\text{distill}} = \frac{1}{2M} \sum_{i=1}^{M} \left[ \mathrm{KL}(p_i^{\text{full}} \| p_i^{\text{res}}) + \mathrm{KL}(p_i^{\text{res}} \| p_i^{\text{full}}) \right].\tag{8}$$

KV budget ratios are uniformly sampled:

$$r \sim \mathcal{U}(r_{\min}, r_{\max}), \qquad 0 < r_{\min} < r_{\max} < 1.\tag{9}$$

Training Configuration

• 8 restore tokens ($n=8$), LoRA rank $r_{\mathrm{LoRA}}=8$, scaling factor $\alpha=16$ (0.4% of Qwen3-4B parameters)
• Training set: 6.2k context–query pairs from 2.5k unique contexts (LongAlpaca, PG-19, Tulu-3 Flan)
• Budget ratios sampled from $\mathcal{U}(0.025, 0.25)$
• Full training run for Qwen3-4B: ~2 hours on a single NVIDIA RTX PRO 6000 GPU

---

Empirical Validation / Results

Main Results (Qwen3-4B, RULER-4K)

Generalization Across Five Eviction Methods

RestoreKV improves SnapKV and H₂O (query-agnostic adaptations) as well:

Comparison with Synthesized-Cache Baseline

At r=0.05, RestoreKV+ achieves 70.7 vs 52.8 on RULER-4K compared to AM-fast, and constructs a 4K-context cache in 0.74 s vs 9.68 s (~13× faster).

Component Ablations (RULER-4K, Qwen3-4B, r=0.05)

Key finding: Attention-side adaptation (LoRA on q/k/v) is the primary source of recovery, contributing 72.4/73.2 points versus only 42.1 from learned embeddings alone.

Effects of Full-Context Conditioning

Full-context conditioning adds +8.8 points at r=0.05 over evicted-cache generation.

Number of Restore Tokens

Efficiency Analysis (Llama-3.1-8B-Instruct, 32K context)

• Compression overhead: +0.03–0.04 s (<0.5% of total compression time), +84 MB peak memory (0.4%)
• Inference: Attention latency nearly identical to KVzip (0.457–0.253 ms); KV-cache memory exactly matched (4.20 → 0.20 GB for both methods)

Context-Attention Recovery

RestoreKV recovers lost context attention:

• Average attention mass: 0.56% → 1.30% (full cache: 2.73%)
• KL divergence from full-cache model: 0.198 → 0.167
• Predictive distribution divergence: 7.3 → 3.8

---

Theoretical and Practical Implications

Theoretical Implications

1. Beyond selection-based caching: RestoreKV demonstrates that the selection-based formulation of KV eviction is not the only viable endpoint. A hybrid approach—combining selection with a small, learned, context-conditioned complement—can substantially outperform pure selection under aggressive budgets.

2. Shared restoration mechanism: The finding that a single, shared mechanism (LoRA adapters + restore tokens) can generate context-specific complements across diverse contexts validates the premise that "the mechanism for generating compact complements can be shared across contexts."

3. Attention-side adaptation is key: The ablations reveal that the recovery primarily stems from adapting attention projections (q/k/v), not from memorized information in embeddings—suggesting that the model learns a generalizable transformation of full-cache information into compact form.

Practical Implications

1. Drop-in compatibility: RestoreKV works as a plug-in for any query-agnostic eviction method without modifying the importance scorer or eviction rule, making it easy to integrate into existing pipelines.

2. Negligible overhead: With <0.5% one-time compression overhead and zero query-time cost, RestoreKV is practical for real-world deployment.

3. Budget-matched: The method preserves the exact query-time KV memory budget and decoding cost of the base method, which is critical for deployment with hard memory constraints.

4. Broad applicability: Improvements across model scales (0.6B–8B), architectures (Qwen3, Llama-3.1), and eviction strategies suggest wide applicability.

---

Conclusion

RestoreKV revisits query-agnostic KV cache eviction from a complementary restoration perspective. It uses a shared, LoRA-adapted restoration mechanism to generate a compact, context-conditioned restore cache before eviction and combines it with retained original KV pairs under the same total budget.

Key takeaways:

• Preserves the base importance scorer and eviction rule while disabling adapters after cache construction
• Adds negligible one-time overhead and no query-time KV-memory or decoding cost
• Consistently reduces compression-induced degradation across four backbones, four benchmarks, and five base eviction methods
• Largest gains appear under aggressive budgets (e.g., +35.0 points on RULER-4K at r=0.05 for KVzip)
• Attention-side adaptation is the primary source of recovery; full-context conditioning provides additional complementary benefit

Future directions suggested by the work include: exploring more sophisticated restore-cache architectures, extending the approach to query-aware settings, and investigating whether larger restore-cache budgets could further close the gap to full-cache performance.