Résumé Technique : KVpop — Compression du Cache Clé-Valeur par Élagage Prédictif en Ligne

1. Énoncé du Problème

Le décodage autorégressif dans les modèles de langage de type Transformer repose sur un cache Clé-Valeur (KV) pour éviter de recomput l'ensemble des représentations des jetons passés. Cependant, le cache KV croît de manière linéaire avec la longueur du contexte, créant un goulot d'étranglement significatif pour la mémoire et la bande passante lors de l'inférence à long contexte.

Les solutions existantes pour réduire ce coût se divisent généralement en deux catégories :

• Éviction Statique/Heuristique : Des méthodes comme la fenêtre glissante (Sliding Window) ou la préservation des jetons "Sink" (ex: StreamingLLM) reposent sur des règles fixes, tandis que les méthodes basées sur des scores utilisent des proxys en ligne (ex: attention cumulative). Celles-ci échouent souvent à suivre l'utilité future des jetons, ce qui rend l'éviction fragile lorsque la pertinence change.
• Éviction Apprise : Des approches comme la Compression de Mémoire Dynamique (DMC) ou la Sparsification de Mémoire Dynamique (DMS) entraînent des politiques pour fusionner ou supprimer des jetons. Cependant, beaucoup de ces méthodes prennent des décisions d'éviction au moment de l'insertion du jeton ou reposent sur des relaxations différentiables qui ne supervisent pas directement la décision spécifique de "garder ou supprimer" basée sur l'utilité future. De plus, les méthodes avec des fenêtres protégées retardent le retrait des jetons mais décident souvent de leur sort dès leur entrée, ignorant les preuves qui s'accumulent pendant que le jeton reste dans la fenêtre protégée.

Le défi central est de prédire quels jetons seront utiles pour de futures requêtes, car des jetons localement saillants peuvent devenir non pertinents, tandis que des jetons avec une attention immédiate faible peuvent devenir critiques plus tard (par exemple, dans des traces de raisonnement).

2. Méthodologie : KVpop

KVpop (KV compression with Predictive Online Pruning) est un retrofit d'attention parcimonieuse conçu pour imposer un budget KV fixe par tête sans modifier l'architecture du modèle de base. Il répond aux limites des travaux précédents grâce à deux mécanismes principaux : la Supervision par Attention Future et le Scoring Étatisé Retardé.

2.1 Architecture et Budget

KVpop maintient un cache borné par tête KV défini par le budget $B = s + w + k$ :

• $s$ Jetons Sink : Les premiers jetons sont toujours conservés.
• $w$ Fenêtre Protégée : Les $w$ jetons les plus récents sont protégés de l'éviction.
• $k$ Longue Portée Top-K : Les jetons plus anciens sont en compétition pour un budget fixe de $k$ emplacements basés sur des scores d'importance appris.

2.2 Supervision par Cible d'Attention Future

Contra�à les méthodes qui scorant les jetons sur la base de l'attention passée ou de statistiques locales, KVpop supervise le scoreur en utilisant la masse d'attention future qu'un jeton reçoit après avoir quitté la fenêtre protégée.

• Définition de la Cible : Pour un jeton $t$, la cible est la masse logarithmique agrégée de l'attention qu'il reçoit des requêtes futures (positions $d \ge t + w$).
• Calcul Efficace : Pour éviter le coût en $O(S^2)$ de la matérialisation de cartes d'attention denses pendant l'entraînement, KVpop utilise un passage d'attention transposé. En inversant les rôles des requêtes et des clés et en réutilisant les normalisateurs log-sum-exp (LSE) renvoyés par les noyaux d'attention parcimonieuse, le système calcule la cible d'attention future efficacement sans matérialisation de matrice dense.
• Fonction de Perte : L'objectif d'entraînement est une perte logistique par paire sensible à la frontière. Elle compare le score d'un nouveau jeton éligible ($t_{new}$) au score du jeton de coupure de l'enseignant ($t_{bnd}$). La perte pousse l'étudiant à classer correctement $t_{new}$ par rapport à la frontière, garantissant que la sélection top-$k$ correspond au classement d'utilité future de l'enseignant.

2.3 Scoring Étatisé avec Décision Retardée

KVpop introduit un scoreur à mémoire retardée (implémenté via un mLSTM) qui diffère la décision de scoring jusqu'à ce que le jeton atteigne la frontière d'éviction.

• Mécanisme : Au lieu de scorer un jeton immédiatement lors de son insertion, le scoreur met à jour son état interne avec la représentation du jeton et le contexte subséquent. Lorsque le jeton sort de la fenêtre protégée ($q - w$), le scoreur lit la mémoire accumulée pour générer le score d'importance.
• Avantage : Cela permet au scoreur d'incorporer le contexte de "proche futur" (les preuves qui s'accumulent pendant que le jeton est dans la fenêtre protégée) avant de prendre la décision de garder ou de supprimer, traitant ainsi la limitation du scoring à l'insertion.

2.4 Efficacité de l'Inférence

Pendant l'inférence, KVpop impose le budget fixe $B$. La sélection des $k$ jetons de longue portée est optimisée à l'aide d'un seuil courant sur des priorités de jetons statiques et d'un arbre de Fenwick, permettant de générer le masque parcimonieux à la volée au sein du noyau d'attention (par exemple, FlexAttention) sans stocker de masques denses $S \times S$.

3. Principales Contributions

1. Politique d'Éviction Supervisée : Les auteurs formulent l'éviction KV à budget fixe comme un problème de prédiction supervisée à la frontière d'éviction, en utilisant une cible d'attention future pour entraîner les décisions top-$k$ de longue portée.
2. Calcul de Cible Efficace : Ils introduisent un passage d'attention transposé uniquement pour l'entraînement afin de calculer les cibles d'attention future sans matérialiser de cartes d'attention denses, garantant aucun surcoût lors de l'inférence.
3. Scoring Retardé : Ils démontrent qu'un scoreur étatisé peut exploiter le contexte du proche futur en différant la décision de scoring jusqu'à ce que le jeton devienne évictable, surpassant les méthodes de scoring immédiat.
4. Validation Empirique : La méthode est évaluée sur le raisonnement mathématique (AIME, HMMT) et des tâches hors domaine (GPQA-D, LiveCodeBench), montrant que les politiques d'éviction apprises peuvent maintenir des performances élevées sous une compression agressive.

4. Résultats Expérimentaux

Les auteurs ont évalué KVpop sur les modèles Qwen3-4B et Qwen3-8B, entraînés sur le jeu de données Nemotron-Math v2 avec une longueur de séquence de 16 384.

4.1 Raisonnement Mathématique (AIME & HMMT)

• Ratio de Compression 75% (Budget 2048) :
  • Qwen3-4B : KVpop a conservé 98% de la performance de l'enseignant à attention complète.
  • Qwen3-8B : KVpop a conservé 100% (performance quasi identique à l'enseignant).
  • Comparaison : KVpop a nettement surperformé les baselines sans entraînement (StreamingLLM, TOVA) et la baseline apprise DMS. Par exemple, à 75% de compression, StreamingLLM est tombé à ~76% de performance relative sur Qwen3-4B, tandis que KVpop maintenait 98%.
• Ratio de Compression 88% (Budget 4096) :
  • Qwen3-4B : KVpop a conservé 97% de la performance de l'enseignant.
  • Qwen3-8B : KVpop a conservé 100% de la performance de l'enseignant.
  • Comparaison : Les méthodes sans entraînement se sont effondrées de manière significative (ex: StreamingLLM est tombé à ~47% de performance relative), tandis que KVpop est resté robuste.

4.2 Généralisation Hors Domaine

Évalué sur GPQA-Diamond (raisonnement scientifique) et LiveCodeBench (génération de code) :

• Les modèles KVpop entraînés uniquement sur des données de raisonnement mathématique se sont bien généralisés à ces domaines.
• À 75% de compression, KVpop sur Qwen3-4B a atteint une précision comparable à l'enseignant dense sur GPQA-D (0.59 vs 0.59) et des résultats compétitifs sur LCB.

4.3 Efficacité de l'Inférence

• Mémoire : Contrairement à l'attention dense, où l'utilisation de la VRAM augmente linéairement avec la longueur de la génération (atteignant 36 Go à 131k jetons), KVpop et DMS maintiennent une empreinte mémoire presque constante (19 Go) quel que soit la longueur de la séquence.
• Latence : KVpop a démontré une latence de décodage de bout en bout plus faible que DMS et l'attention dense pour des longueurs de génération élevées (jusqu'à 131k jetons). Les auteurs attribuent cela au budget fixe par tête de KVpop, qui évite les structures de cache "irrégulières" de DMS qui compliquent la compilation et l'exécution.

4.4 Études d'Ablation

• Scoring Retardé : Supprimer le délai (scoring immédiat lors de l'insertion) a entraîné une chute de 0.2 point de la précision des jetons, confirmant le bénéfice de l'intégration du contexte du proche futur.
• Stateless vs. Stateful : La variante étatisée (KVpop) a généralement égalé ou légèrement surpassé la variante sans état (KVpopmlp), particulièrement à des ratios de compression plus élevés.

5. Signification et Revendications

L'article affirme que superviser l'éviction avec des signaux d'attention future est une stratégie supérieure pour la compression du cache KV par rapport aux heuristiques statiques ou aux scores proxys. En entraînant directement une politique pour prédire quels jetons seront importants pour les requêtes futures, KVpop atteint :

1. Haute Fidélité : Il maintient une performance proche de l'enseignant même à des ratios de compression élevés (jusqu'à 88%), un régime où d'autres méthodes échouent.
2. Robustesse : La politique apprise se transfère efficacement aux tâches hors domaine, suggérant que le modèle apprend des principes généraux d'utilité des jetons plutôt que de sur-apprendre des distributions de données spécifiques.
3. Efficacité : Il impose une limite mémoire stricte, permettant un décodage à long contexte stable avec une utilisation de mémoire constante et un débit amélioré par rapport aux méthodes de gating dynamique comme DMS.

Les auteurs concluent que KVpop démontre que les politiques d'éviction apprises, lorsqu'elles sont correctement supervisées par l'utilité future et qu'elles sont autorisées à utiliser le contexte du proche futur, peuvent résoudre efficacement le goulot d'étranglement de la mémoire dans la génération autorégressive sans sacrifier la qualité du modèle.