Résumé technique : HeadQ : Distortion visible par le modèle et correction dans l'espace des scores pour la quantification du cache KV

1. Énoncé du problème

Les méthodes actuelles de quantification du cache KV optimisent principalement la reconstruction de l'espace de stockage, en minimisant typiquement l'erreur quadratique moyenne (MSE) entre les clés et les valeurs originales et quantifiées. Cependant, les auteurs soutiennent que cet objectif est inadapté à la manière dont les mécanismes d'attention consomment réellement les données du cache.

• Les clés sont lues via les logits (scores requête-clé), ce qui signifie que seules des projections linéaires spécifiques des clés influencent la distribution d'attention.
• Les valeurs sont lues via une lecture pondérée par l'attention ($\sum p_j v_j$).

L'article postule que l'erreur persistante du cache doit être mesurée dans des coordonnées visibles par le modèle plutôt que dans des coordonnées de stockage euclidiennes brutes. Une perturbation peut avoir une grande MSE de stockage mais un impact nul sur l'attention (si elle se situe dans une direction « softmax-nulle »), ou une faible MSE de stockage mais des dommages catastrophiques pour l'attention (si elle affecte la direction « visible par les scores »). Cette discordance conduit à des stratégies de quantification sous-optimales qui gaspillent des bits sur des composants non pertinents tout en échouant à protéger ceux qui sont comportementalement critiques.

2. Fondement théorique : Géométrie visible par le modèle

L'article établit un cadre géométrique distinguant l'erreur de stockage de l'erreur comportementale :

2.1 Espace des scores des clés (Géométrie de Fisher)

Pour une requête fixe $q$, une perturbation de clé $\delta k_j$ induit une perturbation de score $\delta s_j = q^\top \delta k_j / \sqrt{d}$.

• Directions softmax-nulles : Une translation partagée par le token (ajout d'un vecteur constant $\mu$ à toutes les clés) ajoute une constante à tous les scores. Puisque le softmax est invariant aux décalages constants, cette perturbation est softmax-nulle et ne cause aucun dommage à l'attention, malgré une MSE de stockage potentiellement élevée.
• Espace visible par les scores : L'attention dépend uniquement des coefficients actifs centrés de la perturbation de clé par rapport à la requête.
• Information de Fisher : La géométrie locale d'ordre deux de la distribution softmax est capturée par la matrice d'information de Fisher. L'article démontre que la divergence KL entre les distributions d'attention originale et perturbée est localement proportionnelle à une variance pondérée par Fisher des erreurs de scores (Éq. 6), et non à la MSE brute des clés.

2.2 Géométrie de la lecture des valeurs

Les valeurs sont consommées via une somme pondérée. Pour des poids d'attention fixes $p$, l'erreur de lecture est $\Delta o = \sum p_j \Delta v_j$.

• L'article propose une distortion de token pondérée par $A^2$ comme substitut approprié à l'erreur de valeur. Contrairement à la MSE brute des valeurs, cette métrique pondère la distortion de chaque token par le carré de sa probabilité d'attention ($p_j^2$), reflétant le fait que les erreurs sur les tokens fortement attendus contribuent davantage à l'erreur de sortie.

3. Méthodologie : HeadQ

HeadQ est une méthode de quantification côté clé conçue pour préserver le résidu « visible par les scores » plutôt que le vecteur de clé brut.

Aperçu de l'algorithme

1. Calibration : Pour chaque couche et chaque tête, une base de faible rang $U_h$ (rang $r$) est apprise par PCA à partir de requêtes de calibration (post-RoPE).
2. Écriture dans le cache :
   • Un quantificateur standard faible bit $B_b$ génère une clé de base $\hat{k}^{base}_j$.
   • Le résidu $e_j = k_j - \hat{k}^{base}_j$ est calculé.
   • Le résidu est projeté sur la base apprise : $z_j = Q_q(U_h^\top e_j)$, où $Q_q$ est un quantificateur latéral faible bit.
   • Le cache stocke la clé de base et le code latéral $z_j$.
3. Lecture du cache (Inférence) :
   • La requête courante $q$ est utilisée pour déquantifier le code latéral.
   • Une correction additive de logits est appliquée :
     $$\hat{s}_j = \frac{q^\top \hat{k}^{base}_j}{\sqrt{d}} + \frac{q^\top U_h \hat{z}_j}{\sqrt{d}}$$
   • Cette correction cible directement l'erreur de score, évitant la nécessité de reconstruire le vecteur de clé complet.

Distinction clé : HeadQ ne stocke pas la requête courante ; il stocke un code résiduel dans une base apprise par calibration que la requête d'exécution peut décoder en une correction de logits.

4. Contributions clés

1. Formulation de la distortion visible par le modèle : L'article définit la distortion des clés comme une « erreur visible par les scores modulo les décalages constants » et la distortion des valeurs comme une « erreur visible par la lecture avec un substitut diagonal $A^2$ ». Cela explique pourquoi la MSE brute est un mauvais proxy pour la performance comportementale.
2. Algorithme HeadQ : Un algorithme pratique d'écriture/lecture de cache qui stocke un code latéral résiduel de faible rang dans une base de requêtes et l'applique comme une correction additive de logits.
3. Chaîne de preuves falsifiables :
   • Supériorité des métriques : L'erreur dans l'espace de Fisher/des scores prédit la divergence KL de l'attention significativement mieux (gain de Spearman $\rho \approx 0,62\text{--}0,69$) que la MSE brute des clés sur six modèles.
   • Contre-exemples : Des expériences appariées montrent qu'un quantificateur avec une MSE de stockage plus faible peut avoir une KL d'attention et une erreur de sortie plus élevées, prouvant que la MSE de stockage n'est pas un proxy suffisant.
   • Interventions dans l'espace nul : La restauration d'une composante de moyenne « softmax-nulle » (que l'attention dense ignore) augmente nettement le dommage Fisher pour les quantificateurs à grille fixe non équivariants, démontrant que le gaspillage de coordonnées de stockage brise la performance comportementale.
   • Déployabilité : Les bases PCA de requêtes apprises à partir de données de calibration approximent avec succès le résidu actif, alors que des bases aléatoires de même budget échouent.
   • Directionnalité : L'inversion du signe de la correction HeadQ (contrôle « signe erroné ») provoque un échec catastrophique, confirmant que la méthode fonctionne en corrigeant l'objet score spécifique dans la bonne direction.

5. Résultats expérimentaux

Les expériences ont été menées sur WikiText-103 en utilisant six modèles (GPT-2, Pythia, Qwen2, Qwen2.5, TinyLlama, Mistral).

• Expériences K-seul (Valeurs denses) :
  • Sur les lignes 2 bits les plus endommagées, HeadQ (rang 4) élimine environ 84–94 % de la perplexité excédentaire (PPL) par rapport à une base de type KIVI pour des modèles comme Qwen2, Pythia et GPT-2.
  • TinyLlama montre une amélioration matérielle ; Mistral est déjà proche des performances denses dans ce protocole.
• Composition complète KV (2 bits) :
  • La combinaison de HeadQ (côté clé) avec une politique de valeur pondérée par $A^2$ améliore les performances sur les six modèles à une précision de 2 bits.
• Analyse des anomalies :
  • Pythia-160M a montré une corrélation plus faible entre Fisher et KL dans des régimes spécifiques à faible entropie de « basculement de route » (où de petites perturbations inversent le token top-1). Cette anomalie était localisée à une frontière spécifique à une famille et s'est atténuée dans les points de contrôle Pythia plus grands (410M, 1B).

6. Importance et revendications

L'article affirme que l'appariement d'opérateurs est la leçon de conception critique pour la quantification du cache KV :

• Les clés sont des objets de scores ; la compression doit préserver les résidus visibles par la requête et éviter de gaspiller la résolution sur les décalages softmax-nuls.
• Les valeurs sont des objets de lecture ; la compression doit prioriser l'erreur pondérée par $A^2$.

Les auteurs déclarent explicitement que ce travail est mécaniste et diagnostique, et non une revendication d'un système de service prêt pour la production. Ils ne prétendent pas avoir optimisé les noyaux packés, la latence ou le débit. L'importance réside dans :

1. La correction de l'objectif théorique de la quantification du cache KV, passant de la reconstruction de stockage à la distortion visible par le modèle.
2. La fourniture d'une méthode déployable (HeadQ) qui réalise une récupération substantielle de PPL dans des régimes faible bit en ciblant la bonne géométrie d'erreur.
3. La démonstration que les métriques de stockage standard (MSE) peuvent être trompeuses sur le plan comportemental, nécessitant un changement vers des proxies dans l'espace des scores.

L'article conclut que la prochaine étape consiste à étendre cette vision appariée aux opérateurs vers un codec KV complet entièrement optimisé, plutôt que de traiter la MSE de stockage comme la cible par défaut.