Key-Value Means

Le papier présente les Moyennes Clé-Valeur (KVM), un nouveau mécanisme de récurrence par blocs pour l'attention qui unifie les avantages des transformateurs et des RNN linéaires en permettant un entraînement efficace et parallélisable par blocs avec une croissance flexible de l'état et un temps de préremplissage sous-quadratique, le tout en utilisant des opérations standard et un nombre minimal de paramètres supplémentaires.

L'essentiel

Imaginez que vous essayez de lire un livre très long, mais que votre cerveau dispose d'une quantité limitée de « mémoire de travail » pour retenir l'histoire dans votre esprit pendant que vous lisez.

Le problème avec les modèles actuels (Transformers)
Les modèles d'IA actuels (Transformers) agissent comme un élève qui tente de se souvenir de chaque mot individuel qu'il a déjà lu dans le livre.

• Le positif : Ils sont incroyablement précis car ils ont l'histoire entière sous les yeux.
• Le négatif : À mesure que le livre s'allonge, leur « mémoire de travail » devient énorme. Lire un livre de 100 pages demande un effort minime, mais lire un livre de 1 000 pages demande un temps et une énergie considérables. C'est comme essayer de porter un sac à dos qui devient plus lourd à chaque pas que vous faites.

Le problème avec les modèles récurrents (RNN)
Les modèles d'IA récurrents (RNN) adoptent une approche différente : ils maintiennent un petit résumé de taille fixe de ce qu'ils ont lu jusqu'à présent et le mettent à jour au fur et à mesure.

• Le positif : Ils sont super rapides et légers. Leur sac à dos ne devient jamais plus lourd, peu importe la longueur du livre.
• Le négatif : Ils oublient le début de l'histoire. Si vous leur posez une question sur un élément de l'intrigue de la page 10, ils pourraient ne pas s'en souvenir car ils ne retiennent que les dernières pages.

La nouvelle solution : les Moyennes Clé-Valeur (KVM)
Les auteurs de cet article introduisent une nouvelle méthode appelée Moyennes Clé-Valeur (KVM). Considérez la KVM comme un carnet de notes intelligent et magique qui combine le meilleur des deux mondes.

Voici comment cela fonctionne en utilisant une analogie simple :

1. La « Fenêtre Glissante » (Le contexte immédiat)

Imaginez que vous lisez un livre et que vous avez une loupe qui ne vous permet de voir clairement que les dernières pages. C'est la « Fenêtre Glissante ». La KVM porte une attention parfaite aux mots les plus récents, tout comme une IA standard. Cela garantit qu'elle ne manque pas le contexte immédiat.

2. Le « Résumé Compressé » (La mémoire à long terme)

Alors que vous lisez au-delà de ces quelques pages, les anciennes pages glissent hors de votre loupe. Au lieu de les jeter (comme les modèles récurrents) ou d'essayer de porter tout le livre (comme les modèles actuels), la KVM fait quelque chose d'astucieux :

• Elle examine les pages qui viennent de glisser.
• Elle se demande : « Parmi ces pages, lesquelles sont les plus importantes ou les plus uniques ? »
• Elle écrit un résumé court et compressé de ces pages importantes dans un carnet spécial.
• Si une nouvelle page arrive qui est très similaire à ce qui se trouve déjà dans le carnet, elle met simplement à jour la note existante. Si c'est quelque chose de totalement nouveau et surprenant, elle ajoute une nouvelle ligne au carnet.

3. La « Fusion Intelligente » (Le tour de magie)

L'article décrit une méthode spécifique de fusion d'informations appelée règle « Gagnant-Tout-Prendre ».

• Imaginez que vous avez un seau d'eau (les nouvelles informations) et une éponge (le carnet).
• Au lieu de simplement verser l'eau dedans, la KVM trouve l'endroit exact dans l'éponge qui correspond le mieux à l'eau et l'absorbe là.
• Elle utilise également une normalisation « Juste-à-Temps ». Imaginez que vous remplissiez un carnet en écrivant simplement les totaux bruts (les sommes et les comptes) sans faire les divisions compliquées à chaque fois. La KVM retarde cette étape de division : elle ne calcule la moyenne réelle qu'au moment précis où vous ouvrez le carnet pour le lire. Cela évite de devoir recalculer et réajuster tout le carnet à chaque nouvelle note ajoutée.

Pourquoi cela compte

• Taille flexible : Vous pouvez demander à la KVM de garder un carnet de notes minuscule (taille fixe) pour la vitesse, ou laisser le carnet grandir à mesure que le livre s'allonge (taille extensible).
• Vitesse contre Mémoire : Cela vous permet de choisir un juste milieu. Vous n'avez pas à choisir entre « super rapide mais oublieux » ou « super intelligent mais lent ». Vous pouvez l'ajuster pour qu'il soit assez rapide pour une utilisation en temps réel mais assez intelligent pour se souvenir de l'histoire entière.
• Pas de matériel personnalisé : Contrairement à certaines autres nouvelles méthodes qui nécessitent des puces informatiques spéciales et coûteuses pour fonctionner, la KVM peut s'exécuter sur des ordinateurs standards en utilisant des opérations logicielles normales.

Les résultats

Les auteurs ont testé cela sur des modèles de langage (IA qui lit et écrit du texte).

• Contextes courts : Elle a performé aussi bien que les meilleurs modèles d'IA standards.
• Contextes longs : Lorsque l'entrée devenait très longue (milliers de mots), la KVM avec un carnet grandissant se souvenait beaucoup mieux des détails que les modèles à mémoire fixe et était beaucoup plus rapide que les modèles à attention complète.
• L'« Aiguille dans une Botte de Foin » : Dans les tests où l'IA devait trouver une phrase spécifique cachée dans un texte massif, la version de la KVM qui laissait son carnet grandir a très bien performé, prouvant qu'elle pouvait réellement rappeler des informations issues d'un passé lointain.

En bref, la KVM est une nouvelle façon pour l'IA de lire de longs livres sans se fatiguer, sans oublier le début et sans avoir besoin d'un sac à dos qui devient infiniment lourd. Elle y parvient en maintenant une vue claire du présent tout en conservant un résumé intelligent et compressé du passé.

Résumé technique

Résumé Technique : Moyennes Clé-Valeur (KVM)

Énoncé du Problème

Les Transformers offrent un entraînement efficace sur le matériel moderne mais souffrent d'une mise à l'échelle linéaire en mémoire et en temps par token de sortie par rapport à la longueur du contexte (mémoire $O(N)$, temps de décodage $O(N)$). À l'inverse, les RNN linéaires modernes (LRNN) atteignent une mémoire et un temps constants par token ($O(1)$) mais luttent généralement face à une récupération limitée dans des contextes longs. Les architectures existantes tentant de combler ce fossé reposent souvent sur des états de taille fixe (limitant la récupération) ou sur un entraînement complexe au moment du test avec des optimiseurs d'exécution (impactant la vitesse). Il existe un besoin d'une architecture qui équilibre l'efficacité mémoire, la vitesse et la récupération en contexte long sans nécessiter de noyaux personnalisés ni d'ajustement complexe des hyperparamètres pour l'entraînement au moment du test.

Méthodologie : Moyennes Clé-Valeur (KVM)

KVM est un mécanisme d'attention récurrente par blocs novateur qui intègre une attention à fenêtre glissante par blocs (BSWA) avec un état compressible et dynamiquement extensible. Il opère au sein d'une seule couche d'attention softmax, unifiant les avantages des Transformers traditionnels (contexte extensible, parallélisme par morceaux) et des RNN linéaires.

Mécanismes Principaux

1. Fenêtre Glissante par Blocs avec État Compressé :
   KVM traite l'entrée par morceaux. Il maintient une fenêtre BSWA de taille fixe pour les tokens récents et un état séparé, mis à jour périodiquement, pour les tokens plus anciens. Lorsqu'un bloc de tokens déborde de la fenêtre BSWA, il est traité pour mettre à jour l'état plutôt que d'être jeté.

2. Compression et Fusion de l'État :
   Les tokens débordants sont compressés dans l'état en utilisant une règle de fusion de type « gagnant-gagne-tout » basée sur la similarité cosinus.

   • Métrique de Similarité : Au lieu d'un softmax standard, KVM utilise une matrice de mise à jour maximisant la parcimonie (inspirée de la quantification vectorielle en ligne) où chaque clé débordante est assignée à la seule clé d'état la plus corrélée.
   • Renormalisation Juste-à-Temps (JIT) : Pour empêcher la norme des vecteurs d'état de diminuer au fil du temps en raison de la moyenne de vecteurs orthogonaux ou opposés, KVM applique une normalisation JIT. Les clés d'état sont normalisées via LayerNorm avant l'attention. Les valeurs d'état sont normalisées vers un « rayon de lecture » fixe ($\rho_i$) déterminé lors de la création du slot, préservant les magnitudes des valeurs tout en permettant des changements de direction.
   • Porte de Fusion : Une porte scalaire dépendante des données module la quantité de clé/valeur débordante entrante absorbée par l'état.

3. Stratégie d'Extension de l'État :
   Contrairement aux RNN de taille fixe, KVM prend en charge un état extensible. Les tokens débordants les plus « surprenants » (les moins redondants) sont ajoutés directement à l'état, tandis que le reste est fusionné. Cela permet une croissance sous-linéaire de la mémoire (par exemple, $O(\sqrt{N})$) tout en maintenant la récupération du contexte précoce.

4. Gestion du Codage Positionnel :
   Pour maintenir la compatibilité avec les Embarquements Positionnels Rotatifs (RoPE) dans la fenêtre BSWA tout en évitant le RoPE dans l'état compressé (qui agrège des tokens de positions très variées), KVM emploie un remise à zéro partielle du RoPE. Le sous-espace rotatif des clés d'état est mis à zéro, tandis que la fenêtre BSWA conserve le RoPE complet. Cela permet au modèle d'utiliser des requêtes non rotatives pour l'état et des requêtes rotatives pour la fenêtre au sein d'un même passage d'attention.

5. Tokens Puits :
   Un ensemble protégé de lignes d'état initiales (puits) est préservé pour empêcher la dégradation des informations critiques du contexte précoce, résolvant le problème des tokens puits ayant des magnitudes de valeurs distinctes.

Contributions Clés

L'article présente les contributions spécifiques suivantes :

• Formulation Récurrente par Blocs Novatrice : Une méthode pour compresser les tokens débordants dans un état renormalisé dynamiquement en utilisant une règle de fusion « gagnant-gagne-tout », éliminant le besoin de couches de compression séparées.
• Stratégie d'Extension de l'État : Un mécanisme pour ajouter de nouveaux tokens débordants à l'état, permettant une croissance sous-linéaire de la mémoire sans sacrifier la récupération.
• Renormalisation JIT : Un schéma pour normaliser les clés et valeurs d'état juste-à-temps afin de maintenir les normes vectorielles et empêcher les interférences destructrices lors de la moyenne.
• Partage Partiel du RoPE : Une méthode pour partager le codage positionnel entre les régions compressées et non compressées en mettant à zéro la dimension RoPE dans les clés d'état, évitant ainsi le besoin de paramètres supplémentaires ou de mécanismes de re-fusion complexes.
• Architecture Unifiée : Une seule couche d'attention qui interpole entre les RNN à état fixe et les Transformers à attention complète, offrant un compromis continu entre l'efficacité mémoire et la récupération.

Résultats Expérimentaux

Les auteurs ont entraîné des modèles (120M et 350M paramètres) sur le jeu de données Prolong avec des longueurs de contexte de 8k.

• Performance en Contexte Long :
  • KVM à État Fixe (256 tokens) : A surpassé les modèles OVQ/SWA à état plus grand sur la perte de position de séquence et les benchmarks à contexte court. Cependant, il a eu des difficultés dans les tests « Aiguille dans une Botte de Foin » (NIAH) avec des distracteurs nouveaux (NIAH-S2/S3) à des longueurs extrêmes, où la capacité de l'état est devenue un goulot d'étranglement.
  • KVM Extensible (Planifications loi de puissance/Saturantes) : La variante « KVM sqrt » (taille de l'état $\propto \sqrt{N}$) a obtenu des résultats compétitifs sur les benchmarks à contexte long (RULER, LongBench, NIAH), égalant ou surpassant les modèles GPTAlpha non hybrides dans les zones d'extrapolation au-delà du contexte d'entraînement de 8k. Il a nettement surpassé le KVM à état fixe et les LRNN purs (RWKV-7) sur des tâches nécessitant la récupération d'informations nouvelles sur de longues distances.
• Performance en Contexte Court : Les variantes KVM se sont comportées de manière cohérente avec les Transformers standards sur les benchmarks à contexte court (LAMBADA, ARC, HellaSwag, etc.), confirmant que la fenêtre BSWA préserve les capacités d'attention standard.
• Études d'Ablation : La suppression de la normalisation de la longueur des valeurs a causé la dégradation de performance la plus significative. La suppression de la protection des puits et de la porte de fusion a également considérablement affaibli la récupération en contexte long.

Importance et Revendications

L'article revendique que KVM comble avec succès le fossé entre les RNN à état fixe et les Transformers à attention complète.

• Efficacité vs Récupération : Il offre un choix flexible de taille d'état, permettant aux utilisateurs d'ajuster le compromis entre l'efficacité mémoire et la récupération. Avec un état fixe, il offre un comportement récurrent par morceaux $O(N)$ ; avec un état extensible, il atteint une croissance sous-linéaire de la mémoire avec une forte récupération en contexte long.
• Simplicité d'Implémentation : KVM est implémentable en utilisant des opérations standard sans noyaux personnalisés et prend en charge un entraînement et un préremplissage parallélisables par morceaux.
• Potentiel Hybride : L'architecture peut être utilisée dans des solutions hybrides aux côtés de couches LRNN pour les compléter avec une croissance sous-linéaire de la mémoire améliorée et des capacités de décodage en contexte long.
• Absence d'Optimiseurs d'Exécution : Contrairement aux approches d'Entraînement au Moment du Test (TTT), KVM repose sur une règle simple de mise à jour de l'état plutôt que sur des optimiseurs d'exécution comme SGD ou Adam, évitant ainsi les défis d'hyperparamètres associés.

Les auteurs concluent que KVM démontre qu'il est possible d'interpoler de manière fluide entre les RNN à état fixe et l'attention complète de manière simple et efficace, offrant un package unifié pour la modélisation en contexte long.