Résumé Technique : SEKV – Cache KV Adaptatif à la Résolution avec Mémoire Sémantique Hiérarchique

Énoncé du Problème

Les grands modèles de langage (LLM) sont de plus en plus déployés dans des scénarios nécessitant un raisonnement sur de longs contextes (par exemple, l'analyse de documents, le dialogue multi-tours, la planification agentique). Cependant, le cache Clé-Valeur (KV), qui stocke les états intermédiaires pour accélérer le décodage, croît linéairement avec la longueur de la séquence. Conserver un cache KV complet pour des contextes dépassant 128K jetons devient prohibitif en termes de mémoire GPU, rendant souvent l'inférence à long contexte irréalisable sans compression.

Les stratégies de compression existantes font face à un compromis fondamental entre efficacité et fidélité :

1. Éviction de Jetons : Les méthodes qui rejettent définitivement des jetons basées sur les statistiques d'attention (par exemple, H2O, SnapKV) risquent de fragmenter des preuves cohérentes ou de perdre irrémédiablement des informations qui ne deviennent critiques que plus tard dans le processus de génération.
2. Groupement Sémantique Statique : Les méthodes qui compressent les jetons en blocs sémantiques ou en phrases au moment du prefill (par exemple, ChunkKV, SentenceKV) fixent la résolution de compression très tôt. Elles manquent de la capacité de récupérer des détails fins, au niveau du jeton, d'un segment compressé une fois que ce segment devient pertinent lors du décodage.

Le défi central est de concevoir une organisation de la mémoire qui préserve l'accès aux preuves pertinentes sous des contraintes de mémoire strictes sans perte d'information irréversible, permettant au modèle d'ajuster dynamiquement la résolution du contexte stocké en fonction de la pertinence de la requête.

Méthodologie : SEKV

Les auteurs proposent SEKV (Semantic KV Cache), un cache KV sémantique adaptatif à la résolution qui organise le contexte en une mémoire sémantique hiérarchique s'étendant du GPU vers le CPU. SEKV évite l'éviction irréversible de jetons en stockant le contexte à des résolutions variables et en effectuant un « zoom avant » (zoom-in) dynamique vers le niveau de détail approprié lors du décodage.

1. Segmentation de Segments Guidée par l'Entropie

Au lieu d'un découpage de taille fixe, SEKV segmente le contexte d'entrée en segments sémantiquement cohérents en utilisant la surprisal des jetons comme un signal de frontière à coût nul.

• Mécanisme : Pendant la passe forward de prefill, la surprisal $H_t = -\log p(x_t | x_{<t})$ est calculée pour chaque jeton. Les jetons à haute surprisal, qui marquent souvent des transitions sémantiques (changements de sujet, introduction d'entités), sont identifiés comme des frontières de segments.
• Jetons Ancres : Ces jetons de frontière à haute surprisal sont conservés à pleine résolution sur le GPU en tant que « jetons ancres ». Ils portent une information sémantique concentrée, essentielle pour le routage et la reconstruction ultérieurs.
• Segments (Spans) : Les jetons situés entre deux ancres voisines forment un segment sémantique $S$.

2. Représentation de Segment à Double Résolution et Hiérarchie de Mémoire

Chaque segment sémantique est représenté simultanément à deux résolutions à travers une hiérarchie GPU–CPU :

• GPU (Routage Grossier) : Chaque segment est représenté par un vecteur de résumé léger (une moyenne pondérée par la surprisal et normalisée des vecteurs de projection de clés) résidant sur le GPU. Ce vecteur est utilisé pour un routage efficace et grossier afin de déterminer quels segments sont pertinents pour la requête actuelle.
• CPU (Reconstruction Fine) : La structure complète au niveau du jeton du segment est encodée de manière compacte sur le CPU à l'aide d'une base SVD de rang faible. Un prédicteur léger assigne des portes souples (soft gates) par composante aux valeurs singulières, permettant au système d'adapter le budget de rang $r_S$ pour chaque segment en fonction de sa densité d'information. Cela permet une reconstruction fidèle et à la demande des paires KV au niveau du jeton sans stocker la matrice complète sur le GPU.

3. Mécanisme de Zoom Avant Entraîné

Pendant le décodage, SEKV emploie un mécanisme appris pour sélectionner dynamiquement les segments à étendre :

• Routage : La requête actuelle est projetée dans le même espace de routage que les résumés de segments. Une porte de pertinence sigmoïde calcule la probabilité qu'un segment soit pertinent.
• Expansion : Un seuil appris, par tête et par couche, détermine si un segment doit être étendu. Si sélectionné, le système récupère de manière asynchrone la base SVD de rang faible correspondante depuis le CPU, reconstruit les paires KV au niveau du jeton, et les intègre dans le calcul de l'attention.
• Attention à Résolution Mixte : Le mécanisme d'attention opère sur une mémoire à résolution mixte : les jetons ancres contribuent avec des paires KV complètes, les segments non étendus contribuent avec leurs vecteurs de résumé grossiers (avec une correction de taille), et les segments étendus contribuent avec les paires de jetons reconstruites.
• Entraînement : Le LLM de base reste gelé. Seuls les projections de routage légères, les seuils de zoom avant et le prédicteur de porte de rang sont entraînés en utilisant une combinaison de perte de distillation (alignement avec un enseignant à KV complet), de supervision de zoom, de perte de reconstruction et d'une pénalité de budget.

Contributions Clés

1. Architecture SEKV : Un cache KV sémantique adaptatif à la résolution qui stocke le contexte à plusieurs résolutions et ajuste dynamiquement le détail du cache selon la pertinence de la requête, évitant l'éviction irréversible de jetons.
2. Segmentation Guidée par l'Entropie : Une méthode utilisant la surprisal des jetons pour identifier des segments sémantiques cohérents et préserver des jetons ancres à haute surprisal pour une récupération fiable.
3. Structure à Double Représentation : Une nouvelle structure KV combinant des résumés de segments résidant sur le GPU pour le routage grossier, des bases de rang faible résidant sur le CPU pour la reconstruction approximative au niveau du jeton, et un mécanisme de zoom avant entraîné qui étend sélectivement uniquement les segments pertinents pour la requête.

Résultats Expérimentaux

SEKV a été évalué sur quatre benchmarks (LongBench, RULER, InfiniteBench, Needle-in-a-Haystack) et cinq modèles de base (incluant des variantes de Llama-3, Mistral et Qwen).

• Performance : SEKV surpasse la plus forte base de comparaison de compression sémantique (SentenceKV) de 5,9 % en moyenne sur tous les benchmarks et modèles. Il obtient les meilleurs résultats dans 20 cas sur 20 de comparaisons benchmark-modèle.
• Fidélité de Récupération : Sur le benchmark Needle-in-a-Haystack (NIAH), SEKV montre une stabilité de récupération supérieure à travers différentes longueurs de contexte et profondeurs d'aiguille par rapport aux méthodes d'éviction de jetons et de compression sémantique statique.
• Efficacité Mémoire : À une longueur de contexte de 128K, SEKV réduit la consommation de mémoire GPU de 53,3 % par rapport au cache KV complet tout en maintenant des performances proches de la ligne de base non compressée.
• Efficacité : SEKV atteint une latence et un débit compétitifs, surpassant souvent le cache KV complet en latence grâce à la réduction de la pression sur la bande passante mémoire, tout en surpassant nettement les méthodes d'éviction agressives en précision.
• Efficacité des Paramètres : La méthode ajoute moins de 0,05 % de paramètres entraînables (environ 4,3 millions de paramètres pour Llama-3-8B) tout en gardant le modèle de base gelé.

Signification et Revendications

L'article soutient qu'une inférence efficace à long contexte nécessite de passer d'une compression statique à une récupération de résolution dynamique. Les auteurs affirment que si la compression statique préserve le contexte grossier, elle échoue à récupérer des preuves fines lorsqu'elles deviennent pertinentes tardivement dans le processus de génération. SEKV répond à cela en organisant le contexte en segments sémantiques guidés par l'entropie et en utilisant un système de mémoire hiérarchique qui permet au modèle de faire un « zoom avant » sur les informations pertinentes à la demande.

La signification de ce travail réside dans la démonstration que :

1. L'éviction irréversible n'est pas nécessaire : Une inférence à haut niveau de fidélité sur long contexte peut être obtenue en déchargeant des représentations de rang faible récupérables vers le CPU plutôt qu'en rejetant les jetons.
2. La résolution adaptative est critique : Allouer une résolution fine uniquement là où elle est nécessaire (via le mécanisme de zoom avant) permet aux modèles de mieux équilibrer les contraintes de mémoire et la fidélité de récupération que les politiques statiques.
3. La structure sémantique aide à la compression : L'utilisation de la segmentation guidée par l'entropie crée des unités plus compressibles que les blocs de jetons arbitraires, permettant une approximation de rang faible plus efficace.

Les auteurs reconnaissent des limites, notant que l'efficacité dépend de la qualité de la construction des segments sémantiques (qui peut avoir des difficultés avec un discours très fragmenté) et que l'approche introduit une dépendance à la bande passante CPU-GPU, ce qui pourrait impacter la latence dans des scénarios de requêtes adverses. Cependant, ils concluent que SEKV fournit un cadre robuste pour une inférence LLM à long contexte, hautement fidèle et scalable.