Hiperescalado en Tiempo de Inferencia con Compresión de Caché 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

Resumen

El escalado en tiempo de inferencia intercambia eficiencia por una mayor precisión en el razonamiento al generar secuencias más largas o más paralelas. Sin embargo, en los modelos de lenguaje grandes (LLM) basados en Transformers, el costo de generación está limitado por el tamaño de la caché de clave-valor (KV), en lugar del número de tokens generados. Por lo tanto, exploramos el hiper-escalado en tiempo de inferencia: al comprimir la caché KV, podemos generar más tokens dentro del mismo presupuesto computacional y mejorar aún más la precisión del escalado de inferencia. El éxito de este enfoque, sin embargo, depende de la capacidad de los métodos de compresión para preservar la precisión incluso en ratios de compresión elevados. Para hacer que el hiper-escalado sea práctico, introducimos la Esparsificación Dinámica de Memoria (DMS, por sus siglas en inglés), un método novedoso para esparsificar las cachés KV que solo requiere 1,000 pasos de entrenamiento para lograr una compresión de 8 veces, manteniendo una mejor precisión que la atención dispersa sin entrenamiento. En lugar de descartar prematuramente los tokens almacenados en caché, DMS retrasa la eliminación de tokens, fusionando implícitamente representaciones y preservando información crítica. Demostramos la efectividad del hiper-escalado en tiempo de inferencia con DMS en múltiples familias de LLM, mostrando que aumenta la precisión para un tiempo de inferencia y carga de memoria comparables. Por ejemplo, mejoramos Qwen-R1 32B en un promedio de 9.1 puntos en AIME 24, 7.6 en GPQA y 9.6 en LiveCodeBench en diferentes presupuestos computacionales.

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.