Mise à l'échelle hyper-optimisée au moment de l'inférence avec compression du cache KV

Inference-Time Hyper-Scaling with KV Cache Compression

June 5, 2025

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

cs.AI

Résumé

Le dimensionnement au moment de l'inférence échange l'efficacité contre une précision accrue en générant des séquences plus longues ou plus parallèles. Cependant, dans les modèles de langage de type Transformer, le coût de génération est limité par la taille du cache clé-valeur (KV), plutôt que par le nombre de tokens générés. Par conséquent, nous explorons l'hyper-dimensionnement au moment de l'inférence : en compressant le cache KV, nous pouvons générer plus de tokens dans le même budget de calcul et améliorer encore la précision de l'inférence dimensionnée. Le succès de cette approche dépend toutefois de la capacité des méthodes de compression à préserver la précision même à des taux de compression élevés. Pour rendre l'hyper-dimensionnement pratique, nous introduisons la Sparsification Dynamique de la Mémoire (DMS), une méthode novatrice pour sparsifier les caches KV qui ne nécessite que 1 000 étapes d'entraînement pour atteindre une compression de 8 fois, tout en maintenant une meilleure précision que l'attention sparse sans entraînement. Au lieu de supprimer prématurément les tokens mis en cache, DMS retarde l'éviction des tokens, fusionnant implicitement les représentations et préservant les informations critiques. Nous démontrons l'efficacité de l'hyper-dimensionnement au moment de l'inférence avec DMS sur plusieurs familles de modèles de langage, montrant qu'il améliore la précision pour un temps d'inférence et une charge mémoire comparables. Par exemple, nous améliorons Qwen-R1 32B en moyenne de 9,1 points sur AIME 24, 7,6 sur GPQA et 9,6 sur LiveCodeBench pour différents budgets de calcul.

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.