Kv cache eviction improves long context inference in large language models
Jacap: Robust KV Cache Eviction via Jacobian-Based Nonlinear Information Capacity Preservation
Computation and Language
Summary
Large language models use key-value caches to remember what they have seen so far, but these caches have limited space and need to decide what to keep or remove. Existing methods use simple rules that don’t fully consider how attention works in these models. The authors thought about this problem using ideas from information theory and geometry to better understand which cached tokens are most useful. They created a new method called Jacap that smartly chooses which tokens to keep, leading to better performance, especially when the cache is very limited.
What this means in practice
- •For ai infrastructure engineers: Improve long context handling in transformer-based models by using Jacap to optimize cache eviction and maintain model performance with smaller cache sizes.
- •For chatbot developers: Enable chatbots to generate more coherent responses in long conversations by applying Jacap for robust memory management of previous dialogue tokens.
Authors
Jiaming Yang, Chenwei Tang, Liangli Zhen, Chenyang Zhang, Jiancheng Lv
Abstract
Key-value (KV) cache eviction is essential for scaling long-context inference in Large Language Models. However, existing policies predominantly rely on empirical heuristics, lacking a rigorous characterization of token utility under the inherently nonlinear softmax attention mechanism. In this work, we rethink KV cache eviction through the lens of local information geometry, modeling the attention process as a nonlinear Gaussian communication channel. By performing a first-order Taylor expansion of the attention mapping, we derive the Jacobian Information Capacity, a novel objective that explicitly captures query relevance, softmax sensitivity, and structural diversity. Guided by this theory, we introduce Jacap, a capacity-aware eviction method that utilizes softmax-aware importance weighting and statistical leverage scores for subset selection. Extensive experiments across diverse architectures and benchmarks demonstrate that \textsc{Jacap} delivers superior performance in most scenarios, particularly in high-compression regimes.