Résumé Technique : Ablation, Inférence Statistique et Validation pour la Compression du KV-Cache

Énoncé du Problème

L'inférence des Transformers à grande échelle est contrainte par la bande passante mémoire, l'accès au cache Key-Value (KV) dominant le trafic mémoire total lors de la génération de contextes longs. Bien que la quantification des clés et des valeurs à 2–4 bits par élément soit une approche standard pour réduire la taille du cache, l'évaluation de ces schémas de compression présente un défi important. Les différences de qualité observées sont souvent confondues par les propriétés distributionnelles du trafic réel des LLM, les effets matériels et les hypothèses algorithmiques, ce qui rend difficile l'attribution des gains ou des échecs de performance à des choix de conception spécifiques.

Méthodologie

Les auteurs introduisent un cadre d'évaluation contrôlé composé de six régimes statistiques synthétiques conçus pour isoler les hypothèses structurelles spécifiques du pipeline de compression (par exemple, la forme de la distribution, l'alignement du sous-espace K/V, la décroissance des valeurs propres). Ce cadre est couplé à une méthodologie de validation statistique pour distinguer les différences algorithmiques systématiques de la variance induite par l'implémentation.

Configuration Expérimentale

• Schémas Évalués : Deux familles principales de compression sont comparées :
  1. TurboQuant (TQ) : Une approche agnostique aux données utilisant une transformée de Walsh-Hadamard (WHT) aléatoire pour répartir l'énergie des tokens de manière uniforme, suivie d'un codebook Beta Lloyd-Max agnostique aux données. Elle utilise optionnellement un résidu de sketch Johnson-Lindenstrauss à 1 bit (QJL) sur les chemins de la clé (K) ou de la valeur (V).
  2. SpectralQuant (SQ) : Une approche adaptative aux données qui calibre une base propre par tête à partir de tokens représentatifs, alloue les bits via un remplissage par niveau (water-filling) aux dimensions à haute variance, et utilise un sketch QJL sélectif sur les dimensions effectives supérieures.
• Métriques : L'étude utilise une métrique d'erreur 2D ($d^2$) combinant l'erreur quadratique moyenne relative (relMSE) et l'erreur de direction (similitude cosinus). Elle emploie également des métriques d'isolation (relMSE_K, relMSE_V) et une géométrie d'erreur en 6D pour décomposer les erreurs dans le cache K, le cache V et la sortie d'attention finale.
• Validation : Pour séparer la performance algorithmique du bruit d'implémentation, les auteurs comparent les implémentations Python (oracle) et HIP/GPU (production) utilisant différents générateurs de nombres aléatoires (PCG64 vs mt19937). Le test de Kolmogorov-Smirnov (KS) est appliqué pour identifier les différences systématiques par rapport à la variance dépendante de la graine (seed).

Contributions Clés

1. Ablation Systématique des Schémas de Compression

À travers une ablation complète de huit variantes de TQ et de plusieurs configurations de SQ, l'étude identifie trois schémas non dominés :

• Quantification Scalaire (Plain-KV) : Sans rotation, efficace sur des données isotropes.
• TQ-KV : Rotation WHT avec un codebook Beta Lloyd-Max.
• TQ-QKV : TQ-KV augmenté d'un sketch résiduel QJL à 1 bit sur le chemin de la clé.
  L'étude élimine les variantes QJL sur le chemin V, démontrant que l'application de QJL aux valeurs (qui sont sommées sur les tokens) est une heuristique qui dégrade systématiquement les performances par rapport aux corrections sur le chemin K.

2. Validation Statistique et Inégalité de Jensen

Le papier introduit une méthodologie pour séparer les différences de codec du bruit d'implémentation. Une conclusion critique est l'identification de l'inégalité de Jensen comme mécanisme d'amplification de l'erreur.

• La non-linéarité softmax est convexe ; ainsi, la variance d'une erreur de correction de moyenne nulle ($\epsilon$) ajoutée aux scores d'attention est exponentiellement amplifiée ($E[\exp(s+\epsilon)] \approx \exp(s) \cdot \exp(\sigma^2/2)$).
• La correction QJL sur le chemin K introduit une telle variance. L'étude montre que la variance de l'écart Python-HIP pour TQ-QKV est significativement plus grande que pour TQ-KV, confirmant que le chemin K est une signature directe de l'amplification de la variance des scores par l'inégalité de Jensen. Inversement, les corrections sur le chemin V entrent linéairement dans la sortie et ne déclenchent pas cette amplification.

3. Performance Spécifique au Régime et Sensibilité au Calibrage

• Données à Queue Lourde (Heavy-Tailed) : Pour les données avec des queues lourdes (ex: distributions de Student-t), les méthodes basées sur l'autobasis (SQ) échouent de manière catastrophique. Les covariances d'échantillon sont déstabilisées par les valeurs aberrantes, conduisant à une base calibrée systématiquement mal alignée. Aucune augmentation du budget de bits ne permet de récupérer cette perte. TQ-KV domine dans ces régimes car la WHT répartit les valeurs aberrantes des queues lourdes, rendant les marginales projetées approximativement gaussiennes.
• Données Structurées : Sur des régimes de faible rang et structurés, SQ surpasse TQ lorsque K et V sont calibrés sur des ensembles représentatifs distincts. Cependant, l'étude trouve que l'allocation de bits par remplissage par niveau (water-filling) apporte un bénéfice négligeable par rapport à une allocation uniforme dans tous les régimes testés, car les ratios de valeurs propres au sein du sous-espace effectif dépassent rarement le seuil requis pour justifier une allocation non uniforme.
• Dimension Sémantique Effective ($d_{eff}$) : Le papier caractérise $d_{eff}$ comme s'auto-calibrant au budget de calibration disponible plutôt que de récupérer le vrai rang des données. Une seule valeur propre fallacieuse peut causer un échec catastrophique (sortie anti-corrélée), nécessissant un critère de gap spectral pour élaguer la dimension estimée.

Résultats et Recommandations

Performance Comparative

• Budget $b=2$ : Aucun schéma ne domine de manière unique. Les variantes SQ gagnent sur les régimes structurés (faible rang), TQ-KV gagne sur les données à queue lourde, et Plain-KV est compétitif sur les données isotropes.
• Budget $b \ge 4$ : TQ-KV affirme sa dominance dans la plupart des régimes. Les variantes SQ ne conservent un avantage que sur les régimes hautement structurés où l'autobasis calibré peut concentrer efficacement les bits, mais elles ne parviennent pas à se rétablir de l'alignement erroné de l'autobasis sur les données à queue lourde.
• Stockage vs Qualité : Bien que SQ puisse atteindre des coûts de stockage inférieurs à des budgets plus élevés grâce à sa capacité à ignorer les dimensions de queue, le "plancher de queue" (bits fixes pour les dimensions non modélisées) devient une contrainte contraignante. Pour les séquences longues, l'accumulation d'erreurs de queue érode l'avantage sémantique de SQ, faisant de TQ-KV l'option la plus sûre pour les tâches de génération.

Recommandations Pratiques

Les auteurs concluent par des directives spécifiques de déploiement :

1. Utiliser TQ-KV lorsque la distribution des données est inconnue, possède des queues lourdes ou est sujette à un changement de distribution, et pour tout budget supérieur à 2 bits.
2. Utiliser SQU-KV (SpectralQuant avec allocation uniforme) uniquement si les données sont connues pour être structurées (sous-espace KV de faible rang), que le budget est faible (2–3 bits), que la longueur de séquence est courte (charges de travail de prefill) et que des ensembles de calibration distincts pour K et V sont disponibles.
3. Abandonner le Water-Filling : Sur les données synthétiques, l'allocation uniforme est équivalente au water-filling.
4. Éviter le QJL sur K pour les Queues Lourdes : La correction de score QJL sur les clés amplifie la variance via le softmax (effet Jensen) sur les données à queue lourde et doit être évitée dans ces régimes.

Signification

La signification de ce papier réside dans sa méthodologie rigoureuse et contrôlée pour évaluer la compression du cache KV. En isolant les hypothèses structurelles via des régimes synthétiques et en appliquant l'inférence statistique pour séparer le signal algorithmique du bruit d'implémentation, les auteurs fournissent un cadre de diagnostic clair. Ce travail va au-delà du simple rapport de précision pour expliquer pourquoi les méthodes réussissent ou échouent, révélant les limitations fondamentales des méthodes basées sur l'autobasis sur les données à queue lourde et le comportement non monotone de la dimension effective sous des budgets de calibration finis. Ces conclusions offrent un guide définitif pour la sélection des schémas de compression en fonction des caractéristiques des données et des contraintes opérationnelles.