KV Cache Compression Through the Lens of Transform Coding

2026-08-14

AI Summary

This study addresses the KV cache memory bottleneck in long-context reasoning and the limitation of existing quantization methods that overlook attention error propagation. We propose Attention-Aware Transform Coding (AATC), which pioneers the integration of transform coding and reverse water-filling into KV cache compression. Grounded in rate-distortion theory, AATC constructs an attention-aware distortion decomposition model that optimizes bit allocation to minimize perceptual distortion. Experimental results demonstrate that AATC achieves near-lossless accuracy at a compression ratio of approximately 5.8×, significantly outperforming state-of-the-art baselines. Consequently, this work effectively resolves the critical trade-off between storage efficiency and inference precision in long-sequence scenarios.

Abstract

The key-value (KV) cache stores information from past tokens and is a major memory bottleneck in long-context inference. Existing quantization methods address this bottleneck by representing the KV cache uniformly with lower-precision data types and designing quantization schemes to minimize reconstruction error in the cache itself, without accounting for how that error propagates through attention mechanisms. We prove that, under a white-noise quantization model, the expected attention-aware distortion decomposes into additive key and value contributions that factor across tokens and channels. Building on transform coding and reverse water-filling, which are classical tools from signal processing and rate-distortion theory, we introduce Attention-Aware Transform Coding (AATC), which allocates bits over a calibration set to minimize attention-aware distortion. On Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct, evaluated across LongBench, RULER, GSM8K, MMLU-Pro, and MATH-500, our method achieves near-lossless accuracy at approximately $5.8\times$ compression, whereas each baseline degrades in at least some settings.

Problem

• KV Cache Compression
• Long-context Inference
• Memory Bottleneck
• Attention-aware Distortion

Innovation

• Attention-Aware Transform Coding
• KV Cache Compression
• Rate-Distortion Theory
• Reverse Water-Filling
• Attention-Aware Distortion