LKV: End-to-End Learning of Head-wise Budgets and Token Selection for LLM KV Cache Eviction
Authors: Enshuai Zhou, Yifan Hao, Chao Wang et al.
Summary
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
Relevance
Tangential to my work — LKV learns internal resource allocation (which heads get cache budget) in a task-driven way, somewhat analogous to how I'm studying which behavioral patterns get installed during fine-tuning, though the technical connection is loose.
Abstract
arXiv:2605.06676v1 Announce Type: new Abstract: Long-context inference in Large Language Models (LLMs) is bottlenecked by the linear growth of Key-Value (KV) cache memory. Existing KV cache compression paradigms are fundamentally limited by heuristics: heuristic budgeting relies on statistical priors rather than task objectives, causing resource misallocation, while heuristic selection relies on coupled query-key interactions or static inductive biases (e.g., attention sinks). To address this limitation, we introduce LKV (Learned KV Eviction), which formulates KV compression as an end-to-end differentiable optimization problem. LKV integrates LKV-H to learn task-optimized global budgets, and LKV-T to derive intrinsic KV importance without materializing attention matrices. This design bypasses heuristic proxies, strictly aligning compression with task objectives. Extensive evaluations demonstrate that LKV achieves state-of-the-art performance on both LongBench and RULER benchmarks at high compression rates. In particular, on LongBench, LKV achieves near-lossless performance with only 15% KV cache retention. Crucially, our analysis identifies learned budgeting as the dominant driver of fidelity, demonstrating that data-driven allocation is essential to overcome the limitations of hand-crafted heuristics.