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.