Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Aomufei Yuan, Zhiming Wang, Ruijie Miao, Dayu Wang, Yuxuan Tian, Zihan Wang, Yebo Peng, Yuhan Wu, Bairen Yi, Xin Liu, Tong Yang

TL;DR

KVReviver introduces a reversible KV cache compression via a sketch-based reconstruction that eliminates Contextual Amnesia in long-context LLM inference. It partitions storage into Recent, Candidate, and Vague, with the Vague portion stored as a Count Sketch variant to reconstruct evicted tokens during attention calculations. Theoretical error bounds are derived under mild assumptions, and empirical results on 2k and 32k contexts demonstrate near-full accuracy at 10%–25% KV budgets, enabling scalable long-context inference. This approach offers a practical path to high-memory-efficiency, low-loss KV caching compatible with existing compression methods.

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.

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

TL;DR

KVReviver introduces a reversible KV cache compression via a sketch-based reconstruction that eliminates Contextual Amnesia in long-context LLM inference. It partitions storage into Recent, Candidate, and Vague, with the Vague portion stored as a Count Sketch variant to reconstruct evicted tokens during attention calculations. Theoretical error bounds are derived under mild assumptions, and empirical results on 2k and 32k contexts demonstrate near-full accuracy at 10%–25% KV budgets, enabling scalable long-context inference. This approach offers a practical path to high-memory-efficiency, low-loss KV caching compatible with existing compression methods.

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.

Paper Structure

This paper contains 32 sections, 3 theorems, 28 equations, 6 figures, 3 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Motivation
  3. Ideas of Design
  4. Contribution
  5. Background
  6. Related Work
  7. Sketch-based Approximation
  8. Contextual Amnesia
  9. Rethinking Token Sparsity
  10. Algorithm
  11. Overview
  12. Vague Sketch
  13. Recent and Candidate
  14. Mathematical Analysis
  15. Experiments
  16. ...and 17 more sections

Key Result

Lemma 1

Let $X_1, X_2, X_3$ be i.i.d. random variables with distribution $F$ and finite variance $\sigma^2$. Define the median estimator as For any symmetric distribution $F$, we have the following result. with equality if and only if $X_i$ are almost surely constant.

Figures (6)

  • Figure 1: A token-level view of Contextual Amnesia caused by Eviction and Merging.
  • Figure 2: Heatmap and Activating Frequency of different algorithms and Heavy Hitter ratios.
  • Figure 3: Overview of KVReviver system.
  • Figure 4: Rouge Metrics on short-text datasets.
  • Figure 5: NIAH Results on different methods with around 25% KV Cache budgets.
  • ...and 1 more figures

Theorems & Definitions (5)

  • Lemma 1: Variance Reduction by Median
  • proof
  • Lemma 2: Variance Reduction for Median of Three Normals
  • proof
  • Lemma 3: Variance Decoupling by Sign Randomization