ChatPaper.aiChatPaper

Iper-scalabilità al Tempo di Inferenza con Compressione della Cache KV

Inference-Time Hyper-Scaling with KV Cache Compression

June 5, 2025
Autori: Adrian Łańcucki, Konrad Staniszewski, Piotr Nawrot, Edoardo M. Ponti
cs.AI

Abstract

Il ridimensionamento al momento dell'inferenza scambia efficienza con una maggiore accuratezza nel ragionamento, generando sequenze più lunghe o più parallele. Tuttavia, nei modelli linguistici Transformer (LLM), il costo della generazione è limitato dalla dimensione della cache chiave-valore (KV), piuttosto che dal numero di token generati. Pertanto, esploriamo l'iper-ridimensionamento al momento dell'inferenza: comprimendo la cache KV, possiamo generare più token mantenendo lo stesso budget computazionale e migliorare ulteriormente l'accuratezza dell'inferenza scalata. Il successo di questo approccio, tuttavia, dipende dalla capacità dei metodi di compressione di preservare l'accuratezza anche ad alti rapporti di compressione. Per rendere pratico l'iper-ridimensionamento, introduciamo la Sparsificazione Dinamica della Memoria (DMS), un metodo innovativo per sparsificare le cache KV che richiede solo 1K passi di addestramento per ottenere una compressione di 8 volte, mantenendo un'accuratezza migliore rispetto all'attenzione sparsa senza addestramento. Invece di scartare prematuramente i token memorizzati nella cache, DMS ritarda l'evizione dei token, fondendo implicitamente le rappresentazioni e preservando le informazioni critiche. Dimostriamo l'efficacia dell'iper-ridimensionamento al momento dell'inferenza con DMS su diverse famiglie di LLM, mostrando che migliora l'accuratezza mantenendo tempi di inferenza e carico di memoria comparabili. Ad esempio, miglioriamo Qwen-R1 32B in media di 9,1 punti su AIME 24, 7,6 su GPQA e 9,6 su LiveCodeBench, indipendentemente dal budget computazionale.
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.
PDF271June 6, 2025