🤖 machine learning

A JoLT for the KV Cache: Near-Lossless KV Cache Compression via Joint Tucker and JL-Residual Allocation for LLMs

JoLT parvient à une compression de 2 à 3 fois quasi sans perte du cache KV dans les grands modèles de langage en appliant une décomposition de Tucker partielle sur les axes des jetons et des caractéristiques, et en restaurant l'information écartée via un résidu à faible nombre de bits pivoté par Johnson-Lindenstrauss, le tout optimisé sous un budget d'octets unifié pour maintenir les performances de base sur la perplexité et les tâches en aval.

Auteurs originaux : Rahul Krishnan, Volker Schulz

Publié 2026-07-15

📖 7 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Rahul Krishnan, Volker Schulz

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous dirigez un cerveau de robot massif et super intelligent (un grand modèle de langage) qui essaie de raconter une histoire. Pour faire progresser l'histoire, le robot doit se souvenir de chaque mot qu'il a jamais prononcé au cours de la conversation actuelle. Il garde ce souvenir dans un « carnet de notes » spécial appelé KV Cache.

Le problème ? À mesure que l'histoire s'allonge, ce carnet de notes devient énorme. Il devient si gros qu'il commence à dévorer toute la mémoire du robot, ralentissant tout le processus. C'est comme essayer de porter une bibliothèque dans un sac à dos tout en courant un marathon ; finit par s'arrêter.

Des scientifiques ont déjà essayé de rétrécir ce carnet de notes par le passé. Certains ont tenté de l'écraser entièrement dans une petite boîte (quantification), tandis que d'autres ont essayé de résumer les pages en trouvant des motifs dans les lignes ou les colonnes (méthodes de rang faible), mais les auteurs de ce papier, Rahul Krishnan et Volker Schulz, ont remarqué que ces autres méthodes passaient à côté de quelque chose : le carnet de notes n'est pas seulement une pile plate de papier. C'est un bloc 3D avec trois côtés différents : les Têtes (différentes façons de penser), les Tokens (les mots) et les Caractéristiques (les détails).

Ils ont découvert que deux de ces côtés sont remplis de choses ennuyeuses et répétitives qui peuvent être compressées facilement, mais que les autres côtés sont uniques et ne peuvent pas être compressés sans faire perdre au robot son pouvoir cérébral.

La Grande Découverte : La Méthode « JoLT »

L'équipe a inventé un nouveau tour, appelé JoLT (Joint Tucker and JL-residual allocation). Voyez cela comme un service de colisage super intelligent pour votre sac à dos.

L'Écrasement Intelligent (Tucker Partiel) : Au lieu d'essayer d'écraser tout le bloc 3D, JoLT examine les données et dit : « D'accord, les "Têtes" et les "Couches" sont uniques et précieuses ; laissons-les tranquilles. Mais les "Tokens" et les "Caractéristiques" sont pleins de superflu. » Il ne compresse donc que ces deux côtés spécifiques. C'est comme prendre un oreiller géant et moelleux et n'extraire l'air que du milieu, en laissant les bords robustes intacts.
Le Filet de Sécurité (JL-Résiduel) : Quand on écrase un oreiller, un peu d'air s'échappe. Si on le laisse tel quel, l'oreiller est plat et inutile. JoL T attrape cet « air échappé » (l'information perdue) et le stocke dans un filet de sécurité minuscule et super efficace appelé JL-résiduel. Ce filet est si performant pour le compactage qu'il peut contenir les détails manquants en seulement quelques bits.
L'Équilibre Parfait (Dual de Lagrange) : Voici la partie magique. Le robot a un budget strict pour l'espace qu'il peut utiliser (disons, 1 octet). JoLT utilise un « allocateur intelligent » mathématique pour décider exactement de la part de l'oreiller à écraser par rapport à l'espace à accorder au filet de sécurité. Il réalise que certaines parties de la mémoire (les « Clés ») sont faciles à compresser, tandis que d'autres (les « Valeurs ») sont têtues et nécessitent plus d'espace pour le filet de sécurité. Il déplace le budget de manière dynamique pour obtenir le meilleur résultat.

Ce Qu'ils Ont Prouvé (et Ce Qu'Ils N'Ont Pas Prouvé)

Les auteurs ont testé cela sur deux cerveaux de robots célèbres : Mistral-7B (qui utilise un style de « Grouped-Query ») et LLaMA-2-13B (qui utilise un style « Multi-Head »).

La « Zone Libre » : Ils ont trouvé un point idéal où ils pouvaient réduire la mémoire de 2 à 3 fois (2–3×) sans que les performances du robot ne chutent. C'était « quasi-sans perte ».
- Sur des tests comme GSM8K (problèmes mathématiques) et RULER (trouver une aiguille dans une botte de foin textuelle), le robot compressé a obtenu exactement le même score que le robot non compressé, à la marge d'erreur statistique près.
- L'erreur de reconstruction de la mémoire était infime : environ 0,009 pour les clés et 0,006 pour les valeurs. C'est environ 10 fois mieux (un ordre de grandeur) que les méthodes précédentes comme la quantification à 4 bits ou la SVD inter-couches.
Le « Mur » : Ils ont également trouvé une limite. Si vous essayez de trop compresser la mémoire (au-delà de 3×), les choses deviennent chaotiques.
- Le robot Mistral s'est dégradé de manière progressive, perdant lentement un peu de performance à mesure qu'on le compressait.
- Le robot LLaMA, cependant, a percuté un « mur » entre 4× et 5× de compression. Ses performances se sont effondrées, passant d'un score de 5,39 à 9,07 (un énorme bond de la perplexité, ce qui signifie qu'il est devenu bien moins bon pour prédire les mots).

La Version Rapide : FlashJoLT

Calculer l'écrasement parfait prend du temps. Pour corriger cela, ils ont créé FlashJoLT. Au lieu de faire les calculs lourds parfaitement à chaque fois, il utilise un raccourci « randomisé » qui devine les principaux motifs rapidement.

Le Résultat : Il est 5 à 13 fois plus rapide pour compresser la mémoire, mais la qualité reste exactement la même que celle de la version lente et parfaite.

Ce Qu'Ils Ont Éliminé

Le papier est très clair sur ce qui ne fonctionne pas bien pour ce problème spécifique :

Écraser tout : Essayer de compresser les trois côtés (Têtes, Tokens et Caractéristiques) est une mauvaise idée. Les « Têtes » et les « Couches » sont trop uniques ; les écraser nuit au cerveau du robot.
Quantification à bits fixes : Réduire simplement le nombre de bits pour chaque nombre (comme forcer tout à 4 bits) ne peut pas atteindre le « point idéal » de compression de 2–3×. Soit on ne compresse pas assez, soit on perd trop de qualité.
Solution unique pour tous : On ne peut pas traiter les « Clés » et les « Valeurs » de la même manière. Les « Valeurs » sont beaucoup plus difficiles à compresser (2 à 3 fois plus) et nécessitent donc un budget d'espace différent.

En Résumé

Les auteurs ont mesuré cela sur du matériel réel (un GPU A100) et ont trouvé que JoLT crée un moyen quasi-sans perte de réduire la mémoire de 2 à 3 fois sans nuire à l'intelligence du robot.

Cependant, ils précisent bien que ce n'est pas un remède miracle pour tout.

Cela fonctionne très bien pour la « zone libre » (2–3×), mais pousser plus loin sur certains types de robots (comme LLaMA) provoque une chute brutale de la qualité.
Bien que le stockage de la mémoire soit plus petit, le robot doit toujours effectuer des calculs pour « décompresser » la mémoire chaque fois qu'il parle. Ils suggèrent que pour rendre cela vraiment pratique pour une utilisation réelle, les ingénieurs doivent construire des puces spéciales (noyaux fusionnés ou fused kernels) capables de lire directement la mémoire compressée sans avoir à la décompresser au préalable.

En bref, JoLT est un tour de magie mathématique brillant qui économise énormément d'espace pour les longues conversations, mais il a une limite et a besoin d'un petit coup de pouce du futur matériel pour fonctionner à pleine vitesse.

Résumé Technique : Un JOLT pour le KV Cache

Problématique

Dans l'inférence des transformeurs autorégressifs, le cache Clé-Valeur (KV) représente le coût mémoire dominant, évoluant en $O(\text{batch} \times T \times L \times n_h \times d_h)$ . Pour les contextes longs ou les grands lots (batches), le cache KV, plutôt que les poids du modèle, devient le principal goulot d'étranglement de la mémoire et limite le débit. Les méthodes de compression existantes appartiennent à deux familles, toutes deux échouant à exploiter pleinement la redondance structurelle du cache :

Méthodes de bas rang (low-rank) : Elles factorisent des tranches bidimensionnelles (par exemple, des matrices par tête ou des blocs inter-couches) mais s'engagent sur un seul axe de redondance, ignorant les autres.
Méthodes de quantification : Elles réduisent la largeur de bits de chaque entrée mais sont limitées par des taux de bits fixes (par exemple, 4 bits), créant un « plancher de compression » qui empêche d'atteindre des ratios de compression intermédiaires (2–3×) là où la qualité est souvent préservée.

L'observation fondamentale motivant ce travail est que le cache KV à une couche donnée est un tenseur d'ordre trois ( $n_h \times T \times d_h$ ) où les axes transportent des quantités de redondance très différentes. Plus précisément, les axes tête (head) et couche (layer) sont largement incompressibles, tandis que les axes jeton (token) et caractéristique (feature) détiennent la majeure partie de la structure compressible. De plus, les valeurs sont empiriquement 2 à 3 fois plus difficiles à compresser que les clés en raison de leurs propriétés spectrales distinctes.

Méthodologie : JoLT

La méthode proposée, JoLT (Joint Tucker and JL-residual allocation), répond à ces défis en traitant le cache comme un tenseur d'ordre trois et en appliquant une stratégie de compression ciblée.

1. Décomposition de Tucker Partielle

Contra�à les factorisations de tenseurs complètes qui compressent tous les axes, JoLT applique une décomposition de Tucker partielle. Elle tronque uniquement les axes jeton et caractéristique (qui portent la redondance) tout en gardant les axes tête et couche intacts (qui sont incompressibles). Cela est réalisé via une SVD de haut ordre séquentiellement tronquée (ST-HOSVD).

Motivation : Compresser les axes tête ou couche entraîne une erreur significative, tandis que le troncatage des axes jeton et caractéristique capture l'essentiel de la redondance.

2. Résidu Roté par JL

Un squelette de bas rang pur est insuffisif pour une fidélité quasi sans perte, particulièrement pour les valeurs qui possèdent un spectre plat où le troncatage rejette une énergie significative. JoLT restaure cette énergie rejetée en utilisant un résidu de faible bit roté par Johnson-Lindenstrauss (JL) :

Le résidu de troncatage $R = X - \tilde{X}$ est calculé.
Une rotation JL aléatoire est appliquée au résidu.
Le résidu roté est quantifié uniformément à $b \in \{0, 2, 4, 8\}$ bits.
Le décodage implique d'inverser la rotation. Ce résidu agit comme un levier pour convertir un squelette de perte de précision en un squelette quasi sans perte, ce qui est particulièrement précieux pour les valeurs.

3. Allocation Lagrangienne Conjointe

La méthode utilise un unique Lagrangien dual pour allouer conjointement les ressources entre les deux étapes (rangs de Tucker et largeurs de bits du résidu) et entre les différentes cibles (Clés vs Valeurs, Groupes de Couches).

Objectif : Minimiser l'erreur de reconstruction totale sous un budget global d'octets.
Mécanisme : La variable duale permet à l'allocateur de déplacer des bits entre les rangs et les résidus, et entre les clés et les valeurs, exploitant dynamiquement l'asymétrie de compressibilité (par exemple, financer lourdement le résidu des valeurs tout en réduisant les bits du résidu des clés).
Granularité : L'allocation est effectuée par groupe de couches et séparément pour les Clés et les Valeurs.

4. FlashJoLT (Variante Rapide)

Pour traiter le coût computationnel d'une SVD exacte sur de longs contextes, FlashJoLT remplace la SVD exacte du mode jeton par une SVD de bas rang randomisée.

Elle calcule uniquement les $q_{cap}$ directions principales du mode jeton.
Elle emploie une comptabilisation de la masse de la queue (tail-mass accounting), ajoutant la queue spectrale omise comme une valeur singulière de résidu synthétique pour éviter que l'allocateur ne sous-estime l'énergie capturée par un rang donné.
Cette variante permet un accélération de 5 à 13× du temps de compression avec une perte de qualité négligeable dans la « zone gratuite ».

Contributions Clés

Analyse Spectrale : Une étude empirique identifiant que les axes tête et couche sont incompressibles, tandis que les axes jeton et caractéristique détiennent de la redondance, les valeurs étant nettement plus difficiles à compresser que les clés.
Algorithme JoLT : Un compresseur novateur combinant une décomposition de Tucker partielle avec un résidu roté par JL, couplés par un unique dual lagrangien pour l'allocation conjointe des rangs et des bits.
FlashJoLT : Une variante de SVD randomisée qui réduit drastiquement le temps de compression tout en préservant la qualité.
Validation Empirique : Une étude complète sur les modèles GQA (Mistral-7B) et MHA (LLaMA-2-13B) établissant une « zone gratuite » de compression quasi sans perte.

Résultats

La « Zone Gratuite » (Compression 2–3×)

L'article identifie un ratio de compression de 2–3× comme une « zone gratuite » où la qualité du modèle est présalisée dans le bruit statistique de la référence non compressée :

Perplexité : Sur Mistral-7B (GQA) et LLaMA-2-13B (MHA), la dérive de la perplexité est négligeable à 2–3× pour des longueurs de contexte allant jusqu'à 1024.
Erreur de Reconstruction : À 2× de compression, JoLT atteint une erreur de Frobenius relative de 0,009 (Clés) et 0,006 (Valeurs). C'est environ un ordre de grandeur inférieur à la quantification 4-bits et aux bases SVD inter-couches.
Tâches Aval (Downstream) :
- GSM8K : Les modèles compressés (2–3×) restent dans l'intervalle de confiance à 95 % de la référence.
- RULER (Needle-in-a-Haystack) : Mistral maintient une précision de récupération de 100 % à 2–3× ; LLaMA égale son plafond de base.

Divergence d'Architecture

Au-delà de la zone gratuite de 2–3×, les deux architectures divergent :

Mistral (GQA) : Se dégrade avec grâce, perdant environ 4 % de perplexité par étape de ratio entier au-delà de 4×.
LLaMA (MHA) : Se dégrade brutalement entre 4× et 5×, la perplexité augmentant de ~27 % à 4× et de ~68 % à 5×. Les tâches de raisonnement (GSM8K) pour LLaMA s'effondrent de manière significative à des ratios plus élevés.

Efficacité

Mémoire : L'empreinte persistante du KV diminue linéairement avec le ratio de compression (par exemple, de 268,4 Mo à 134,6 Mo à 2× pour Mistral avec un contexte de 2048). La mémoire de processus de pointe reste limitée par les poids et les activations.
Vitesse : FlashJoLT offre une accélération de 5 à 13× du temps de compression par rapport au squelette exact.

Signification et Revendications

L'article affirme que JoLT offre une compression de 2–3× quasi sans perte du cache KV, éliminant efficacement le goulot d'étranglement de la mémoire pour l'inférence à long contexte sans sacrifier la qualité du modèle. En traitant le cache comme un tenseur et en exploitant l'asymétrie spectrale spécifique entre les clés et les valeurs, JoLT surpasse les méthodes existantes de bas rang et de quantification dans le régime de compression intermédiaire.

Les auteurs positionnent la méthode comme un levier direct pour réduire le coût de service des contextes longs. Ils notent deux limites principales :

Limites d'Architecture : La méthode est robuste pour les modèles GQA mais fait face à une dégradation brutale pour les modèles MHA au-delà de 4×, suggérant que la bande 2–3× est la recommandation sûre et indépendante du modèle.
Coût de Déploiement : Bien que le stockage soit réduit, le décodage nécessite actuellement la reconstruction du cache à partir des facteurs à chaque étape. Les auteurs identifient le développement d'un noyau fusionné de reconstruction et d'attention comme la prochaine étape critique pour atteindre la parité d'inférence et réaliser pleinement le potentiel de la méthode.