Summary

• KV-SURGEON introduces a compensation-aware offline KV-cache compaction method that jointly optimizes which cache entries to delete and how the remaining entries should be adjusted, rather than treating selection and refitting as separate sequential steps.
• The method derives a quadratic surrogate of the attention-output error that admits closed-form expressions for greedy deletion saliencies and compensating updates, inspired by Optimal Brain Surgeon (OBS) pruning techniques.
• KV-SURGEON substantially advances the cost–accuracy Pareto frontier across QuALITY, LongHealth, and Long Code Arena benchmarks, with models from 4B to 30B parameters and compaction ratios from 5× to 100×.
• At 50× compaction, KV-SURGEON improves mean question-answering accuracy by 3.0 percentage points over the prior state of the art (attention matching), while its fast variant (KV-SURGEON-fast) exceeds all prior methods at only two-fifths of the cost.
• A controlled ablation demonstrates that the performance gain comes specifically from compensation during selection (updating surviving entries after each deletion), not merely from the surrogate objective.

Introduction and Theoretical Foundation

Background and Motivation

Large language model (LLM) applications frequently query a fixed context (documents, patient records, codebases). The key-value cache (KV-cache) grows linearly with context length and dominates memory usage during long-context inference. Offline compaction processes a context once, compresses its cache, and reuses the compacted cache for all subsequent requests—amortizing a one-time cost over many queries.

The Core Problem: Softmax Coupling

The central challenge is deciding which cache entries to delete. For one attention head with $T$ cache entries, each containing a key $k_j \in \mathbb{R}^{d_k}$ and value $v_j \in \mathbb{R}^{d_v}$, a query $q \in \mathbb{R}^{d_k}$ produces:

$$P \triangleq \text{softmax}\left(\frac{q^\top k_\ell}{\sqrt{d_k}}\right)_{\ell=1}^T \in \mathbb{R}^T, \quad o \triangleq \sum_{j=1}^T P_j v_j \in \mathbb{R}^{d_v}$$

The key insight is that softmax renormalization couples deletion decisions nonlinearly: removing one entry changes the attention weights of all remaining entries. This makes selection a combinatorial problem rather than a simple ranking problem.

Limitations of Existing Approaches

Existing methods suffer from a critical restriction: they evaluate each candidate deletion with the rest of the cache held fixed. Specifically:

• Attention matching (Zweiger et al., 2026) refits the retained cache after selection, but selection itself ignores compensation potential.
• Output-aware methods (CAOTE, OBCache, ReST-KV) estimate single-deletion error but evaluate deletions in isolation.
• Eviction-by-proxy methods (H2O, StreamingLLM) rank entries by scores that ignore the rest of the cache entirely.

What is missing is compensation-aware selection: each candidate deletion should be scored by the error remaining after surviving entries have been adjusted to compensate for it.

Methodology

Problem Formulation

For one attention head with $T$ cache entries and $N$ calibration queries $q_1, \ldots, q_N$, the attention weights and outputs are:

$$P_{ij} \triangleq \frac{e^{Z_{ij}}}{\sum_{\ell=1}^T e^{Z_{i\ell}}}, \quad o_i \triangleq \sum_{j=1}^T P_{ij} v_j$$

where $Z_{ij} \triangleq q_i^\top k_j / \sqrt{d_k}$.

Given a budget $t < T$, the goal is to retain a support set $S$ with $|S| = t$, assigning each retained entry $j \in S$ an adapted value $\hat{v}_j$ and a positive attention multiplier $w_j > 0$. The compacted cache produces:

$$\hat{P}_{ij} \triangleq \frac{w_j P_{ij}}{\mu_i}, \quad \hat{o}_i \triangleq \frac{1}{\mu_i} \sum_{j \in S} P_{ij} w_j \hat{v}_j$$

where $\mu_i \triangleq \sum_{j \in S} P_{ij} w_j$ is the ratio between compacted and original normalization terms.

The Quadratic Surrogate

The key technical contribution is deriving a tractable surrogate. By introducing the variable $U \triangleq \text{diag}(w)\hat{V}$, the output can be written as a ratio of two linear terms:

$$\hat{o}_i = \frac{n_i}{\mu_i}, \quad n_i \triangleq \sum_{j \in S} P_{ij} u_j, \quad \mu_i = \sum_{j \in S} P_{ij} w_j$$

Instead of matching the ratio directly, the method asks the compacted cache to reproduce numerator and denominator separately. The resulting quadratic surrogate is:

$$\|O - P_S U\|_F^2 + \gamma^2 \|1 - P_S w\|_2^2 = \|[O \; \gamma 1] - P_S [U \; \gamma w]\|_F^2$$

The right-hand side forms a single linear least-squares problem with $d_v + 1$ target columns. With $Y \triangleq [O \; \gamma 1] \in \mathbb{R}^{N \times (d_v+1)}$ and a ridge penalty $\lambda > 0$, the support score is:

$$F(S) \triangleq \min_B \|Y - P_S B\|_F^2 + \lambda \|B\|_F^2$$

Greedy Elimination with Closed-Form Compensation

For a fixed support $S$, the optimal coefficients and shared inverse are:

$$B = G P_S^\top Y \in \mathbb{R}^{|S| \times (d_v+1)}, \quad G \triangleq (P_S^\top P_S + \lambda I)^{-1} \in \mathbb{R}^{|S| \times |S|}$$

The saliency of entry $j$ (the increase in optimal surrogate objective when removed) is:

$$s_j \triangleq F(S \setminus \{j\}) - F(S) = \frac{\|B_{j,:}\|_2^2}{G_{jj}}$$

After removing entry $j$, the updates are:

$$B \leftarrow B_{-j,:} - \frac{G_{-j,j} B_{j,:}}{G_{jj}}, \quad G \leftarrow G_{-j,-j} - \frac{G_{-j,j} G_{j,-j}}{G_{jj}}$$

These give the exact minimizer of the surrogate on $S \setminus \{j\}$ without solving another regression. The algorithm repeatedly removes the entry with smallest saliency, updating $B$ and $G$ until $|S| = t$.

Implementation Details

Key practical additions include:

• Candidate pool: Pre-filter to $m = \min\{T, \lceil ct\rceil\}$ entries using a leave-one-out score ($c = 4$)
• Block deletion: Delete $n_b = 32$ entries per step using block-update formulas
• Final refit: Optimize multipliers $w$ (box-constrained) then values $\hat{V}$ (least-squares)
• Mild-compaction fallback: For heads retaining $> 30\%$ of entries, skip greedy elimination

Empirical Validation / Results

Main Results at 50× Compaction

Table 1: Accuracy (%) at 50× compaction across model–benchmark pairs

Key findings:

• KV-SURGEON reaches 65.1% mean accuracy in 145 seconds, 3.0 points above AM-HAK at comparable cost
• KV-SURGEON-fast reaches 63.9% in 54 seconds, exceeding AM-OMP (63.2% in 763 s) with a 14.1× reduction in pipeline time
• Gains are consistent across all 8 model–benchmark pairs and all compaction ratios (5× to 100×)

Source of Performance Gain

Table 2: Impact of compensation during selection (Qwen3-4B, QuALITY, 50×)

Disabling only the compensating update (keeping the surrogate, pool, and refit fixed) reduces accuracy by 3.0 points, demonstrating that compensation during selection—not the surrogate itself—drives the improvement.

Robustness to Calibration

Replacing expensive self-study calibration with a single repeat pass costs KV-SURGEON only 1.2 points of mean accuracy (from 65.1% to 63.9%), while the same change costs AM-HAK 4.9 points. This robustness persists across all compaction ratios.

Theoretical and Practical Implications

Theoretical Contributions

1. Quadratic surrogate derivation: The surrogate in Equation (8) is shown to be the second-order Taylor expansion of the exact loss around the original cache (Appendix A.1), providing a principled approximation.

2. Closed-form compensation: The saliency formula $s_j = \|B_{j,:}\|_2^2 / G_{jj}$ has a natural interpretation: it equals the squared coefficient magnitude of entry $j$ times the ridge-regularized residual of regressing its attention column on the other retained columns. An entry with a large coefficient is cheap to delete when its attention column lies close to the span of others.

3. Block deletion exactness: The block updates in Equation (14) give the exact minimizer on $S \setminus J$ for any selected block $J$; only the block selection (by individual saliencies) is approximate.

Practical Implications

• Cost–accuracy frontier: KV-SURGEON establishes a new Pareto frontier for offline KV-cache compaction, making aggressive compaction (50×–100×) practical for production long-context applications.
• Calibration flexibility: The method's insensitivity to calibration quality means practitioners can use cheap repeat-pass calibration without significant accuracy loss.
• Model generality: Gains hold across dense and mixture-of-experts architectures from 4B to 30B parameters, suggesting broad applicability.
• Compute efficiency: The O($m^3$) complexity (bounded by the pool size $ct$, not cache size $T$) makes the method scalable to very long contexts.

Conclusion

KV-SURGEON demonstrates that deletion and adaptation can be coupled efficiently in offline KV-cache compaction, contrary to prior methods that treated them sequentially. The key insight is that a quadratic surrogate of the attention-output error enables closed-form compensation-aware selection, where each candidate deletion is scored after surviving entries have been adjusted to compensate.

The method improves the cost–accuracy frontier at every evaluated compaction ratio (5×–100×) across three benchmarks and five models, with the largest gains at aggressive compaction levels (50× and 100×). Its fast variant exceeds all prior methods at two-fifths of the cost of the strongest baseline.

Future directions suggested by this work include:

• Exploring whether compensation-aware selection extends to other attention variants (e.g., sliding-window, linear attention)
• Investigating joint optimization across heads rather than independent per-head compaction
• Applying the OBS-inspired framework to other cache-compression axes (quantization, low-rank projection)