PaperlayerView exact revision on arXiv

Source-linked AI summary

KV$^2$: A Self-Refining KV Cache

Johannes Wesch, Danni Liu, Jan Niehues

arXiv:2610.03198v1cs.AIcs.CL

TL;DR

Reusable long-context deployments must retain a growing KV cache for many later queries, while reconstruction-based scoring can require processing the full context. KV² selects informative tokens with a lightweight proxy and reconstructs selectively; its gains over baselines widen as budgets tighten, with a better quality–efficiency trade-off than full-context scoring.

  • Problem

    Reusable long-context caches must serve many later queries despite a footprint that grows with context, and reconstruction-based methods score cached entries broadly.

    Source: §1 · pp. 1–2

  • Method

    KV² uses a lightweight proxy to select informative context tokens, then uses only those tokens as reconstruction queries against the full cache to compute final eviction scores.

    Source: §3 · p. 3

  • Results

    KV²'s gains over baselines increase as cache budgets tighten, and it offers a better quality–efficiency trade-off than full-context reconstruction scoring.

    Source: §1 · p. 3

  • Takeaways & Limitations

    A small, well-chosen reconstruction set can preserve much of the useful long-context signal under severe compression without reprocessing the full context.

    Source: §6 · p. 8

    The masking-based implementation does not yet materialize physically compact reusable caches, so full decoding-time memory and latency gains require a compact serving backend.

    Source: Limitations · p. 9

Abstract

from arXiv · show

The memory footprint of the key-value (KV) cache constrains the practical use of long-context models, and it dominates cost when one prefilled context must later serve many different queries. In this reusable setting, query-agnostic compression trades cost against quality: lightweight estimators are cheap but less accurate, whereas full-context reconstruction scoring is more accurate yet reprocesses the entire prompt. We introduce KV$^2$, a query-agnostic KV-cache compression method based on selective reconstruction. KV$^2$ first uses a lightweight proxy scorer to identify informative in-context tokens, then reprocesses only this subset to compute final eviction scores. On RULER, Needle-in-a-Haystack, and LongBench, KV$^2$'s margin over baselines widens as the budget tightens: on RULER 16K at a 2% KV-cache budget it improves the average score over the next-best baseline by more than 40 percentage points, and on LongBench it attains the highest average across 2%-10% budgets at lower compression-stage runtime and peak memory than full-context reconstruction. Reusable KV-cache compression thus does not require reprocessing the full context. Our code is available at https://anonymous.4open.science/r/KVsquared-0B97.

1 Introduction

KV² targets reusable long-context caches, where memory grows with context and one prefilled cache must serve many later queries. It selectively reconstructs a small set of informative queries, with gains over baselines widening as budgets tighten and a better quality–efficiency trade-off than full-context scoring.

1 Introduction

  • Existing query-aware eviction can be less suitable when one compressed cache must support many later questions; query-agnostic compression scores entries without a specific downstream query.The work focuses on query-agnostic compression for reusable caches.Source: §1 · p. 2
  • Reusable caches must serve many later queries at budgets that can fall well below 10%, while the KV footprint grows with context length.The introduction describes this as a deployment setting rarely studied in prior work.Source: §1 · pp. 1–2
  • KV² uses a lightweight proxy to select informative context tokens, then reconstructs only with those selected queries against the full cache to assign eviction scores.The method evicts low-scoring entries; the selective scoring design concentrates computation on a subset of queries.Source: §1 · p. 2, Figure 1 · p. 2
  • The two-stage design is modular in its proxy choice and supports iterative self-refinement when more compression-time compute is available.The introduction also frames informative anchors as sparse at extreme budgets, motivating selective reconstruction.Source: §1 · p. 2
  • KV²'s gains over baselines increase as the cache budget tightens, and it offers a better quality–efficiency trade-off than full-context reconstruction scoring.The reported comparison includes compression-stage runtime and peak memory.Source: §1 · p. 3

2 Related work

Prior KV-cache work spans inference-system optimizations and cache compression, including query-aware eviction and query-agnostic methods for reusable caches. Reconstruction-based approaches align with reusable-cache objectives but can be costly; KV² retains reconstruction while selecting fewer queries.

2 Related work

  • Early training-free eviction methods include StreamingLLM's recent and attention-sink tokens [Xiao et al., 2024] and H2O's heavy hitters, highlighting unequal token utility and the value of inference-time signals.The related-work discussion also notes that recency matters but is not sufficient on its own.Source: §2 · p. 3
  • Query-aware methods such as SnapKV [Li et al., 2024], LAQ [Wang et al., 2025], and Judge Q [Liu et al., 2026] can fit an anticipated continuation but are less suited to caches serving multiple later questions.Expected Attention [Devoto et al., 2025a] and KeyDiff instead use query-agnostic signals.Source: §2 · p. 3
  • KVzip [Kim et al., 2025] uses context reconstruction for reusable caches but incurs broad reconstruction cost; KV² keeps an actual reconstruction pass while shrinking its query set.Fast KVzip and KVzap seek cheaper reconstruction-derived signals [Kim et al., 2026; Jégou and Jeblick, 2026].Source: §2 · p. 3

3 Method

KV² compresses a reusable cache by selecting informative reconstruction queries with a lightweight proxy, then reprocessing only those queries against the full cache to score entries. Its chunk-local scoring and optional iterative refinement make selective reconstruction tractable while preserving a single global eviction decision.

3 Method

  • KV² uses a lightweight proxy to select informative in-context queries, then reprocesses only those against the full cache to assign final eviction scores.The core design choice is which queries to use for reconstruction, rather than whether to reconstruct.Source: §3 · p. 3
  • KeyDiff is the default proxy; KV² averages its layer- and head-specific scores, then selects the highest-scoring positions as ordered reconstruction queries.The selected positions may be a sparse set rather than a contiguous span; ρ is the query-selection ratio.Source: §3 · pp. 3–4
  • Each chunk proposes its own reconstruction queries, and its local scores are merged into one global eviction decision.Chunking keeps reconstruction-based scoring tractable on long prompts while allowing different context regions to propose queries.Source: §3 · p. 3, Figure 2 · p. 4
  • Selected queries are appended after the prefilled cache and processed with access to the full cache, but their attention scores are restricted to positions in the current chunk.This global-conditioning, local-scoring split is KV²’s main design choice.Source: §3 · p. 4
  • A cached position’s chunk score is its maximum attention across reconstruction queries; chunk scores are then gathered, sink positions pinned, and the lowest-scoring entries evicted to meet budget.The maximum retains a position when any reconstruction query depends strongly on it.Source: §3 · p. 4
  • Stage 2 scores can replace the Stage 1 proxy for another pass, but single-pass KV² is the default because it offers the strongest quality-efficiency trade-off.The optional multi-pass variant is studied as an extension.Source: §3 · p. 4
  • The proxy pass is linear in prompt length, while reconstruction-stage runtime and peak activation footprint scale with the selected fraction ρ rather than full-chunk reconstruction.Score extraction is a smaller chunk-local postprocessing step.Source: §3 · pp. 4–5

4 Experiments

Experiments evaluate KV² on retrieval stress tests and long-context benchmarks using two main model backbones. Across the reported tasks, its advantage is clearest under tight cache budgets, with a boundary when useful evidence is spread broadly across the prompt.

4.1 Experimental setup

  • The main evaluation uses RULER 4K/16K, LongBench, and a Needle-in-a-Haystack diagnostic, mainly with Llama-3.1-8B-Instruct and Qwen3-8B; default settings use 4K-token chunks and KeyDiff.The comparisons include proxy scoring, expected future-query scoring, and reconstruction-based compression baselines.Source: §4.1 · p. 5

4.2 RULER 16K

  • On RULER 16K, KV² degrades more gracefully than baselines at 5% and especially 2% KV-cache budgets, with the advantage strongest on sparse retrieval and aggregation.Distractor-heavy variants remain difficult for all methods; Figure 4 highlights the severe-budget regime for Qwen3-8B.Source: §4.2 · p. 5, Figure 4 · p. 5

4.3 Needle-in-a-Haystack (NIAH)

  • On Needle-in-a-Haystack at a fixed 5% cache budget, KV² is the most stable method as context length and needle depth increase; KeyDiff is patchy and KVzip degrades faster on longer haystacks.The diagnostic uses Llama-3.1-8B-Instruct and per-example ROUGE-2 recall.Source: §4.3 · pp. 5–6, Figure 5 · p. 5

4.4 LongBench

  • At a 2% KV-cache budget, LongBench average scores rise from 24.16 to 33.59 on Llama-3.1-8B-Instruct and from 16.59 to 32.97 on Qwen3-8B relative to the next-best baseline.Table 1 reports LongBench results across budgets for both backbones.Source: §4.4 · p. 6, Table 1 · p. 6
  • LongBench gains concentrate in multi-document question answering and synthetic retrieval and counting, while RepoBench-P is a recurring exception where broader reconstruction is more helpful.Selective reconstruction helps most when evidence is sparse and anchor-like, and less when useful evidence is spread broadly across the prompt.Source: §4.4 · p. 6

5 Analysis

The analysis examines compression-stage efficiency and how proxy choice, reconstruction budget, and iterative refinement affect KV2’s operating regime.

5.1 Single-pass efficiency

  • KV2 is consistently faster and uses less peak memory than KVzip across relevant chunk sizes, with 4K capturing nearly all available quality efficiently.These measurements concern compression-stage costs in a masking-based backend, not realized decode-time savings from physically compact caches.Source: Figure 6 · p. 7, §5.1 · p. 7, §5 · p. 7

5.2 Proxy scorer

  • KeyDiff is the strongest proxy throughout, while CUR and StreamingLLM fall below a random-proxy baseline at the highest compression ratios studied.Content-adaptive proxy selection can cover useful signal better under severe compression; the second-stage scorer is held constant in this comparison.Source: Figure 7 · p. 7, §5.2 · p. 7

5.3 Reconstruction budget

  • Optimal reconstruction fractions rise with cache budget: Llama peaks at 10%, 50%, and 100% reconstruction for 2%, 5%, and 10% final budgets.For Qwen3-8B, the corresponding maxima are 10%, 25%, and 75%; no single reconstruction fraction is optimal across budgets.Source: Figure 8 · p. 8, §5.3 · p. 8

5.4 Iterative refinement

  • Iterative-refinement gains concentrate in the first few passes, especially with a weak proxy or 2% reconstruction, and diminish thereafter.The gains mainly cost runtime, while peak memory changes much less.Source: Figure 9 · p. 8, §5.4 · p. 8

6 Conclusion

KV2 shows that reusable KV-cache compression can preserve useful long-context signal without reconstructing everything.

6 Conclusion

  • A small, well-chosen reconstruction set can preserve much of the useful long-context signal under severe compression.The two-stage design and sparse informative anchors at extreme budgets point toward compression that is both cheaper and more accurate.Source: §6 · p. 8

Limitations

The paper identifies limitations in refinement scheduling, proxy quality, and the current serving implementation.

Limitations

  • Iterative refinement is fixed rather than adaptive, and its gains mainly occur with weak initial proxies or very small reconstruction budgets, at a runtime cost.The authors suggest adaptive schedules that continue while the reconstruction set or scores are still changing meaningfully.Source: Limitations · p. 9
  • The first-stage proxy remains an improvement target; content-adaptive scorers outperform structurally fixed ones at low budgets.The authors suggest that learned or model-aware proxies may improve performance further.Source: Limitations · p. 9
  • The masking-based realization does not yet materialize physically compact reusable caches, so full decoding-time memory and latency gains require a compact serving backend.The authors identify headwise cache compaction, optimized attention kernels, and repeated-query serving pipelines as needed components.Source: Limitations · p. 9

A Complexity Analysis

The complexity analysis separates KV2’s proxy scoring, reconstruction forward pass, and final score extraction; refinement adds runtime while peak memory remains tied mainly to chunk size.

A Complexity Analysis

  • KeyDiff proxy scoring is O(n_c); selective reconstruction forward costs O(ρn_c^2 + ρ^2n_cm/2) across chunks.Each reconstruction block has length n_q ≈ ρm, and the forward pass remains an ordinary causal pass that can use optimized attention implementations.Source: A Complexity Analysis · pp. 10–11
  • Score extraction costs O(ρm^2) per chunk and O(ρmn_c) overall, but its softmax-over-keys then max-over-queries dependency resists standard blockwise attention.For each chunk, extraction recomputes attention on a reduced n_q × (m + n_q) matrix.Source: A Complexity Analysis · p. 11

B Iterative refinement cost

  • Refinement runtime rises roughly with pass count, while peak memory remains driven mainly by chunk size because passes reuse essentially the same working buffer.This makes iterative refinement an optional extension for harder settings rather than the default configuration.Source: B Iterative refinement cost · p. 11

C LongBench-E

The appendix extends evaluation across LongBench-E and RULER context lengths and backbones. KV2 variants remain competitive across these settings, with smaller advantages on shorter RULER contexts.

C LongBench-E

  • Across LongBench-E length buckets, all three KV2 variants remain ahead of KeyDiff and KVzip on overall average and stay relatively close to one another.The results, together with Figure 6, support the paper’s 4K default.Source: C LongBench-E · p. 11

D Additional benchmark results

  • Appendix D provides full per-benchmark results, including RULER 4K and 16K tables and the 200K-token InfiniteBench evaluation.The InfiniteBench results are covered in the D.3 summary.Source: D Additional benchmark results · p. 11

D.1 RULER 4K

  • At RULER 4K, KV2’s advantage is smaller: KVzip remains competitive at 10% and 5%, while KV2 degrades more gracefully under more aggressive compression.The passage identifies shorter prompts, where global reconstruction is less wasteful, as the boundary of the main result.Source: §D.1 · p. 11

D.2 RULER 16K

  • The tables below report the full per-backbone RULER 16K results summarized in the main text, and include results covering a wider range of KV budgets.Source: §D.2 · p. 11

D.2.1 Llama-3.1-8B-Instruct

  • On Llama-3.1-8B-Instruct, KV2 preserves aggregation and retrieval accuracy better than the baselines as the RULER 16K budget tightens toward 2%.The passage describes the pattern as consistent with the main text.Source: §D.2.1 · p. 11

D.2 RULER 16K

  • On Qwen3-8B, KV2’s margin widens at 2%–5% budgets and is largest on sparse retrieval and aggregation subtasks.The same pattern holds on Llama-3.1-8B-Instruct as budgets tighten toward 2%.Source: §D.2.2 · p. 11, §D.2.1 · p. 11

D.3 InfiniteBench

On InfiniteBench, KV2 improves over the strongest baseline across the tested budgets on Llama-3.1-8B-Instruct and stays slightly ahead on Qwen2.5-7B-Instruct-1M. Related results show task sensitivity and chunk-size trade-offs.

D.3 InfiniteBench

  • On InfiniteBench, KV2 improves over the strongest baseline at 2%, 5%, and 10% budgets on Llama-3.1-8B-Instruct and stays slightly ahead on Qwen2.5-7B-Instruct-1M.Retrieval tasks remain comparatively robust under compression, while QA is more sensitive to reduced budgets.Source: §D.3 · pp. 11–12
  • For iterative KV2 at a fixed 2% reconstruction percentage, runtime rises roughly with refinement passes, while peak memory depends mainly on chunk size and changes less with iterations.Figure 10 plots runtime and peak memory against chunk size.Source: Figure 10 · p. 12
  • On LongBench-E, KV2 consistently outperforms KeyDiff and KVzip on average; larger chunks tend to be strongest on longer inputs, while smaller chunks remain more competitive on shorter inputs.The input-length buckets are measured in words, whereas the KV2 chunk sizes are measured in tokens.Source: Table 2 · p. 12

E Token-level reconstruction example

The token-level example shows KeyDiff selecting reconstruction positions concentrated on the prompt’s instruction, assignment chain, and question-answer span, with sparse selections in repetitive filler.

E Token-level reconstruction example

  • Selected positions cluster in the instruction, variable-assignment chain, and question-answer span, while repetitive filler contributes only sparse selections.This pattern illustrates the signal sparsity that selective reconstruction is designed to exploit.Source: E Token-level reconstruction example · p. 12, Figure 12 · p. 16
  • The RULER 4K figure reports scores by KV retention ratio for Llama-3.1-8B-Instruct and Qwen3-8B, with KV2 more robust than lightweight heuristics at small budgets.KVzip becomes more competitive as more KV entries are retained.Source: Figure 11 · p. 13
  • Figure 12 highlights the first-stage KeyDiff selections in a condensed RULER variable-tracking prompt, with positions chosen after averaging over layers and KV heads.Selection is shown at token granularity, so some highlights cover only parts of words.Source: E Token-level reconstruction example · p. 12, Figure 12 · p. 16
Loading 2610.03198v1…