🤖 machine learning

Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction

This paper introduces a global retention-based KV cache eviction method that learns to selectively discard irrelevant tokens to reduce memory usage while improving long-context reasoning performance by mitigating attention dilution, thereby surpassing full-cache inference across diverse benchmarks.

Original authors: Ngoc Bui, Hieu Trung Nguyen, Arman Cohan, Rex Ying

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ngoc Bui, Hieu Trung Nguyen, Arman Cohan, Rex Ying

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

The Big Problem: The "Too-Much-Information" Bottleneck

Imagine you are a brilliant detective trying to solve a complex mystery. To do your job, you have a giant whiteboard where you write down every single clue, witness statement, and piece of evidence you encounter.

In the world of AI (specifically Large Language Models), this whiteboard is called the KV Cache. As the AI reads a long story or looks at a long video, it writes every word and image pixel onto this board so it doesn't forget the beginning of the story while writing the end.

The Catch: The whiteboard gets huge. If the story is 100,000 words long, the board is massive.

It runs out of space: Eventually, the board is so full that the AI can't fit new clues on it.
It gets too slow: To find the one important clue (like "the butler did it"), the detective has to scan through thousands of irrelevant notes (like "the butler had a red tie" or "it was raining"). This slows everything down.

The Old Solution: The "First-In, First-Out" Trash Can

To fix the space issue, previous methods acted like a strict janitor. They would say, "We are full! Throw away the oldest notes to make room for new ones."

The Flaw: This is dangerous. Sometimes the oldest note is the most important one (e.g., "The butler owns a gun"). Throwing it away just because it's old makes the AI stupid. Other methods tried to be smarter by looking at which words the AI was currently looking at, but they were often short-sighted, throwing away things that would be useful five minutes later.

The New Solution: The "Smart Retention" System (TrimKV)

This paper introduces a new method called TrimKV (or DBTrimKV). Instead of just throwing away old things, it asks a better question: "Which clues will actually help me solve the mystery in the future?"

Here is how it works, using three simple concepts:

1. The "Future Utility" Score

Imagine every piece of evidence on your whiteboard has a little sticky note on it. This sticky note predicts how useful that clue will be later in the investigation.

High Score: "This gun is crucial. Keep it forever."
Low Score: "This red tie is irrelevant. Throw it away."

The AI learns to write these scores automatically. It doesn't just look at what is happening right now; it looks at what will be needed tomorrow.

2. The "Global Competition" (The Great Filter)

In the past, different parts of the AI (different "layers" and "heads") had their own separate little whiteboards with their own limits. If one board was full, it threw things away even if another board had space.

This new method creates one giant, shared whiteboard for the whole AI.

The Analogy: Imagine a VIP party. In the old way, every room had a bouncer who let in 10 people. If a VIP was in the hallway, they couldn't get in because the room was full.
The New Way: There is one bouncer for the whole club. The VIPs (the most useful clues) get to enter, regardless of which room they came from. The "distractors" (irrelevant background noise) get kicked out, even if they were in a "VIP" room.

3. Fighting "Attention Dilution"

The paper argues that having too much information actually hurts the AI.

The Analogy: Imagine trying to hear a friend whispering in a quiet room. You can hear them perfectly. Now, imagine that same friend whispering in a stadium full of 10,000 people screaming. You can't hear them anymore, even though they are still whispering.
The Result: By aggressively throwing away the "screaming crowd" (irrelevant tokens), the AI can hear the "whisper" (the important evidence) much better. Sometimes, deleting information actually makes the AI smarter because it stops it from getting distracted.

What Did They Find?

The researchers tested this on difficult tasks, like solving math problems, analyzing long videos, and having long conversations.

It saves space: They could reduce the memory usage by a huge amount (sometimes keeping only 10-20% of the original data) without the AI getting confused.
It improves performance: In many cases, the AI performed better with less memory than with full memory. This proves that "less is more" when it comes to removing distractions.
It works for images and text: It successfully managed both text and visual data (like images in a video) together, deciding which pixels and words to keep.

Summary

This paper teaches AI to be a better detective. Instead of hoarding every single piece of paper it finds, it learns to identify the "Golden Nuggets" of information and throw away the "Rock and Dirt." By doing this, it runs faster, uses less memory, and actually solves problems more accurately because it isn't distracted by the noise.

Based on the paper "Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction," here is a detailed technical summary.

Problem Statement

The Key-Value (KV) cache is essential for efficient autoregressive decoding in Transformer-based Large Language Models (LLMs) and Vision-Language Models (VLMs). However, as context windows expand to tens of thousands of tokens (including text and visual tokens), the KV cache becomes a primary bottleneck due to linear memory growth and quadratic attention computation costs.

Existing solutions typically treat KV eviction as a compression problem: removing tokens to approximate full-cache inference under a memory budget. These methods often rely on heuristics (e.g., recency, attention magnitude) or fixed per-layer/per-head budgets. The authors argue that this approach is suboptimal because it assumes full-cache inference is the ideal target. In reality, long contexts often contain many irrelevant "distractor" tokens that dilute attention mass away from useful evidence, degrading model performance. The paper posits that selective eviction, if done correctly, can not only reduce memory but also improve reasoning by suppressing these distractors.

Methodology: Global Retention-Based Eviction

The authors propose DBTrimKV (Dynamic Budget TrimKV), a method that learns to predict the future utility of cached tokens and allocates a single global memory budget across all layers, heads, and modalities.

1. Theoretical Insight: Attention Dilution
The paper provides a theoretical analysis showing that full-cache attention suffers from "attention dilution." When a large number of distractor tokens exist in the cache, the softmax normalization forces useful tokens to compete with them, causing the attention mass assigned to relevant evidence to vanish. The authors prove that preferentially retaining useful tokens while evicting distractors can reduce this dilution and theoretically improve attention quality.

2. Geometric Retention as Future Utility
Instead of using query-dependent attention scores (which reflect short-term relevance), the method learns a geometric retention score ( $\beta_i$ ) for each token.

Concept: A token's utility is modeled as a persistence event. If a token is likely to remain useful for future decoding steps, it receives a high retention score.
Formulation: The retention weight for token $i$ at step $t$ is defined as $r_{t,i} = \beta_i^{t-i}$ .
Justification: The authors theoretically justify this geometric decay as a query-agnostic proxy for future utility, supported by empirical observations that token attention survival drops rapidly over time in long contexts.

3. Weight-Tied Global Gates
To enable a unified eviction policy, the authors introduce weight-tied retention gates:

Architecture: A lightweight gate predicts a scalar retention score $\beta$ for every KV entry. Crucially, the final projection layer of these gates is shared (tied) across all transformer layers and attention heads.
Global Calibration: This weight sharing forces retention scores from different layers and heads onto a common scale. Consequently, tokens from a text-heavy head can be directly compared against tokens from a visual head or a deep layer.
Global Budget: Instead of assigning fixed budgets per layer or head, the system maintains a single global KV budget ( $M_{global}$ ). At any eviction step, the system ranks all cached KV entries globally by their predicted future utility and retains the top $M_{global}$ tokens.

4. Implementation
The method uses a PagedAttention layout to support variable-length KV sequences per head, allowing for dynamic, head-specific cache sizes without contiguous physical memory requirements.

Key Contributions

Reframing Eviction: The paper challenges the assumption that KV eviction is merely a compression technique. It demonstrates that selective eviction can act as a corrective mechanism against attention dilution, potentially improving reasoning performance over full-cache inference.
Global Retention Policy: It introduces a learnable, globally calibrated eviction policy that dynamically allocates memory across layers, heads, and modalities based on a single utility signal, replacing rigid, hand-designed budget rules.
Theoretical and Empirical Validation: The authors provide theoretical bounds showing that preferential retention reduces attention dilution and justify geometric retention as a valid surrogate for future utility.

Experimental Results

The method was evaluated on diverse benchmarks including long-horizon reasoning (MathVision, AIME), multi-turn dialogue (MMDU), and visual question answering (VQAText, MME).

Performance vs. Full Cache: In many scenarios, particularly under tight memory budgets, DBTrimKV matched or surpassed the performance of full-cache (Vanilla) inference. For example, on the MMDU multi-turn dialogue benchmark, DBTrimKV with a budget of 512 tokens outperformed the full-cache baseline by 14%.
Comparison to Baselines: DBTrimKV consistently outperformed existing KV eviction methods (e.g., SnapKV, AdaKV, R-KV) and visual token pruning baselines (e.g., VisionZip, FastV).
Efficiency: The method significantly reduced KV memory usage and decoding time compared to full-cache inference, scaling efficiently with context length and generation length.
Ablation Studies: Removing the weight-tied projection (forcing local scores) resulted in significant performance drops, confirming the necessity of global calibration.

Significance and Claims

The paper claims that learned, globally calibrated KV eviction is not just a compression technique, but a mechanism for improving long-context reasoning.

By treating eviction as a way to suppress distractors rather than just save memory, the authors show that "forgetting" irrelevant tokens can sharpen attention and enhance model accuracy. The results suggest that dynamic, learned memory management can outperform static, full-cache inference in long-context settings, offering a path toward more efficient and capable long-horizon AI systems. The authors note limitations regarding the scale of models tested (up to 8B parameters) and the potential for future work in end-to-end joint training of retention gates and base models.