KVSculpt: Reformulating KV Cache Compression as a Knowledge Distillation Problem

The long-context reasoning capability of large language models supports many applications, but the memory overhead of KV cache has become a bottleneck for deployment. Existing compression methods have limitations such as anchoring to original KV entries. KVSculpt innovatively reformulates KV cache compression as a knowledge distillation problem: it breaks away from anchoring to original entries, optimizes KV pairs in a continuous embedding space to preserve attention behavior, and introduces an adaptive budget allocation mechanism. Experiments show that on the Qwen2.5-1.5B model, KVSculpt achieves a 3.5-4.1x reduction in KL divergence, significantly improving compression effectiveness.

Memory Dilemma of Long-Context Reasoning

In the Transformer architecture, KV cache is a key data structure for self-attention. Each generated token requires storing the corresponding key and value vectors, leading to linear memory expansion as context length increases (e.g., a 70B model with 8192 tokens in FP16 precision requires tens of GB of GPU memory), limiting deployment and inference efficiency.

Limitations of Existing Compression Methods

Existing methods fall into two categories:

1. Pair-wise Compression (quantization, low-rank decomposition): Reduces storage per KV pair, but aggressive quantization loses information and the low-rank assumption does not always hold;
2. Sequence Length Compression (pruning, merging): Reduces the number of KV entries but anchors to original entries, limiting compression flexibility.

The core innovation of KVSculpt lies in breaking away from anchoring to original KV entries and reformulating compression as a knowledge distillation problem:

• Distillation Perspective: Treat compression as a distillation task where a small number of KV pairs approximate the attention behavior of the original model, optimizing new KV pairs in a continuous embedding space;
• Alternating Optimization Strategy: Key vectors are iteratively optimized using L-BFGS (a quasi-Newton algorithm suitable for nonlinear problems), while value vectors are obtained via a closed-form solution using least squares, balancing efficiency and stability.

KVSculpt introduces an adaptive budget allocation mechanism to address the non-uniformity of compression difficulty:

• Compression Difficulty Disparity: The mean squared error (MSE) of compression varies by 100x or even hundreds of times across different layers and heads of the model; a uniform compression ratio easily leads to resource misallocation;
• Adaptive Allocation: Evaluate the compression difficulty of each component via offline trial runs, allocate budgets on demand (retain more capacity for hard-to-compress components), which does not increase inference overhead and further reduces KL divergence by 1.3x.

Experimental validation on the Qwen2.5-1.5B-Instruct model:

• Comparison with Existing Methods: Under 2048-token context, KVSculpt outperforms the Select+Fit method significantly across three compression ratios, achieving a 3.5-4.1x reduction in KL divergence (KL divergence reflects the degree of deviation in attention behavior; a lower value means quality is closer to the original model);
• Effect of Adaptive Allocation: At the same compression ratio, adaptive allocation further reduces KL divergence by an additional 1.3x compared to uniform allocation, with no increase in inference overhead.

Implications for Long-Context Reasoning

1. The distillation perspective can break through traditional compression limitations and open up new optimization spaces for model compression;
2. Adaptive resource allocation is a universal principle that can be extended to other resource allocation scenarios;
3. The application of continuous optimization methods to discrete problems is worth exploring.

Conclusion

By reformulating compression as a distillation problem and combining continuous space optimization with adaptive allocation, KVSculpt achieves efficient compression. As the context length of large models increases, such technologies will become key enablers for long-context applications.

KVSculpt has the following limitations and future directions:

1. Currently focuses on KV cache compression in the pre-filling phase; dynamic cache management in the decoding phase needs optimization;
2. The computational cost of offline optimization is relatively high and needs further acceleration;
3. Can be combined with other compression techniques such as quantization to achieve more aggressive compression;
4. The adaptive budget allocation strategy can be improved (e.g., more efficient difficulty estimation, online adjustment).

Paper link: http://arxiv.org/abs/2603.27819v1