Kv cache compression improves by optimizing bits and rank per head
KV-COBRA: KV Cache Compression via Co-Optimized Bit-Rank Allocation
Machine LearningArtificial Intelligence
Summary
KV-cache stores information for AI models to remember past details, but compressing it too much can reduce accuracy. The authors found that instead of compressing all parts the same way, it works better to adjust compression differently for each attention head. They designed KV-COBRA to smartly split compression effort by balancing two types of loss within each head, improving performance especially at very low bit rates. This lets AI models keep longer memories with less accuracy loss and no extra delay per token.
What this means in practice
- •For machine learning engineers: Compress attention cache in transformer models to reduce memory use with minimal accuracy loss in low bit-rate scenarios.
- •For cloud infrastructure teams: Optimize memory footprint for AI services supporting long-context queries by adapting compression per attention head.
Authors
Sihyeon Ha, Jaeho Lee, Yo-Seb Jeon
Abstract
What limits KV-cache compression at extreme bit-rates? We argue that it is not the choice of compression scheme, but how its budget is allocated across attention heads. Existing methods apply rank and bit-width uniformly, ignoring that each head has a different optimal mix of rank truncation and quantization. We show that co-optimizing rank and bit-width per head, using only standard low-rank projection and scalar quantization, dominates uniform allocation, with the largest gains at low bit-rates. Our method, KV-COBRA (Co-Optimized Bit-Rank Allocation), formalizes this as a resource-allocation problem: it balances rank-truncation loss against quantization loss within each head, then redistributes budget across heads to minimize total distortion. A fused Hadamard rotation equalizes per-channel variance, and reordering the SVD basis by attention-KL importance makes the solver query-aware. The same allocator extends to joint $K{+}V$ compression. On perplexity, zero-shot, and long-context benchmarks from $0.5$ to $4$ bits per dimension (bpd), KV-COBRA shows the smallest accuracy degradation among evaluated methods at low bpd, with no per-token overhead.