Hiperdimensionamento em Tempo de Inferência com Compressão de Cache KV

Inference-Time Hyper-Scaling with KV Cache Compression

June 5, 2025

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

cs.AI

Resumo

A escalonamento no tempo de inferência troca eficiência por maior precisão de raciocínio ao gerar sequências mais longas ou mais paralelas. No entanto, em LLMs baseados em Transformers, o custo de geração é limitado pelo tamanho do cache de chave-valor (KV), em vez do número de tokens gerados. Portanto, exploramos o hiper-escalonamento no tempo de inferência: ao comprimir o cache KV, podemos gerar mais tokens dentro do mesmo orçamento computacional e melhorar ainda mais a precisão da inferência escalonada. O sucesso dessa abordagem, no entanto, depende da capacidade dos métodos de compressão em preservar a precisão mesmo em altas taxas de compressão. Para tornar o hiper-escalonamento prático, introduzimos a Esparsificação Dinâmica de Memória (DMS), um método novo para esparsificar caches KV que requer apenas 1K passos de treinamento para alcançar uma compressão de 8 vezes, mantendo uma precisão melhor do que a atenção esparsa sem treinamento. Em vez de descartar prematuramente os tokens armazenados em cache, o DMS atrasa a evicção de tokens, implicitamente mesclando representações e preservando informações críticas. Demonstramos a eficácia do hiper-escalonamento no tempo de inferência com DMS em várias famílias de LLMs, mostrando que ele aumenta a precisão para tempos de inferência e carga de memória comparáveis. Por exemplo, aprimoramos o Qwen-R1 32B em média 9,1 pontos no AIME 24, 7,6 no GPQA e 9,6 no LiveCodeBench em diferentes orçamentos computacionais.

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.