Que chaque token compte : Vers l'amélioration des performances en contexte long avec l'éviction du cache KV

Résumé

Le cache clé-valeur (KV) est un goulot d'étranglement majeur dans l'inférence sur de longs contextes, où la mémoire et le calcul augmentent avec la longueur de la séquence. Les méthodes existantes d'éviction du cache KV réduisent ce coût mais dégradent généralement les performances par rapport à l'inférence avec cache complet. Notre idée clé est que l'attention avec cache complet n'est pas toujours optimale : dans les longs contextes, des tokens non pertinents peuvent diluer l'attention en l'éloignant des preuves utiles ; ainsi, une éviction sélective et apprenable peut améliorer la génération plutôt que de simplement approximer le cache complet. Nous introduisons une méthode d'éviction du cache KV basée sur la rétention globale qui apprend l'utilité future de chaque token sous un budget mémoire unifié. Des portes de rétention légères attribuent des scores d'utilité aux entrées KV mises en cache, et une projection de score finale partagée calibre ces scores à travers toutes les couches et toutes les têtes. Cela permet une politique d'éviction globale unique dans laquelle des tokens de différentes couches, têtes et modalités entrent en compétition directe pour la capacité du cache. Nous fournissons en outre une analyse théorique montrant que la rétention préférentielle des tokens utiles réduit la dilution de l'attention, et nous justifions la rétention géométrique comme un proxy indépendant de la requête pour l'utilité future. À travers divers benchmarks de raisonnement linguistique et vision-langage sur de longs contextes, et de dialogues multi-tours, notre méthode réduit considérablement la mémoire KV tout en égalant ou surpassant l'inférence avec cache complet. Ces résultats suggèrent que l'éviction apprise et globalement calibrée du cache KV n'est pas seulement une technique de compression, mais aussi un mécanisme pour améliorer le raisonnement sur de longs contextes.

English

The key-value (KV) cache is a major bottleneck in long-context inference, where memory and computation grow with sequence length. Existing KV eviction methods reduce this cost but typically degrade performance relative to full-cache inference. Our key insight is that full-cache attention is not always optimal: in long contexts, irrelevant tokens can dilute attention away from useful evidence, so selective, learnable eviction can improve generation rather than merely approximate the full cache. We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget. Lightweight retention gates assign utility scores to cached KV entries, and a shared final scoring projection calibrates these scores across all layers and heads. This enables a single global eviction policy in which tokens from different layers, heads, and modalities compete directly for cache capacity. We further provide theoretical analysis showing that preferentially retaining useful tokens reduces attention dilution, and we justify geometric retention as a query-agnostic proxy for future utility. Across diverse long-context language and vision-language reasoning, and multi-turn dialogue benchmarks, our method substantially reduces KV memory while matching or surpassing full-cache inference. These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.