Résumé Technique : xKV - Compression de Cache KV Inter-couche par Extraction de Vecteurs Singuliers Alignés

1. Énoncé du Problème

Les modèles de langage de grande taille (LLM) à contexte long permettent des applications avancées telles que la recherche d'information à grande échelle et le débogage de bases de code, mais se heurtent à un goulot d'étranglement critique : la consommation mémoire des caches Clé-Valeur (KV) lors de l'inférence. À mesure que les longueurs de contexte s'étendent à des centaines de milliers ou des millions de tokens, l'empreinte mémoire du cache KV devient prohibitif, limitant sévèrement les demandes d'inférence concurrentes et le débit.

Les stratégies d'atténuation existantes relèvent généralement de deux catégories :

• Méthodes intra-couche : Des techniques telles que la quantification, l'éviction de tokens et la décomposition de rang faible sur une seule couche compressent le cache KV de chaque couche indépendamment. Bien qu'efficaces dans une certaine mesure, elles ne parviennent pas à exploiter la redondance potentielle entre les couches.
• Méthodes inter-couche : Des approches comme l'Attention Inter-couche (CLA) et YOCO nécessitent des modifications architecturales coûteuses et un pré-entraînement. D'autres, comme MiniCache, reposent sur une fusion a posteriori de couches adjacentes basée sur la similarité cosinus au niveau des tokens. Cependant, une analyse empirique montre que la similarité au niveau des tokens entre les couches est souvent trop faible pour soutenir une compression robuste sans dégradation significative de la précision.

2. Méthodologie : xKV

Les auteurs proposent xKV, une méthode de compression post-entraînement, plug-and-play, qui exploite la redondance structurelle profonde entre les couches sans nécessiter de réentraînement du modèle ni de modifications architecturales.

2.1. Observation Principale : Vecteurs Singuliers Alignés

L'article remet en question l'hypothèse selon laquelle les couches sont distinctes en se basant sur la similarité cosinus au niveau des tokens. En utilisant l'Alignement du Noyau Centré (CKA), les auteurs démontrent que, bien que les embeddings de tokens individuels puissent diverger, les vecteurs singuliers dominants (la base définissant les variations principales dans l'espace des tokens) sont hautement alignés entre les couches adjacentes. Ce score CKA élevé indique que la structure géométrique des sous-espaces du cache KV est partagée, même si les similarités point par point sont faibles.

2.2. Factorisation Inter-couche (CLF)

En exploitant cet alignement, xKV partitionne les $N$ couches du modèle en fenêtres contiguës de taille $W$. Pour chaque fenêtre, elle effectue une Factorisation Inter-couche :

1. Concaténation : Les caches KV de toutes les couches au sein d'une fenêtre sont concaténés horizontalement.
2. SVD : Une Décomposition en Valeurs Singulières (SVD) de rang $r$ est calculée sur le cache concaténé.
3. Décomposition : Le résultat produit une base de tokens partagée ($A_k$) pour la fenêtre et des matrices de reconstruction spécifiques à la couche ($B_\ell$).
4. Stockage : Au lieu de stocker $W$ caches pleine dimension, le système stocke une base partagée et $W$ matrices de reconstruction plus petites, réduisant la complexité mémoire de $O(W \cdot L \cdot d)$ à $O(L \cdot r + W \cdot r \cdot d)$.

2.3. Reconstruction Sélective (SR)

Pour résoudre le goulot d'étranglement computationnel de la reconstruction de tous les tokens lors de la phase de décodage (qui évolue linéairement avec la longueur de séquence $L$), xKV introduit la Reconstruction Sélective (SR).

• Mécanisme : Au lieu d'une reconstruction dense, SR identifie un sous-ensemble épars de tokens ($S$) les plus pertinents pour la requête actuelle en utilisant une sélection de blocs guidée par des points de repère (similaire à ShadowKV).
• Exécution : Seules les lignes correspondant à ces tokens critiques sont reconstruites à l'aide de la base partagée et des matrices spécifiques à la couche.
• Avantage : Cela réduit le coût de reconstruction à une valeur constante par rapport à la longueur de la séquence, permettant des accélérations de bout en bout.

3. Contributions Clés

1. Découverte de la Redondance Inter-couche : L'article révèle que les caches KV partagent des vecteurs singuliers dominants hautement alignés entre les couches, une forme de redondance masquée par une faible similarité au niveau des tokens.
2. Algorithme xKV : Une méthode sans entraînement qui extrait une base partagée via la Factorisation Inter-couche (CLF), permettant un partage efficace d'informations entre $W$ couches ($W > 2$).
3. Reconstruction Sélective : Une optimisation au moment du décodage qui reconstruit uniquement les tokens pertinents pour la requête, éliminant la surcharge computationnelle linéaire de la reconstruction dense.
4. Compatibilité Plug-and-Play : La méthode ne nécessite aucun ajustement fin ni modification architecturale, la rendant applicable aux modèles pré-entraînés existants (par exemple, Llama-3.1, Qwen2.5, DeepSeek-V2).

4. Résultats Expérimentaux

Les auteurs ont évalué xKV sur Llama-3.1-8B-Instruct, Qwen2.5-7B/14B-Instruct et DeepSeek-Coder-V2-Lite en utilisant les benchmarks RULER, LongBench et Multi-turn NIAH.

• Compression vs Précision :
  • xKV atteint jusqu'à 8× de compression du cache KV avec une perte de précision minimale (≤3% en moyenne) par rapport aux bases de référence d'attention complète.
  • Sur Llama-3.1-8B, xKV avec une compression 8× atteint 88,50% de précision moyenne sur RULER, surpassant nettement MiniCache (45,04% à 1,3×) et égalant ou dépassant les méthodes d'éviction de tokens comme SnapKV et PyramidKV.
  • Dans les tâches NIAH multi-tours, xKV maintient des performances stables à travers les tours de conversation, tandis que les méthodes d'éviction de tokens subissent des baisses abruptes après le premier tour.
• Efficacité et Débit :
  • Mémoire : xKV permet des tailles de lots plus importantes sous le même budget mémoire GPU.
  • Latence : La reconstruction dense réduit la latence d'attention jusqu'à 3,6×.
  • Débit : Combinée à la Reconstruction Sélective (xKV-SR), le système atteint jusqu'à 4,23× d'accélération de bout en bout par rapport aux bases de référence d'attention complète sur les GPU NVIDIA A100.
  • Comparaison : xKV-SR dépasse ShadowKV de 30% en débit sous des niveaux de précision similaires tout en maintenant à la fois les Clés et les Valeurs sur le GPU, évitant ainsi les goulots d'étranglement des transferts PCIe.
• Robustesse : La méthode reste efficace sur différentes architectures de modèles, y compris l'Attention par Groupe de Requêtes (GQA) et l'Attention Latente Multi-têtes (MLA) dans DeepSeek-V2.

5. Importance et Revendications

L'article revendique que xKV fournit une approche plug-and-play pratique pour réduire simultanément l'empreinte mémoire et la latence de l'inférence LLM à contexte long. En déplaçant le focus de la similarité au niveau des tokens vers l'alignement des vecteurs singuliers, xKV débloque une redondance inter-couche inexploitée que les méthodes précédentes ont manquée.

Les auteurs soulignent que leur approche :

• Évite le coût élevé du pré-entraînement requis par les changements architecturaux.
• Surmonte les limitations de précision de la fusion basée sur la similarité (MiniCache).
• Résout le goulot d'étranglement computationnel de la reconstruction dense de faible rang via la Reconstruction Sélective.
• Permet une inférence entièrement sur appareil (GPU) pour les caches compressés, contournant les limites de bande passante PCIe qui contraignent les méthodes basées sur le déchargement comme ShadowKV.

Ce travail est présenté comme une étape fondamentale vers une inférence à contexte long efficace, le code étant rendu publiquement disponible pour faciliter son adoption.