arxiv.org web signal

AbridgeKV cuts reasoning KV cache to a quarter via pause tokens

TL;DR

  • AbridgeKV trains a model to compress each completed block of generated reasoning into the KV entries of a few learned pause tokens, then frees the block's original cache.
  • On 446 competition math problems with 32K-token generations, keeping a quarter of the generated cache matched an identically trained full-KV model on accuracy and output length.
  • The paper reports 3.5x cache reduction per sequence and 3.2x faster problem completion on saturated H100 hardware.

"Long chain-of-thought decoding is memory-bound: the KV cache of the generated tokens grows linearly with the test-time compute being spent." That is the opening claim of a new preprint on arxiv, "AbridgeKV: Compressing Long Reasoning into the KV Cache of Learned Pause Tokens," by Jack Cai, Lijie Yang, Vasilis Kontonis, Dimitris Papailiopoulos and Shivam Garg. The fix the authors propose is to train the model to compress its own output.

"AbridgeKV compresses each completed block of generated text into the KV entries of a few learned pause tokens," they write. "These are extra input tokens with trained embeddings that produce no output. Once the pause tokens have encoded the block, AbridgeKV frees the block's original cache; subsequent tokens attend to the pause tokens' KV entries in its place."

On 446 competition math problems drawn from AIME, HMMT, AMC, BRUMO, CMIMC and SMT, with 32K-token generations, the paper reports that keeping a quarter of the generated cache matches an identically trained full-KV model in accuracy and output length. The hardware numbers the authors claim: 3.5x cache reduction per sequence and 3.2x faster problem completion on saturated H100s.

Prior cache-eviction tricks fail here, the authors argue, because a model that was never trained with entries removed just keeps writing. "Evicting cache entries does not fix this, because a model never trained on a cache with entries removed generates longer, often up to the length limit." They credit the gains to three design choices: positioning pause-token KV entries within the replaced blocks, allocating cache budget across layers by measured value, and operating on large blocks rather than token-by-token.

The abstract reports results only on competition math and does not benchmark AbridgeKV against other KV compression schemes.

Shared on Bluesky by 2 AI experts