🤖 machine learning

DistillCache: KL-Guided Adaptive KV-Cache Eviction for Memory-Efficient LLM Inference

DistillCache introduces a reinforcement learning framework that learns a lightweight policy to adaptively evict KV-cache tokens based on internal model signals and KL-divergence rewards, achieving superior memory efficiency and accuracy retention for long-context LLM inference compared to existing heuristic and concurrent methods.

Original authors: Asaad Althoubi

Published 2026-08-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Asaad Althoubi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a very long story while you are telling it to a friend. As the story gets longer, your brain gets crowded with every single detail you've ever said, making it hard to think about what to say next. This is exactly the problem facing modern "Large Language Models" (AI chatbots). These AI brains are amazing at writing, coding, and solving puzzles, but they work by remembering everything they've said so far in a special memory area called a "KV cache." The trouble is, this memory grows bigger and bigger the longer the conversation gets, eventually filling up the computer's memory and slowing everything down to a crawl.

To fix this, scientists have tried to be like a strict librarian, throwing away old pages of the story that seem unimportant. Some librarians just throw away the oldest pages, while others throw away pages that nobody looked at much. But these rules are often too simple; they might accidentally toss out a tiny word that holds the key to the whole ending, causing the AI to forget the plot or make up nonsense. The big question is: How can we teach the AI to be a smarter librarian, one that knows exactly which memories to keep and which to forget, so it can tell long stories without running out of brain space?

Enter DistillCache, a new method that teaches the AI to make these decisions using a clever form of trial-and-error learning. Instead of following a rigid rulebook, DistillCache acts like a student learning to play a video game. It tries different strategies for throwing away old memories and gets a "score" based on how well the AI's next guess matches what it would have guessed if it had kept all the memories. If the AI forgets something important and its guess goes off track, it gets a low score. If it keeps the right memories and stays on track, it gets a high score. Over time, the AI learns a "policy"—a set of instincts—that tells it exactly which tokens (the tiny pieces of text) to keep and which to evict to stay within a strict memory limit.

The researchers tested this on a popular AI model called Mistral-7B. They found that when they forced the AI to keep only 25% of its usual memory (a huge cut!), DistillCache was still able to get 94.2% of the accuracy it would have had with the full memory. This is a big deal because other methods that use simple rules (like H2O or SnapKV) dropped much further in accuracy under the same pressure. Even compared to other smart, learning-based methods that were developed around the same time, DistillCache came out on top in many long-story tasks, beating them by up to 2.7 points on accuracy.

What makes DistillCache special is how it learns. While other methods might look at how often a word was mentioned in the past, DistillCache looks at the "future impact" of a word. It asks, "If I throw this word away right now, will the AI's next sentence sound weird?" It uses a mathematical measure called KL-divergence to check if the AI's predictions are still sounding like the original, full-memory version. The paper shows that this approach is particularly good at keeping the AI's logic intact, even when the memory is squeezed tight.

However, the paper is careful to note that this isn't a magic wand that solves everything. Training this smart librarian takes time and extra computer power (about 130 hours on powerful graphics cards) because the AI has to practice twice for every step—once with the full memory and once with the pruned memory—to learn the difference. Also, while DistillCache is great at general long stories, another method called RLKV was still slightly better at specific logic puzzles at moderate memory levels, though DistillCache caught up when the memory was squeezed even tighter.

In the end, DistillCache suggests that the best way to manage an AI's memory isn't with a static rule, but with a learned intuition that constantly checks, "Am I still making sense?" By doing this, it allows AI models to handle much longer conversations and documents without needing massive amounts of computer memory, potentially making long-context AI faster and more accessible for everyone. The researchers also found that this "smart librarian" trained on one type of AI model could be used on a different model without any extra training, suggesting these memory instincts might be universal.

Technical Summary: DistillCache

Problem Statement

Transformer-based Large Language Models (LLMs) rely on Key-Value (KV) caches to avoid recomputing token representations during autoregressive decoding. However, the size of the KV cache grows linearly with sequence length, creating a severe memory bottleneck that limits the context window and throughput of long-context inference.

Existing solutions primarily rely on heuristic eviction strategies (e.g., H2O, SnapKV) that retain tokens based on static attention statistics or positional signals. The author argues that these heuristics often fail to capture a token's future predictive influence, particularly in scenarios involving long-range dependencies or structured reasoning where token importance is contingent on downstream generation consequences rather than past attention patterns.

Methodology: DistillCache

The paper proposes DistillCache, a reinforcement learning (RL) framework that formulates KV-cache eviction as a sequential decision-making problem under strict memory budgets. Unlike heuristic methods, DistillCache learns a policy to retain a subset of tokens that preserves the output distribution of a full-cache model.

Core Components

State Space (Feature Representation):
For each token in the cache, the policy observes a 5-dimensional feature vector derived from internal model signals available during decoding:
- $\alpha_i$ : Average attention mass received by the token across all heads and layers.
- $\|v_i\|_2$ : The $\ell_2$ norm of the token's value vector.
- $\sigma_i$ : Variance of the attention mass received across heads.
- $H_i$ : Attention-weighted average entropy of recent queries attending to the token (capturing context-sensitive importance).
- $pos_i$ : Normalized sequence position.
Action Space:
A lightweight Multi-Layer Perceptron (MLP) policy network assigns a scalar score to each cached token. During training, a stochastic ranking is induced by adding Gumbel noise to these scores, allowing the policy to sample a top- $B$ subset of tokens (where $B$ is the budget) from the Plackett–Luce distribution. At inference, the top- $B$ tokens are selected deterministically.
Training Architecture (Teacher-Student Setup):
To train the policy without ground-truth labels, the author employs a teacher-student distillation approach:
- Teacher Cache: Maintains the full, unpruned history of tokens to compute the "gold" next-token distribution ( $P_t$ ).
- Student Cache: Maintains the pruned history according to the policy's decisions to compute the pruned distribution ( $\hat{P}_t$ ).
- Reward Signal: The reward at each step is defined as the negative Kullback-Leibler (KL) divergence between the teacher and student distributions: $R_t = -D_{KL}(P_t \parallel \hat{P}_t)$ . This objective directly optimizes for distributional fidelity at every decoding step, rather than relying on proxy signals like attention scores.
Optimization:
The policy is optimized using the REINFORCE algorithm with entropy regularization and a variance-reduction baseline (Leave-One-Out). The training process is performed offline, doubling the memory footprint temporarily, but inference requires only a single forward pass with negligible overhead for feature extraction and policy scoring.

Key Contributions

RL Framework for Cache Eviction: Introduction of DistillCache, a general-purpose RL framework that optimizes token retention for downstream distributional fidelity under memory constraints, moving beyond static heuristic scoring.
Distribution-Preserving Objective: Demonstration that optimizing per-step KL divergence yields measurable gains over heuristic and supervised alternatives by explicitly learning long-horizon trade-offs.
Comprehensive Evaluation: Extensive evaluation on Mistral-7B-Instruct-v0.3 and zero-shot transfer to Llama-3-8B-Instruct across LongBench, reasoning benchmarks (GSM8K, BBH), and throughput metrics.

Experimental Results

Evaluated on Mistral-7B-Instruct-v0.3 with a 25% KV-cache budget:

LongBench Performance: DistillCache retains 94.2% of the full-cache accuracy, outperforming strong heuristic baselines (H2O, SnapKV) by up to 2.7 absolute points and concurrent RL-based methods (ForesightKV, RLKV) by up to 1.4 points.
Reasoning Tasks: On reasoning benchmarks, DistillCache is competitive with the best concurrent methods and surpasses them under aggressive compression (12.5% budget).
Cross-Model Transfer: The policy trained on Mistral-7B transfers zero-shot to Llama-3-8B-Instruct, maintaining competitiveness with heuristic baselines without requiring per-model retraining (a requirement for the concurrent RL baselines).
Efficiency: The method delivers up to 2.1× full-cache throughput. While it introduces a small overhead (~1.4 ms/token) for feature extraction and policy evaluation, it remains comparable to other learned methods and significantly faster than full-cache inference.
Error Analysis: DistillCache exhibits the lowest rates of repetitive errors and incorrect answers compared to heuristics and other RL methods, suggesting superior preservation of sequence coherence and factual accuracy.

Significance and Claims

The paper posits that DistillCache demonstrates the effectiveness of learned, distribution-aware policies for memory-efficient long-context LLM inference. By optimizing token retention decisions against future distributional fidelity rather than retrospective attention statistics, the framework learns budget-constrained policies that account for downstream generation consequences.

The author emphasizes that DistillCache is not merely a learned proxy for heuristic metrics but a distinct approach that directly optimizes for the preservation of the model's output distribution. While acknowledging limitations such as training costs (approx. 130 GPU-hours) and a slight performance gap with RLKV on reasoning tasks at moderate budgets, the paper concludes that distribution-preserving adaptive eviction is a promising direction for enabling longer-context inference on resource-constrained hardware. The work suggests that combining token-level distributional fidelity with head-level gating strategies (as explored in hybrid experiments) offers a complementary path for future optimization.