KVReviver: Reversible KV Cache Compression with Sketch-Based Token Reconstruction

TL;DR

KVReviver introduces a reversible cache compression technique for large language models that significantly reduces memory usage while preserving inference accuracy, enabling efficient processing of longer contexts.

Contribution

The paper presents KVReviver, a novel sketch-based reversible KV cache compression method that prevents information loss and improves memory efficiency in LLMs.

Findings

Requires only 10% of KV cache for 2k contexts with no accuracy loss.

Achieves ~2% accuracy loss with 25% cache for 32k contexts.

Enables full-scale computation within limited memory constraints.

Abstract

As the context length of current large language models (LLMs) rapidly increases, the memory demand for the Key-Value (KV) cache is becoming a bottleneck for LLM deployment and batch processing. Traditional KV cache compression methods typically involve permanently evicting or irreversibly merging "less important" tokens with low attention scores. This approach results in the unrecoverable loss of token information, which we call Contextual Amnesia, significantly degrading the model's information retrieval capability. To address this issue, we propose KVReviver, a reversible KV cache compression method based on the sketch algorithm. This method allows reconstructing compressed tokens from an additional data structure, thus enabling full-scale computation within limited memory. Experiments showed that in 2k-length contexts, it requires only 10% of KV Cache budget while maintaining identical end-to-end inference accuracy. For 32k-length contexts, it achieves equivalent or comparable accuracy ~2% accuracy loss) using merely 25% of KV Cache budget.