Résumé Technique : COMPILERKV

1. Énoncé du Problème

Les modèles de langage de grande taille (LLM) font face à un goulot d'étranglement mémoire dû à la croissance linéaire du cache clé-valeur (KV) avec la longueur du contexte. La compression KV uniquement en préremplissage (prefill-only) est apparue comme une solution compatible avec le déploiement, où un sous-ensemble de tokens est sélectionné à la fin de l'étape de préremplissage et figé pour l'ensemble du passage de décodage. Cette approche convertit l'utilisation de la mémoire à l'exécution en un coût fixe, sans intervention au moment du décodage.

Cependant, l'article identifie un décalage statistique fondamental dans les méthodes existantes de préremplissage uniquement (par exemple, SnapKV, PyramidKV, HeadKV). Ces méthodes tentent d'estimer des signaux correctifs — spécifiquement la fiabilité par tête (quelles têtes d'attention sont critiques pour la récupération) et la sensibilité à la compression au niveau du prompt (dans quelle mesure un prompt spécifique peut être compressé) — en ligne, en utilisant des statistiques issues d'un seul prompt bruyant. Les auteurs soutiennent qu'il s'agit de la mauvaise unité statistique :

• E1 (Bruit du prompt unique) : Les scores d'attention par prompt sont à queue lourde et biaisés par l'échelle ; de grandes magnitudes ne garantissent pas une importance sémantique, et de petites erreurs d'estimation provoquent des erreurs de sélection irréversibles.
• E2 (Hétérogénéité des têtes) : Les têtes d'attention sont fonctionnellement spécialisées. Certaines sont « critiques pour la récupération », tandis que d'autres sont constamment bruyantes. Une agrégation naïve permet aux têtes bruyantes de dominer les scores.
• E3 (Variance du risque au niveau du prompt) : Différents prompts ont des limites de compression différentes. Une règle de rétention fixe échoue sur les prompts à haute entropie ou à haute perplexité, entraînant une perte d'information catastrophique.

Les méthodes existantes réestiment ces signaux en ligne pour chaque prompt, échouant à tirer parti du fait que ces signaux présentent une forte régularité inter-prompt plutôt qu'une stabilité intra-prompt.

2. Méthodologie : COMPILERKV

Les auteurs proposent COMPILERKV, une politique de rétention KV qui compile des tables correctives hors ligne à partir d'un corpus de calibration, plutôt que de les estimer en ligne. L'idée centrale est que les classements de fiabilité des têtes et la cartographie des caractéristiques de risque vers des seuils sûrs sont des a priori architecturaux stables qui peuvent être appris une fois et réutilisés.

Le cadre se compose de trois composants couplés opérant pendant l'étape de préremplissage :

A. Utilité des Tokens Stabilisée (Résolution de E1)

Pour éliminer les pics d'attention transitoires et le biais d'échelle inter-têtes sans apprentissage, la méthode calcule un score d'utilité sans paramètre :

• Attention agrégée par fenêtre : Lisse les pics transitoires en agrégeant l'attention sur une fenêtre d'observation $\Omega$ à la queue du prompt.
• Normalisation relative à la tête : Normalise les scores des tokens par la norme de la valeur moyenne de la tête, éliminant ainsi le biais d'échelle.
• Formule : $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$, où $\alpha_t$ est la masse de la fenêtre et $\rho$ la norme de la valeur normalisée.

B. Table d'Hétérogénéité des Têtes (Résolution de E2)

Au lieu de traiter toutes les têtes de manière égale ou de réestimer l'importance en ligne, COMPILERKV compile une table figée $W_{head}[l, h]$ de poids de fiabilité continus.

• Compilation Hors Ligne : Apprise via Q-Learning Conservateur (CQL) sur un corpus de calibration. L'état est $(l, h, B_l)$, et l'action est un poids de fiabilité discrétisé.
• Mécanisme : Pendant l'inférence, l'utilité d'un token est redimensionnée par la fiabilité de sa tête. Crucialement, l'agrégation utilise un max-pooling pondéré ($\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$). Cela donne un « pouvoir de veto » aux têtes critiques pour la récupération : si une seule tête critique s'attache fortement à un token, il est retenu, indépendamment du consensus des têtes bruyantes.

C. Porte de Seuil Adaptative au Risque (Résolution de E3)

Pour gérer les prompts ayant des limites de compression variables, une Porte de Seuil Adaptative au Risque mappe les caractéristiques de risque au niveau du prompt vers un seuil de rétention.

• Signaux de Risque : Utilise l'Entropie d'Attention (risque structurel) et la Perplexité Locale (risque sémantique) calculées sur la fenêtre d'observation.
• Compilation Hors Ligne : Une table de recherche (LUT) mappe les bins de risque discrétisés $(b_{ent}, b_{ppl})$ vers un seuil de rétention $\tau^{(l)}$. Ceci est également compilé via CQL.
• Biais Conservateur : L'objectif CQL pénalise les seuils agressifs sur les états rares à haut risque (haute entropie/perplexité), biaisant la table vers une rétention plus sûre là où les heuristiques en ligne échouent généralement.
• Sélection Élastique : La sélection finale retient tous les tokens dont l'utilité dépasse le seuil $\tau^{(l)}$. Si l'ensemble des candidats dépasse le budget $B_l$, il est limité aux $B_l$ meilleurs tokens. S'il est inférieur, tous les candidats sont conservés (sous-rétention élastique).

D. Fondement Théorique

L'article fournit une analyse formelle montrant que l'erreur d'approximation est contrôlée par la masse de queue rejetée (Proposition 4.1). Il soutient que l'erreur de calibration de table finie est contrôlée au niveau du prompt (à travers des prompts indépendants) plutôt que d'assumer l'indépendance au sein des lignes d'attention d'un seul prompt, réduisant ainsi la variance.

3. Contributions Clés

1. Recadrage du Problème : Identifie que l'unité statistique pour la correction dans la compression uniquement en préremplissage doit être l'échelle inter-prompt, et non le prompt unique.
2. Compilation Hors Ligne : Introduit une méthode pour compiler la fiabilité des têtes et les seuils adaptatifs au risque hors ligne, les transformant en a priori architecturaux portables.
3. Politique Couplée : Propose une optimisation conjointe où les poids de fiabilité des têtes et les seuils de risque sont calibrés par rapport à la même récompense en aval, plutôt que d'être des modules indépendants.
4. Portabilité : Démontre que les tables compilées se transfèrent à travers des corpus disjoints et même entre différentes familles de modèles (par exemple, de LLaMA à Mistral) avec une dégradation de performance minimale.

4. Résultats Expérimentaux

La méthode a été évaluée sur quatre architectures de base de 7B–8B (LLaMA-3, Mistral, Qwen2, InternLM) à travers les benchmarks LongBench, RULER et Needle-in-a-Haystack (NIAH).

• LongBench (budget de 512 tokens) : COMPILERKV atteint un SOTA compressé sur les quatre architectures de base, surpassant la meilleure base de référence (par exemple, ChunkKV, HeadKV) d'une moyenne de +1,67 points (IC 95 % [+1,08, +2,37]). Sur certaines architectures, il dépasse légèrement les performances de FullKV, ce que les auteurs attribuent à un débruitage plutôt qu'à un gain d'information.
• Régimes de Pression :
  • RULER (contexte 128k) : Sous un ratio de cache fixe de 512/32k, COMPILERKV maintient un score d'environ 73, surpassant significativement SnapKV (38) et KVZip (66,5), tout en restant en dessous de FullKV (~79).
  • NIAH (32k) : Atteint une précision de 0,89 contre 0,42 pour SnapKV.
  • Regroupement (Batching) : À une entrée de 32k avec une taille de lot de 16, COMPILERKV reste faisable, alors que FullKV épuise la mémoire (OOM).
• Transférabilité :
  • Inter-Corpus : Les classements de fiabilité des têtes montrent une forte corrélation de Spearman ($\bar{\rho} \approx 0,90$) à travers des corpus de calibration disjoints.
  • Inter-Modèle : Le transfert direct de tables entre modèles (par exemple, de LLaMA-3 à Mistral) entraîne une perte d'à peine 0,4–0,8 points sur LongBench par rapport à une compilation spécifique à la cible.
• Efficacité : La surcharge en ligne est minimale : $O(1)$ recherches de table par couche et un seul tri sur le pool de candidats. Il atteint jusqu'à 2,45× le débit de décodage par rapport à FullKV à un contexte de 32k.

5. Signification et Revendications

L'article revendique que COMPILERKV comble l'écart avec FullKV à budgets fixes en corrigeant le décalage statistique dans la manière dont les politiques de rétention sont calculées. En déplaçant le calcul des signaux correctifs d'une estimation en ligne, bruyante et par prompt unique, vers un a priori hors ligne, compilé et inter-prompt, la méthode stabilise le processus de sélection contre le bruit et l'hétérogénéité.

Les auteurs soulignent que la nouveauté ne réside pas simplement dans la « conscience des têtes » (que possèdent des méthodes antérieures comme HeadKV), mais dans la compilation conjointe hors ligne de la fiabilité continue des têtes et des seuils de risque des prompts. Cela transforme la fiabilité des têtes d'un artifice de réglage en un objet scientifique réutilisable qui peut être compilé une fois et déployé à travers des modèles et des contextes avec une surcharge en ligne de $O(1)$. Le travail suggère que pour la compression uniquement en préremplissage, les « bons » signaux sont des propriétés architecturales du modèle, et non des caractéristiques transitoires du prompt d'entrée.