Inferentie-Tijd Hyper-Scaling met KV Cache Compressie

Inference-Time Hyper-Scaling with KV Cache Compression

June 5, 2025

Auteurs: Adrian Łańcucki, Konrad Staniszewski, Piotr Nawrot, Edoardo M. Ponti

cs.AI

Samenvatting

Inferentie-schaalbaarheid wisselt efficiëntie in voor een hogere nauwkeurigheid in redenering door langere of meer parallelle sequenties te genereren. Echter, in Transformer LLM's wordt de generatiekosten beperkt door de grootte van de key-value (KV) cache, in plaats van het aantal gegenereerde tokens. Daarom onderzoeken we inferentie-hyperschaalbaarheid: door de KV cache te comprimeren, kunnen we meer tokens genereren binnen hetzelfde rekenbudget en de nauwkeurigheid van geschaalde inferentie verder verbeteren. Het succes van deze aanpak hangt echter af van het vermogen van compressiemethoden om de nauwkeurigheid te behouden, zelfs bij hoge compressieverhoudingen. Om hyperschaalbaarheid praktisch te maken, introduceren we Dynamic Memory Sparsification (DMS), een nieuwe methode voor het verspreiden van KV caches die slechts 1K trainingsstappen vereist om een 8-voudige compressie te bereiken, terwijl een betere nauwkeurigheid wordt behouden dan trainingsvrije sparse attention. In plaats van tokens voortijdig te verwijderen, stelt DMS het verwijderen van tokens uit, waarbij representaties impliciet worden samengevoegd en kritieke informatie wordt behouden. We demonstreren de effectiviteit van inferentie-hyperschaalbaarheid met DMS op meerdere families van LLM's, waarbij we aantonen dat het de nauwkeurigheid verhoogt bij vergelijkbare inferentie-runtime en geheugenbelasting. Zo verbeteren we bijvoorbeeld Qwen-R1 32B gemiddeld met 9,1 punten op AIME 24, 7,6 op GPQA en 9,6 op LiveCodeBench over verschillende rekenbudgetten.

English

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 8times 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.