Long-context language models store key-value pairs for every past token, and this cache becomes a memory bottleneck. Existing compression methods use hand-crafted rules to decide which tokens to keep (like "attention sinks" or recent tokens), but these rules don't optimize for the actual task. LKV learns both which attention heads need more cache budget and which tokens matter most, treating compression as an end-to-end optimization problem. On LongBench, it retains only 15% of the cache with almost no quality loss, and the learned budget allocation turns out to be more important than the token-selection method.
Main takeaways:
- Existing KV cache compression uses heuristics (recent tokens, high-attention tokens) instead of task objectives
- LKV learns head-wise budgets and token importance directly from task loss
- Achieves near-lossless performance keeping only 15% of key-value pairs
- Learned budget allocation across heads is the main driver of quality
- Training-based approach but applies at inference without recomputing attention matrices