Hold Onto That Thought: Assessing KV Cache Compression On Reasoning

2025-12-12

AI Summary

This work systematically evaluates KV cache compression for multi-step reasoning tasks with short prompts and long decoding (e.g., GSM8K, MATH500), focusing on its impact on LLM inference quality and reasoning chain length. It is the first to concentrate specifically on *decoding-phase* compression, proposing a decoding-enhanced variant of SnapKV and benchmarking it against major existing strategies. Results show that hit-rate–driven methods (e.g., H2O, SnapKV) significantly outperform prefill-oriented approaches. Surprisingly, cache budget and reasoning chain length exhibit a counterintuitive negative correlation: certain strategies generate longer and more accurate chains under *lower* cache budgets. Experiments on Llama-3.1-8B-Instruct confirm strong task dependency—no universally optimal compression strategy exists. Key contributions include: (1) establishing the critical importance of decoding-phase KV compression; (2) introducing a lightweight, inference-chain–aware compression paradigm; and (3) discovering a novel cache–reasoning-length trade-off phenomenon.

Technology Category

• Data Mining & Knowledge Management: Data Compression
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning
• Search and Optimization: Metareasoning and Metaheuristics

Application Category

• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Large language models (LLMs) have demonstrated remarkable performance on long-context tasks, but are often bottlenecked by memory constraints. Namely, the KV cache, which is used to significantly speed up attention computations, grows linearly with context length. A suite of compression algorithms has been introduced to alleviate cache growth by evicting unimportant tokens. However, several popular strategies are targeted towards the prefill phase, i.e., processing long prompt context, and their performance is rarely assessed on reasoning tasks requiring long decoding. In particular, short but complex prompts, such as those in benchmarks like GSM8K and MATH500, often benefit from multi-step reasoning and self-reflection, resulting in thinking sequences thousands of tokens long. In this work, we benchmark the performance of several popular compression strategies on long-reasoning tasks. For the non-reasoning Llama-3.1-8B-Instruct, we determine that no singular strategy fits all, and that performance is heavily influenced by dataset type. However, we discover that H2O and our decoding-enabled variant of SnapKV are dominant strategies for reasoning models, indicating the utility of heavy-hitter tracking for reasoning traces. We also find that eviction strategies at low budgets can produce longer reasoning traces, revealing a tradeoff between cache size and inference costs.

Problem

• Assess KV cache compression on reasoning tasks
• Evaluate compression strategies for long decoding sequences
• Analyze tradeoff between cache size and inference costs

Innovation

• Benchmarking KV cache compression strategies
• Heavy-hitter tracking for reasoning traces
• Tradeoff between cache size and inference costs