🤖 machine learning

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

Ce document introduit Looped Latent Attention (LLA), un codec de post-entraînement qui exploite la structure de bas rang des caches KV indexés par boucle dans les Transformers à poids partagés afin d'atteindre une compression extrême (jusqu'à 32x) et d'augmenter significativement la capacité de traitement par lot tout en maintenant une performance quasi sans perte grâce à l'initialisation par SVD et à la distillation.

Auteurs originaux : James O' Neill, Fergal Reid

Publié 2026-07-20

📖 7 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : James O' Neill, Fergal Reid

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

Le problème de la mémoire dans les cerveaux d'IA

Imaginez que vous essayiez d'apprendre à un robot à écrire une histoire. Le robot possède un cerveau composé de millions de minuscules interrupteurs (paramètres) qui savent comment les mots s'assemblent. Mais voici le piège : pour écrire une longue histoire, le robot ne peut pas se contenter de regarder son cerveau ; il a besoin d'un bloc-notes. Chaque fois qu'il écrit un nouveau mot, il doit se souvenir de tout ce qu'il vient d'écrire pour ne pas se répéter ou perdre le fil de l'intrigue. Dans le monde de l'intelligence artificielle, ce bloc-notes est appelé le KV cache (cache Clé-Valeur). C'est un carnet de notes privé pour chaque conversation, et il grossit à chaque mot que le robot prononce.

Maintenant, imaginez un type spécial de robot conçu pour être super efficace. Au lieu d'avoir un énorme cerveau avec de nombreuses pièces différentes (couches) pour différentes tâches, ce robot possède une seule petite pièce intelligente qu'il visite encore et encore. Il réutilise le même ensemble d'instructions, bouclant à travers elles pour réfléchir plus profondément. C'est ce qu'on appelle un Transformer en boucle (Looped Transformer). Cela permet d'économiser de l'espace dans le cerveau, ce qui est une bonne chose. Mais il y a un problème sournois : même si le robot réutilise la même pièce, il écrit toujours une nouvelle page dans son bloc-notes à chaque fois qu'il boucle. Si le robot boucle quatre fois pour résoudre un problème de mathématiques, il se retrouve avec quatre pages de notes distinctes pour le même moment de l'histoire. Cela signifie que le robot « en boucle » a toujours besoin d'une quantité massive de mémoire pour fonctionner, ce qui déjoue l'objectif d'être efficace. La grande question pour les scientifiques est la suivante : pouvons-nous rétrécir ce bloc-notes sans faire oublier au robot comment réfléchir ?

La découverte du papier : Un raccourci secret dans les notes

Ce papier présente une astuce ingénieuse appelée Looped Latent Attention (LLA) pour résoudre ce problème de mémoire. Les chercheurs ont découvert que les notes du robot ne sont pas aussi désordonnées qu'elles en ont l'air. Lorsque le robot boucle à travers son processus de réflexion, les notes qu'il écrit pour le même mot ne changent pas radicalement ; elles suivent un chemin fluide et prévisible, comme une bille roulant sur une colline douce. Au lieu de sauvegarder chaque page du bloc-notes du robot, la nouvelle méthode sauvegarde simplement un petit « résumé » de ce chemin et un ensemble d'instructions sur la façon de reconstruire les pages spécifiques uniquement lorsque le robot en a réellement besoin pour les lire.

Pensez-y de cette manière : Imaginez que vous regardez un film où le personnage principal marche dans un couloir. Dans un ordinateur normal, vous enregistreriez une photo haute définition complète du personnage à chaque pas qu'il fait. Cela prend énormément de stockage. Mais avec la LLA, vous réalisez que le personnage marche simplement en ligne droite. Ainsi, au lieu de sauvegarder des milliers de photos, vous sauvegardez une photo du personnage et une note simple disant : « Avance de 1 pouce pour l'étape 2, 2 pouces pour l'étape 3 ». Quand l'ordinateur a besoin de voir le personnage à l'étape 3, il dessine rapidement cette image en utilisant la note. C'est exactement ce que fait ce papier pour l'IA : il compresse la partie « boucle » de la mémoire, transformant une pile de quatre pages en un seul petit résumé qui peut être étendu pour redevenir les pages complètes instantanément.

Les chercheurs ont testé cela sur plusieurs modèles d'IA, dont un appelé Ouro-1.4B. Ils ont constaté que cette méthode fonctionne incroyablement bien. En fait, sur une puce informatique puissante appelée H200, ils ont pu faire tenir 768 conversations différentes dans la mémoire à la fois, alors qu'auparavant, ils ne pouvaient en faire tenir que 32. C'est un bond massif dans la capacité de personnes à utiliser l'IA simultanément. Ils ont également montré que cette astuce fonctionne même lorsque l'IA résout des problèmes mathématiques très longs et difficiles, bien qu'elle ait besoin d'un entraînement supplémentaire pour ne pas s'embrouiller lorsqu'elle écrit sa propre histoire plutôt que de simplement copier un enseignant.

Ce à quoi ce papier dit « Non »

Il est important de savoir ce que cette méthode n'est pas. Les chercheurs ont d'abord testé une idée très simple : « Et si nous jetions simplement les trois premières pages de notes et ne gardions que la toute dernière ? » Ils pensaient que peut-être le robot se stabilise vers une réponse finale après quelques boucles, donc les notes précédentes n'auraient pas d'importance. Ils ont essayé cela, et ce fut un désastre. L'IA a complètement oublié comment faire des mathématiques, obtenant un score de zéro aux tests. Cela prouve que le « chemin » que prennent les notes est important ; on ne peut pas simplement prendre le point final. Il faut capturer le voyage.

Ils ont également comparé leur méthode à d'autres façons de réduire la mémoire, comme l'écrasement des notes (quantification) ou le partage de notes entre différentes parties du cerveau (compression axe-tête). Bien que ces méthodes aident un peu, le papier montre que la compression de l'axe de la « boucle » est la recette secrète. Si vous essayez de compresser les autres parties au lieu des boucles, l'IA perd sa capacité de raisonnement. Le papier est très clair : la boucle est la partie la plus redondante de la mémoire, et c'est le seul endroit où vous devriez compresser.

À quel point sont-ils sûrs ?

Les auteurs sont assez confiants dans ces résultats car ils les ont mesurés directement. Ils n'ont pas seulement deviné ; ils ont fait tourner l'IA sur des tâches réelles comme la résolution de problèmes mathématiques (GSM8K) et l'écriture de code (HumanEval). Ils ont montré qu'avec une compression de 4x (rendant la mémoire 4 fois plus petite), l'IA est presque aussi performante que la version originale non compressée. Lorsqu'ils l'ont poussée à une compression extrême (21,3x), l'IA pouvait encore gérer de nombreuses tâches, bien qu'elle ait commencé à éprouver des difficultés lors des étapes de raisonnement les plus longues et les plus complexes.

Ils ont également testé cela sur un autre type de modèle d'IA appelé Huginn-3.5B, qui possède une structure de cerveau différente et boucle jusqu'à 32 fois. Même là, la méthode a fonctionné, prouvant que ce n'est pas un simple coup de chance pour un robot spécifique. Le papier suggère que cette « trajectoire de rang faible » (le chemin fluide des notes) est une propriété fondamentale de la façon dont ces robots en boucle réfléchissent. Bien qu'ils aient prouvé que cela fonctionne pour économiser de la mémoire et augmenter la capacité, ils notent que pour obtenir les meilleures performances sur des tâches très longues, l'IA a besoin d'un second tour d'entraînement où elle s'exerce avec ses propres notes compressées, plutôt que de simplement copier un enseignant. Cet entraînement « on-policy » corrige les petites erreurs qui surviennent lorsque le robot commence à écrire sa propre histoire.

En résumé, ce papier montre que nous pouvons rendre les modèles d'IA en boucle beaucoup plus efficaces en réalisant que leur mémoire est répétitive d'une manière spécifique. En stockant un petit résumé de cette répétition au lieu des notes complètes, nous pouvons faire tenir des centaines de fois plus de conversations sur une seule puce informatique, rendant l'IA puissante accessible à plus de personnes sans avoir besoin de matériel super coûteux.

Résumé Technique : Looped Latent Attention (LLA)

Problématique

Dans les Transformers à boucles avec partage de poids (weight-tied), un bloc unique est appliqué de manière répétée (récurrence) pour réduire le nombre de paramètres. Bien que cela compresse efficacement les poids du modèle, le goulot d'étranglement de la mémoire lors de l'inférence reste le cache Clé/Valeur (KV). Dans une inférence en boucle standard, le modèle recompute l'attention à chaque étape de récurrence, stockant un cache KV distinct pour chaque étape, couche et tête. Par conséquent, un modèle effectuant $T$ boucles sur $D$ couches transporte un cache KV équivalent à un décodeur de $T \times D$ couches, ce qui annule les bénéfices de la réduction de paramètres par partage de poids.

L'article postule que ce cache indexé par la boucle est hautement structuré. Pour un jeton (token), une couche et une tête donnés, les vecteurs K/V à travers les étapes de récurrence tracent une trajectoire courte et de faible rang plutôt que d'explorer un espace de grande dimension. La redondance est spécifique à l'axe de la boucle ; les spectres le long des axes des têtes et des couches restent significativement plus plats (rang plus élevé). Le défi consiste à compresser cet axe de récurrence sans faire s'effondrer la trajectoire en un point fixe unique, ce qui dégraderait les performances.

Méthodologie : Looped Latent Attention (LLA)

Les auteurs introduisent la Looped Latent Attention (LLA), un codec de cache post-entraînement qui remplace les $T$ vecteurs par tête, couche et jeton par une représentation latente compacte et des cartes de reconstruction spécifiques à la boucle.

Mécanisme Central

Stockage Latent : Au lieu de stocker $T$ $T$ vecteurs ( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ ) pour un jeton, une couche et une tête spécifiques, la LLA stocke un seul vecteur latent de faible rang centré $c \in \mathbb{R}^r$ $c \in R^{r}$ .
- $c = (x - \mu)W_{\text{down}}$
- Ici, $\mu$ est une moyenne calculée sur un ensemble de calibration, et $W_{\text{down}}$ est une matrice de projection descendante.
Reconstruction Spécifique à la Boucle : Le vecteur K/V spécifique pour l'étape de récurrence $t$ $t$ est reconstruit à l'aide d'une projection ascendante spécifique à la boucle :
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
Rangs Asymétriques : La méthode utilise des rangs distincts pour les Clés ( $r_k$ ) et les Valeurs ( $r_v$ ). Une analyse empirique montre que les vecteurs de Valeur se stabilisent plus lentement à travers les boucles que les vecteurs de Clé, ce qui motive l'utilisation de $r_v > r_k$ .
Variante LLA-2D : Pour une compression extrême, la LLA-2D intègre l'axe des têtes dans le latent, stockant un seul latent par jeton et par couche pour reconstruire l'ensemble des $H$ têtes. Cela réduit la taille du cache d'un facteur $H$ mais couple les têtes entre elles.

Entraînement et Initialisation

Initialisation par SVD : Le codec est initialisé via la Décomposition en Valeurs Singulières (SVD) des activations de l'enseignant sur un ensemble de calibration. La projection descendante est définie par les $r$ vecteurs singuliers de droite dominants. Cette initialisation est critique ; une initialisation aléatoire entraîne une divergence KL nettement plus élevée.
Distillation : Le codec est affiné en figeant le modèle enseignant et en minimant la divergence KL directe entre la distribution de l'enseignant et celle du modèle substitué, augmentée d'un terme de correspondance de la sortie d'attention.
Raffinement On-Policy : Pour traiter le biais d'exposition lors des déploiements autoregressifs longs (où le modèle se conditionne sur ses propres préfixes générés plutôt que sur ceux de l'enseignant), une seconde étape de distillation on-policy est employée. Le codec est mis à jour en utilisant des préfixes échantillonnés à partir du modèle étudiant lui-même, stabilisant ainsi les performances de génération longue.

Contributions Clés

Définition de la LLA : Un codec de K/V latent transversal aux boucles pour les Transformers à poids partagés qui atteint un taux de compression $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ . Il stocke $(r_k + r_v)H$ scalaires par jeton par couche au lieu de $2Td_{\text{head}}$ .
Compression Spécifique à l'Axe : L'article démontre que l'axe de récurrence (boucle) est la dimension la plus compressible dans les modèles à boucles (Ouro-1.4B, Ouro-2.6B-Thinking, et Huginn-3.5B). La compression de l'axe de la boucle surpasse la compression de l'axe des têtes (MLA), de l'axe des couches (partage inter-couches), ou l'application de la quantification à des budgets de cache équivalents.
Rejet de l'Effondrement en Point Fixe : Les auteurs montrent que la trajectoire de la boucle ne peut être remplacée par un état unique (la dernière boucle). Un contrôle à zéro paramètre réutilisant uniquement le K/V de la dernière boucle provoque l'effondrement des performances de génération (l'exactitude sur GSM8K chute à 0.000), prouvant que la "trajectoire de faible rang courte" est nécessaire.
Impact sur le Déploiement : La réduction exacte des scalaires se traduit directement par une augmentation de la capacité de service. Sur un seul GPU H200, la LLA augmente la taille de batch maximale pour Ouro-1.4B à 4k de contexte de 32 à 768 séquences (avec une compression de 21,3×).

Résultats Expérimentaux

Compression vs Précision

Performance GSM8K : Sur Ouro-1.4B, une compression légère (1,33× à 2×) égale la précision de l'enseignant (environ 0,79–0,80). À 4× de compression, la LLA maintient une précision de 0,59, surpassant nettement la compression de l'axe des têtes (0,52) et la réutilisation de la dernière boucle (0,00).
Généralisation d'Échelle : Sur Ouro-2.6B-Thinking, la LLA reste quasi-sans perte (0,81–0,83) malgré une compression de 4× à 10,67×, alors que les bases de comparaison sur l'axe des têtes perdent une précision substantielle.
Transfert vers Huginn : La méthode est transférable à Huginn-3.5B (32 étapes de récurrence). Un codec SVD sans entraînement atteint une fidélité quasi parfaite jusqu'à 32× dans les métriques indépendantes du décodeur.

Stabilité de la Génération Longue

MATH-500 : Les longs déploiements mathématiques sont sensibles à la dérive des préfixes. L'entraînement off-policy (utilisant les préfixes de l'enseignant) entraîne une dégradation de la précision et des taux élevés de "non-réponse" à mesure que la longueur de génération augmente.
Gain On-Policy : L'application du raffinement on-policy fait passer l'exactitude de MATH-500 de 0,43 à 0,66 à 4× de compression et réduit significativement les générations sans réponse.

Capacité de Service

Réduction de Mémoire : La réduction du cache est exacte. À 21,3× de compression, le cache KV pour Ouro-1.4B à 1M de contexte passe de 768 GiB à 36 GiB.
Débit (Throughput) : Le "chemin de stockage latent" (stockage des seuls latents) est limité par le calcul mais permet des batchs et des contextes auparavant impossibles dus à l'OOM (Out of Memory). Le "chemin rapide" (reconstruction et mise en cache des K/V complets) préserve le débit mais ne réalise pas l'économie de mémoire.

Signification et Revendications

L'article affirme que la récurrence est l'axe de cache le plus compressible dans les Transformers à boucles avec partage de poids. La signification réside dans :

Insight Mécaniste : Il identifie que l'axe de la boucle est spectralement distinct (rang inférieur) des axes des têtes et des couches, permettant une compression efficace sans faire s'effondrer le comportement dynamique du modèle en un point fixe statique.
Déploiement Pratique : Il fournit une solution post-entraînement qui augmente directement la taille des batchs et les capacités de longueur de contexte sur le matériel existant sans réentraîner l'architecture de base.
Direction Future : Les résultats suggèrent que les architectures récurrentes natives pourraient bénéficier de caches latents découplés, où la récurrence calcule un contenu partagé plus des deltas peu coûteux par boucle, plutôt que de rematérialiser des états K/V complets.

Les auteurs soulignent que leur méthode ne repose pas sur des implémentations d'attention absorbée pour ses affirmations de précision ; le résultat central est la validité de l'hypothèse de la trajectoire de faible rang et l'efficacité du codec à préserver le comportement de l'enseignant sous des contraintes de mémoire strictes.