Inferenzzeit-Hyper-Skalierung mit KV-Cache-Kompression

Inference-Time Hyper-Scaling with KV Cache Compression

June 5, 2025

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

cs.AI

Zusammenfassung

Inference-Time-Scaling tauscht Effizienz gegen eine erhöhte Genauigkeit der Schlussfolgerungen ein, indem längere oder stärker parallele Sequenzen erzeugt werden. Bei Transformer-LLMs wird die Generierungskosten jedoch durch die Größe des Key-Value (KV)-Caches begrenzt und nicht durch die Anzahl der erzeugten Tokens. Daher untersuchen wir Inference-Time-Hyper-Scaling: Durch die Komprimierung des KV-Caches können wir innerhalb desselben Rechenbudgets mehr Tokens erzeugen und die Genauigkeit des skalierten Inferenzprozesses weiter verbessern. Der Erfolg dieses Ansatzes hängt jedoch von der Fähigkeit der Kompressionsmethoden ab, die Genauigkeit auch bei hohen Kompressionsraten zu bewahren. Um Hyper-Scaling praktikabel zu machen, führen wir Dynamic Memory Sparsification (DMS) ein, eine neuartige Methode zur Sparsifizierung von KV-Caches, die nur 1.000 Trainingsschritte benötigt, um eine 8-fache Kompression zu erreichen, während sie eine bessere Genauigkeit als trainingsfreie Sparse-Attention beibehält. Anstatt zwischengespeicherte Tokens vorzeitig zu verwerfen, verzögert DMS die Token-Entfernung, indem es Repräsentationen implizit zusammenführt und kritische Informationen bewahrt. Wir demonstrieren die Wirksamkeit von Inference-Time-Hyper-Scaling mit DMS anhand mehrerer LLM-Familien und zeigen, dass es die Genauigkeit bei vergleichbarer Inferenzlaufzeit und Speicherlast steigert. Beispielsweise verbessern wir Qwen-R1 32B im Durchschnitt um 9,1 Punkte auf AIME 24, 7,6 auf GPQA und 9,6 auf LiveCodeBench über verschiedene Rechenbudgets hinweg.

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.