Learning to Evict from Key-Value Cache

TL;DR

This paper tackles the memory bottleneck of autoregressive KV caches in large language models by reframing eviction as a learning-to-rank problem. It introduces KV Policy (KVP), a framework of lightweight per-head RL agents trained offline on pre-computed KV traces using only keys and values, producing a budget-agnostic ranking for eviction that requires no changes to the underlying model. A global offline RL reward evaluates the quality of the ranking across all budgets, enabling robust eviction decisions that preserve decoding utility across long contexts and multi-turn dialogues. Empirical results on RULER and OASST2-4k show that KVP outperforms baselines and generalizes to unseen domains and longer contexts, with strong downstream performance and minimal online overhead, highlighting a scalable, query-independent approach to adaptive KV cache management.

Abstract

The growing size of Large Language Models (LLMs) makes efficient inference challenging, primarily due to the memory demands of the autoregressive Key-Value (KV) cache. Existing eviction or compression methods reduce cost but rely on heuristics, such as recency or past attention scores, which serve only as indirect proxies for a token's future utility and introduce computational overhead. We reframe KV cache eviction as a reinforcement learning (RL) problem: learning to rank tokens by their predicted usefulness for future decoding. To this end, we introduce KV Policy (KVP), a framework of lightweight per-head RL agents trained on pre-computed generation traces using only key and value vectors. Each agent learns a specialized eviction policy guided by future utility, which evaluates the quality of the ranking across all cache budgets, requiring no modifications to the underlying LLM or additional inference. Evaluated across two different model families on the long-context benchmark RULER and the multi-turn dialogue benchmark OASST2-4k, KVP significantly outperforms baselines. Furthermore, zero-shot tests on standard downstream tasks (e.g., LongBench, BOOLQ, ARC) indicate that KVP generalizes well beyond its training distribution and to longer context lengths. These results demonstrate that learning to predict future token utility is a powerful and scalable paradigm for adaptive KV cache management.

Figures (15)

• Figure 1: Future token importance for KV cache eviction. Effective KV cache eviction requires identifying tokens that will receive little or no future attention. (center) We roll out a sample KV cache and measure the true cumulative future attention for each token, then rank tokens by this importance and color them accordingly (bright = high rank, white = low). (right) Importance estimated by the fixed sink-and-recency heuristic of StreamingLLM deviates substantially from true importance ranking. (left) Our learned policy closely recovers the complex, non-local structure of future attention despite using only past keys and values, without access to queries, attention scores, or future tokens. We expand the comparison with all strategies in \ref{['fig:teaser_ext']}.
• Figure 2: Overall accuracy on the RULER benchmark (Left) and perplexity on the OASST2-4k test set (Right), as a function of the absolute KV cache size. achieves the highest accuracy and lowest perplexity across most of the possible cache sizes.
• Figure 3: Average test accuracy on (Left) BoolQ and (Right) ARC Challenge as a function of KV cache size. Higher is better.
• Figure 4: Average test L-Rouge score on the GovReport as a function of KV cache size. Higher is better.
• Figure 5: Cost per-budget ($-\mathcal{R}^b$) on the OASST2 test set for a representative head (layer 10, head 0).

Theorems & Definitions (2)

• Proposition 1
• proof