Inference-Time Hyper-Scaling with KV Cache Compression

2025-06-05

AI Summary

In Transformer-based large language models, the KV cache size dominates computational and memory overheads, limiting generation length and accuracy. To address this, we propose Dynamic Memory Sparsification (DMS), a differentiable, inference-efficient KV cache compression method. DMS integrates delayed token pruning with implicit representation fusion, achieving 8× KV compression after only 1K fine-tuning steps—substantially outperforming training-free sparse attention baselines. Crucially, DMS enables inference-time scaling without increasing latency or memory footprint, thereby enhancing long-context generation capability. Evaluated on Qwen-R1-32B, DMS improves performance by +9.1, +7.6, and +9.6 points on AIME 2024, GPQA, and LiveCodeBench, respectively. These results demonstrate a significant advancement in the accuracy–efficiency trade-off for KV cache management.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Data Mining & Knowledge Management: Data Compression
• Natural Language Processing: (Large) Language Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Inference-time scaling trades efficiency for increased reasoning accuracy by generating longer or more parallel sequences. However, in Transformer LLMs, generation cost is bottlenecked by the size of the key-value (KV) cache, rather than the number of generated tokens. Hence, we explore inference-time hyper-scaling: by compressing the KV cache, we can generate more tokens within the same compute budget and further improve the accuracy of scaled inference. The success of this approach, however, hinges on the ability of compression methods to preserve accuracy even at high compression ratios. To make hyper-scaling practical, we introduce Dynamic Memory Sparsification (DMS), a novel method for sparsifying KV caches that only requires 1K training steps to achieve 8$ imes$ compression, while maintaining better accuracy than training-free sparse attention. Instead of prematurely discarding cached tokens, DMS delays token eviction, implicitly merging representations and preserving critical information. We demonstrate the effectiveness of inference-time hyper-scaling with DMS on multiple families of LLMs, showing that it boosts accuracy for comparable inference runtime and memory load. For instance, we enhance Qwen-R1 32B by an average of 9.1 points on AIME 24, 7.6 on GPQA, and 9.6 on LiveCodeBench across compute budgets.

Problem

• KV cache compression for efficient inference scaling
• Maintaining accuracy at high compression ratios
• Dynamic Memory Sparsification for delayed token eviction

Innovation

• KV cache compression for efficient scaling
• Dynamic Memory Sparsification (DMS) method
• Delayed token eviction preserves critical information