Résumé Technique : Extension du Contexte des LLM via la Mémoire Récurrente Associative

Énoncé du Problème

Les modèles de langage de grande taille (LLM) sont de plus en plus sollicités pour traiter des entrées s'étendant sur des centaines de milliers ou des millions de jetons (tokens) pour des tâches telles que l'analyse de rapports techniques, le développement logiciel et le raisonnement multi-documents. Cependant, les architectures Transformer standards font face à un goulot d'étranglement fondamental : les coûts computationnels et de mémoire de l'auto-attention augmentent de manière quadratique avec la longueur de la séquence ($O(N^2)$). De plus, les performances se dégradent souvent à mesure que la longueur du contexte augmente, même à l'intérieur de la fenêtre nominale du modèle. Bien que les architectures récurrentes (par exemple, Mamba, RWKV) offrent une mise à l'échelle linéaire, elles nécessitent généralement un entraînement à partir de zéro, ce qui empêche de tirer parti des LLM pré-entraînés existants, et peinent souvent à réaliser des tâches algorithmiques complexes ou du suivi d'instructions par rapport aux Transformers.

Méthodologie

Les auteurs proposent l'Associative Recurrent Memory Transformer (ARMT) comme une solution pratique pour étendre la longueur du contexte tout en maintenant une mise à l'échelle de la mémoire constante et efficace. L'ARMT agit comme une enveloppe (wrapper) autour d'un LLM pré-entraîné, permettant un traitement segmenté des entrées.

Architecture Centrale

L'ARMT divise les entrées à long contexte en segments non chevauchants de longueur fixe. À l'intérieur de chaque segment, le modèle utilise l'auto-attention complète (mémoire de travail à court terme). Crucialement, il introduit un module de mémoire associative par couche qui propage l'information entre les segments (mémoire à long terme). Le mécanisme opère en trois étapes :

1. Extraction de Mémoire : Chaque couche du transformer compresse un segment d'entrée en plongements (embeddings) de mémoire.
2. Consolidation de Mémoire : Ces plongements sont consolidés en une matrice associative par couche sous forme de paires clé-valeur.
3. Association : Les plongements dans les segments suivants sont transformés en vecteurs de requête (query) et multipliés par la matrice associative pour récupérer les informations pertinentes des segments précédents.

Recette d'Entraînement

Le document détaille une stratégie d'entraînement complète pour adapter les LLM pré-entraînés à l'ARMT :

• Pré-entraînement Continu : Les paramètres de la mémoire associative non initialisés sont initialisés via une modélisation de langage non supervisée sur des contextes longs (par exemple, 19 milliards de jetons provenant de FineWeb-Edu) afin d'apprendre une propagation de mémoire efficace avant le réglage fin (fine-tuning) spécifique à la tâche.
• Apprentissage par Curriculum : Pour remédier à la difficulté d'apprendre des dépendances à longue portée à partir de zéro, le modèle est ajusté en augmentant progressivement le nombre de segments (par exemple, de 2 à 4 puis à 8) et en appliquant un lissage (annealing) du taux d'apprentissage.
• Génération de Données Synthétiques : Pour surmonter la rareté des données dans les scénarios de long contexte, les auteurs génèrent des instances d'entraînement synthétiques en concaténant des passages courts de documents longs et en générant des paires QA (questions-réponses) pour chacun, créant ainsi de grands bacs de longueurs de contexte spécifiques.
• Élagage et Sélection de Couches : L'étude examine si la mémoire associative est nécessaire dans chaque couche. Elle propose une stratégie consistant à ne conserver la mémoire associative que dans un sous-ensemble de couches (par exemple, des couches intermédiaires et finales spécifiques), réduisant ainsi les paramètres entraînables et le coût computationnel sans perte de performance significative.

Contributions Clés

1. Jeux de Données Spécifiques au Domaine : La construction de deux nouveaux jeux de données, ManyTypes-long (MT) pour la prédiction de types variables dans le code et GovReport-long (GR) pour le questionnement sur de longs documents, conçus pour évaluer des charges de travail réalistes et à domaine étroit.
2. Recette d'Entraînement : Un cadre novateur pour étendre le contexte des LLM via l'ARMT, combinant pré-entraînement continu, génération de données synthétiques, apprentissage par curriculum et intégration sélective de couches.
3. Validation Empirique : Une étude expérimentale approfondie démontrant que les modèles augmentés par l'ARMT :
   • Traitent des entrées bien au-delà des limites de contexte dées (jusqu'à 64k jetons) sans dégradation des performances par rapport aux bases de référence dans les limites.
   • Font preuve d'une capacité de généralisation supérieure aux longueurs de contexte hors distribution (OOD) par rapport aux modèles de base.
   • Nécessitent environ 30 % de FLOPs en moins tout en préservant les performances de la fenêtre de contexte originale.

Résultats Expérimentaux

Les auteurs ont évalué l'ARMT en utilisant les backbones Gemma-3-1B-IT et SmolLM-2-360M-IT sur les jeux de données MT et GR.

• Performance : Les modèles ARMT ont maintenu des performances stables sur des longueurs de contexte allant jusqu'à 65k jetons. En revanche, les modèles de base (même après fine-tuning) ont montré des chutes de performance brutales au-delà de leurs fenêtres de contexte natives (par exemple, 8k ou 32k). L'ARMT a nettement surpassé les modèles de base dans les régimes Long-OOD (32k–65k).
• Efficacité : L'ARMT a démontré une utilisation constante de la mémoire GPU quel que soit la longueur du contexte, alors que l'utilisation de la mémoire des modèles de base croissait linéairement. Pour une séquence de 32k jetons, l'ARMT a permis ditte augmentation de la taille de lot (batch size) (32 contre 8) pour un même budget mémoire.
• Réduction des FLOPs : L'analyse théorique et les mesures empiriques du temps d'inférence ont confirmé une réduction des FLOPs d'attention globale par un facteur de $T/S$ (longueur de séquence divisée par la taille du segment), conduisant à une réduction globale d'environ 30 % des FLOPs totaux par rapport aux modèles à attention complète.
• Études d'Ablation :
  • Élagage de Couches : Les modèles avec une mémoire associative dans seulement ~20 % des couches (spécifiquement des couches intermédiaires et finales pré-sélectionnées) ont obtenu des performances comparables ou supérieures aux modèles ARMT complets.
  • Pré-entraînement : Le pré-entraînement continu s'est avéré essentiel pour initialiser la mémoire associative, améliorant significativement les performances tant in-domain qu'OOD.
  • Baselines : L'ARMT a surpassé d'autres baselines de long contexte, notamment Mamba-2, DeltaNet et xLSTM, particulièrement en matière de généralisation de long contexte, tout en nécessitant un pré-entraînement moins étendu que les modèles entraînés à partir de zéro.

Signification et Revendications

L'article positionne l'ARMT comme une approche pratique et efficace en termes de calcul pour permettre le traitement de contextes longs dans les LLM de petite à moyenne taille (jusqu'à 1 milliard de paramètres). Les auteurs affirment que cette approche comble le fossé entre les fortes performances en court contexte des Transformers et la mise à l'échelle linéaire des modèles récurrents.

Les principales revendications concernant la signification incluent :

• Confidentialité et Déploiement Local : En permettant un traitement de long contexte efficace sur des modèles plus petits, l'ARMT facilite les applications respectant la vie privée qui ne dépendent pas de LLM distants et de grande échelle via des API.
• Scalabilité : La méthode permet de traiter des contextes arbitrairement longs avec une mémoire constante, une exigence critique pour l'analyse de documents réels.
• Efficacité : La combinaison d'une mise à l'échelle de la mémoire constante et de la réduction des FLOPs fait de l'ARMT une alternative viable pour les environnements aux ressources limitées.
• Généralisation : La recette d'entraînement proposée traite efficacement le phénomène de "perte au milieu" (lost in the middle) et la dégradation des performances observées chez les Transformers standards lors de la manipulation de séquences longues.

Les auteurs reconnaissent des limites, notant que les expériences ont été restreintes à des modèles allant jusqu'à 1 milliard de paramètres et à un ensemble spécifique de tâches (code et QA de documents), et que les mécanismes sous-jacents d'interaction de la mémoire associative avec les représentations du transformer restent partiellement compris. Cependant, ils soutiennent que la capacité démontrée à étendre le contexte dans les petits modèles est une étape significative vers un IA de long contexte pratique et locale.