KVSculpt: KV Cache Compression via Distillation

• KVSculpt is a KV-cache compression method that reformulates sequence-length reduction as a distillation problem by generating unconstrained key-value pairs.
• It employs an alternating optimization procedure, using L-BFGS for key updates and closed-form least-squares for value fitting, preserving layer attention.
• Adaptive budget allocation based on pilot MSE analysis improves compression efficiency, yielding up to a 4× KL reduction compared to baseline methods.

KVSculpt is a KV-cache compression method for efficient long-context LLM inference that formulates sequence-length reduction as a distillation problem rather than as eviction or merging of original cache entries. In the taxonomy given for KV-cache compression, approaches that reduce the per-pair footprint, such as quantization and low-rank decomposition, are orthogonal to those that reduce the sequence length of the cache. Along the sequence-length dimension, existing methods range from pure eviction to merging, but both remain anchored to the original cache entries. KVSculpt instead replaces older cache content with a smaller set of unconstrained key-value pairs optimized in continuous embedding space so as to preserve each layer’s attention behavior (Jiang et al., 29 Mar 2026).

1. Position within KV-cache compression

KVSculpt addresses the setting in which a transformer has already processed a long context and must retain a cache for subsequent decoding. The central problem is that KV cache compression is critical for efficient long-context LLM inference, and that compression can be pursued along at least two axes: reducing the representation size of each stored pair, or reducing the number of stored pairs themselves (Jiang et al., 29 Mar 2026).

Within sequence-length compression, KVSculpt is positioned at the opposite end of the spectrum from methods that simply decide which original KV pairs to keep. Pure eviction selects a subset of existing pairs; merging combines similar existing pairs into fewer ones. KVSculpt departs from both by allowing the compressed cache to consist of newly optimized pairs rather than a subset or recombination of the original entries. This makes the compressed representation unconstrained by the discrete support of the original token positions.

A plausible implication is that KVSculpt should be interpreted less as cache pruning and more as layerwise attention-preserving distillation. That interpretation is explicit in the title formulation, “KV Cache Compression as Distillation,” and is reinforced by the use of unconstrained optimization over compressed keys and least-squares fitting of compressed values (Jiang et al., 29 Mar 2026).

2. Formal problem statement

The formulation is defined for a single KV head in a transformer layer after processing a context of length $T$. The full cache is

$$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$

with $d_k=d_v=d$ for grouped-query attention (Jiang et al., 29 Mar 2026).

The cache is partitioned into two zones. A retain zone contains the most recent $m$ pairs $(K_{\mathrm{ret}},V_{\mathrm{ret}})$, which are kept unchanged. A compress zone contains the older $T-m$ pairs $(K_{\mathrm{old}},V_{\mathrm{old}})$, which are replaced by only $k$ “free” pairs $(K_c,V_c)$. The overall compression ratio is

$$r = \frac{k+m}{T} \quad\text{with}\quad k = \lfloor rT \rfloor.$$

The compressed variables are

$$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$0

For any future query batch $$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$1, the compressed cache is formed by concatenation,

$$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$2

with the objective that the attention output under $$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$3 closely match that of the full cache $$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$4 (Jiang et al., 29 Mar 2026).

The loss combines an output-MSE term and an LSE-matching term on the pre-softmax log-sum-exp scores. Defining

$$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$5

and

$$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$6

the per-head loss is

$$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$7

The formulation also notes that one may measure KL divergence on the final token logits, and that in practice MSE on $$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$8 correlates well with KL. This establishes an explicit bridge between layer-local reconstruction and end-task decoding fidelity (Jiang et al., 29 Mar 2026).

3. Alternating optimization procedure

KVSculpt exploits an asymmetry in the loss landscape: $$K \in \mathbb{R}^{T\times d_k}, \qquad V \in \mathbb{R}^{T\times d_v},$$9 is smooth in $d_k=d_v=d$0 through the softmax, but quadratic in $d_k=d_v=d$1. The resulting solver alternates between a key-optimization step and a value-fitting step (Jiang et al., 29 Mar 2026).

The $d_k=d_v=d$2-step updates keys with L-BFGS:

$d_k=d_v=d$3

using L-BFGS on $d_k=d_v=d$4. The gradients are computed by backpropagation through the softmax. This preserves the continuous optimization character of the method: keys are not restricted to original positions, and the search operates directly in embedding space.

The $d_k=d_v=d$5-step freezes $d_k=d_v=d$6, computes the attention-weight matrix

$d_k=d_v=d$7

partitions it as $d_k=d_v=d$8, and solves

$d_k=d_v=d$9

with $m$0 as a small ridge penalty (Jiang et al., 29 Mar 2026).

The per-head pseudocode is specified as follows. First, initialize $m$1 by selecting the top-$m$2 original keys by accumulated attention score. Second, solve $m$3 via the ridge-regression closed form. Third, repeat for $m$4 outer steps: if $m$5, update $m$6 in closed form; otherwise, compute $m$7 and take one L-BFGS iteration on $m$8. Fourth, return the distilled $m$9 (Jiang et al., 29 Mar 2026).

This alternating structure is significant because it separates the nonlinear search over attention geometry from the linear fitting of values. The reported analysis further states that continuous key optimization is crucial: a joint-optimization baseline that fits only $(K_{\mathrm{ret}},V_{\mathrm{ret}})$0 at original key positions yields almost no improvement over Select+Fit, whereas freeing $(K_{\mathrm{ret}},V_{\mathrm{ret}})$1 gives the reported $(K_{\mathrm{ret}},V_{\mathrm{ret}})$2 gain. This suggests that the main source of improvement is not merely better value regression, but the ability to relocate support points in key space (Jiang et al., 29 Mar 2026).

4. Adaptive allocation of compression budget

KVSculpt augments the core optimizer with adaptive budget allocation. The motivation is that uniformly allocating the same $(K_{\mathrm{ret}},V_{\mathrm{ret}})$3 to every layer and head is suboptimal because some components are much harder to compress than others (Jiang et al., 29 Mar 2026).

The mechanism is a two-level pilot procedure. In the first stage, a pilot run under uniform $(K_{\mathrm{ret}},V_{\mathrm{ret}})$4 uses short L-BFGS optimization—given as, for example, 60 steps per layer or 30 steps per head—to measure per-component reconstruction difficulty. The metrics are $(K_{\mathrm{ret}},V_{\mathrm{ret}})$5 per layer and $(K_{\mathrm{ret}},V_{\mathrm{ret}})$6 per head within a layer (Jiang et al., 29 Mar 2026).

In the second stage, a total budget

$(K_{\mathrm{ret}},V_{\mathrm{ret}})$7

is redistributed proportionally to a dampened difficulty signal,

$(K_{\mathrm{ret}},V_{\mathrm{ret}})$8

leading to

$(K_{\mathrm{ret}},V_{\mathrm{ret}})$9

The choice $T-m$0 is described as square-root dampening, intended to prevent outlier layers or heads from consuming all of $T-m$1 (Jiang et al., 29 Mar 2026).

The associated analysis reports that per-layer pilot MSE varies by up to $T-m$2 across layers, and that per-head pilot MSE within one layer can differ by up to $T-m$3. On that basis, uniform allocation is characterized as starving hard components and over-allocating easy ones, while pure proportional weighting with $T-m$4 is said to overfit to outliers. Dampened weighting is therefore presented as a compromise between sensitivity to component difficulty and robustness to extreme heterogeneity (Jiang et al., 29 Mar 2026).

A plausible implication is that KVSculpt’s contribution is not only the unconstrained optimization of compressed pairs, but also the recognition that compression is a highly non-uniform resource-allocation problem across both layers and KV heads.

5. Empirical evaluation

The reported experiments use Qwen2.5-1.5B-Instruct with context length $T-m$5, retain zone $T-m$6, and evaluation by KL divergence on 128 continuation tokens (Jiang et al., 29 Mar 2026). The principal baseline is “Select+Fit,” defined as selecting the top-$T-m$7 by attention score and fitting values by least squares.

The reported KL values at compression ratios $T-m$8 are as follows.

Compression ratio $T-m$9	Select+Fit KL	KVSculpt KL
0.3	0.233	0.0575
0.5	0.186	0.0463
0.7	0.125	0.0358

These results correspond to a $(K_{\mathrm{old}},V_{\mathrm{old}})$0–$(K_{\mathrm{old}},V_{\mathrm{old}})$1 reduction in KL for KVSculpt relative to Select+Fit (Jiang et al., 29 Mar 2026). The addition of adaptive allocation provides an extra $(K_{\mathrm{old}},V_{\mathrm{old}})$2 KL reduction “for free” at inference time. Compression of a 2048-token context is reported to take approximately 170 s on one A100 GPU with 100 L-BFGS steps per layer, after which inference speed and memory footprint improve by factor $(K_{\mathrm{old}},V_{\mathrm{old}})$3 (Jiang et al., 29 Mar 2026).

The evaluation highlights a characteristic tradeoff. KVSculpt incurs an upfront optimization cost during compression, but the post-compression runtime characteristics are those of a shorter cache. This suggests a deployment regime in which offline or amortized cache distillation is acceptable, while decoding efficiency remains the dominant operational objective.

6. Analysis, bottlenecks, and interpretation

The analysis emphasizes that compression difficulty is highly non-uniform. The reported pilot-MSE variation—up to $(K_{\mathrm{old}},V_{\mathrm{old}})$4 across layers and up to $(K_{\mathrm{old}},V_{\mathrm{old}})$5 between two KV heads within a single layer—supports the claim that fine-grained budget allocation is essential (Jiang et al., 29 Mar 2026). This reframes the compression problem from a homogeneous pruning task into a structured optimization problem with substantial inter-component variability.

Another reported result concerns local optimization quality. Single-run L-BFGS is stated to be within $(K_{\mathrm{old}},V_{\mathrm{old}})$6 of a 100-restart oracle on per-layer MSE, indicating that cross-layer error propagation, rather than per-layer suboptimality, is now the bottleneck (Jiang et al., 29 Mar 2026). This is an important diagnostic distinction: it suggests that, within the layerwise objective, optimization is already close to the attainable local optimum, and that further gains may depend more on joint coordination across layers than on stronger single-layer solvers.

The method also addresses a common misconception that value fitting alone may be sufficient. The reported comparison states that fitting only $(K_{\mathrm{old}},V_{\mathrm{old}})$7 while keeping original key positions yields almost no improvement over Select+Fit, whereas freeing the keys in continuous space produces the major gain. In that sense, KVSculpt identifies key placement as the critical degree of freedom in preserving attention behavior under aggressive sequence-length compression (Jiang et al., 29 Mar 2026).

Taken together, these results support a specific interpretation of KVSculpt. It is not merely an eviction heuristic with a better regressor, nor simply a merging scheme in another form. Rather, it treats the compressed cache as a learned surrogate set of support points for future attention computation. The stated practical takeaway is that, by viewing KV-cache compression as a continuous distillation problem, optimizing keys with an L-BFGS/least-squares loop, and adapting budgets via a cheap pilot, KVSculpt achieves large gains in attention fidelity and model output KL at aggressive compression ratios, while making long-context inference more memory-efficient (Jiang et al., 29 Mar 2026).