Back from the Future: Key-Value Cache Management by Counter-Causal Surprise

2026-07-29

AI Summary

This work addresses the substantial GPU memory consumption caused by the linear growth of key-value (KV) cache with context length in large language model inference, which severely limits long-context generation efficiency. The authors propose a training-free KV cache compression method that leverages anti-causal attention to compute a “surprise score” for each token, dynamically pruning redundant entries that are highly predictable from subsequent tokens. To further reduce computational overhead, they introduce a single-layer Transformer approximation applied exclusively to the final layer. By integrating anti-causal masking with online cache reuse, the approach achieves state-of-the-art or comparable performance across multiple open-source large language models and benchmarks, significantly reducing memory usage and accelerating inference while preserving generation quality.

Abstract

Key-value (KV) cache management through compression and eviction strategies has emerged as an important research direction in recent years. Computational demands of large language models (LLMs) and their multi-modal variants during output generation can be partially alleviated by caching previous key and value calculations needed by subsequent scaled dot-product attention operations. However, this leads to another problem: the size of the resulting KV cache grows linearly with context length and quickly consumes all available GPU memory when either the prompt or the generated output are long. KV cache management periodically prunes entries from the cache thereby reducing its memory footprint while attempting to retain sufficient information for accurate generation. A by-product is faster inference speed. We propose a simple yet effective KV eviction scheme motivated by the insight that past tokens which can be well-predicted from more recent tokens are redundant and their associated keys and values can be removed from the cache. To score entries for eviction we run the model on the tokens in their original order, reusing the key and value representations already stored in the KV cache, and applying a counter-causal attention mask so that each position attends only to its future context. This is in-distribution, tied directly to the actual cache contents, and requires no additional training. To further reduce cost, we additionally propose a fast single-layer approximation that restricts the counter-causal pass to the last transformer layer, achieving a significant speedup per refresh cycle at marginal accuracy cost. We evaluate our strategy on various open-source LLMs and benchmark datasets showing competitive or improved performance over other state-of-the-art methods. Reference code is available at https://github.com/metacognitionai/counter_causal.

Problem

• KV cache
• memory consumption
• large language models
• context length
• cache management

Innovation

• counter-causal attention
• KV cache eviction
• in-context redundancy
• memory-efficient LLMs
• approximate inference