Value-Aware Stochastic KV Cache Eviction for Reasoning Models
2026-06-02 • Machine Learning
Machine LearningComputation and Language
AI summaryⓘ
The authors study how to make reasoning models more efficient by removing unnecessary information from their memory (KV cache) without losing accuracy. They found that some parts of the memory are very important and removing them causes the model to make repeated mistakes. They also discovered that randomly choosing some parts to remove helps keep the model accurate by keeping the memory diverse. Based on this, they created VaSE, a method that smartly avoids removing important memory pieces and adds randomness when deciding what to remove. When tested, VaSE made reasoning models both faster and more accurate compared to other approaches.
Reasoning ModelsKV CacheEviction MethodsSparse AttentionKey-Value PairsCache CompressionFlashAttention2Stochastic EvictionModel AccuracyMemory Footprint
Authors
Ting-Yun Chang, Harvey Yiyun Fu, Deqing Fu, Chenghao Yang, Jesse Thomason, Robin Jia
Abstract
Reasoning models improve accuracy through extended chains of thought, but their long outputs create a memory and compute bottleneck. KV cache eviction methods reduce this cost by evicting unimportant key-value pairs from the cache, yet they often yield worse accuracy than selection-based sparse attention alternatives, which keep the full KV cache. We identify key factors crucial to KV cache eviction accuracy. First, a small fraction of value states have abnormally large magnitudes, and evicting them causes catastrophic failure where models enter repetitive reasoning loops. Second, introducing stochasticity during eviction improves accuracy by increasing cache diversity. Based on these findings, we propose Value-aware Stochastic KV Cache Eviction (VaSE), a training-free recipe that protects large-magnitude value states and promotes diverse eviction decisions. Across six reasoning tasks, Qwen3 models using VaSE with 4x KV cache compression yield higher average accuracies than SOTA selection method at the same sparsity, while outperforming the strongest eviction method by more than 4%. Overall, VaSE bridges the gap between efficiency and accuracy, supporting FlashAttention2 and enabling a static memory footprint for reasoning models.