Résumé technique : Mémoire des modèles de langage et modèles de mémoire pour le langage

Énoncé du problème

L'article comble une lacune fondamentale dans la compréhension des capacités de « mémoire » des grands modèles de langage (LLM). Bien que les modèles d'apprentissage automatique stockent les informations d'entrée dans des vecteurs d'encodage de couches cachées, le contenu informationnel de ces vecteurs dans les modèles de langage causaux est mal caractérisé. Les auteurs postulent que les modèles de langage causaux actuels, optimisés pour la prédiction du token suivant, sont intrinsèquement inadaptés à la formation de mémoires haute fidélité. En effet, la prédiction du token suivant est une fonction de plusieurs vers un, non inversible ; le modèle n'a besoin de retenir que suffisamment d'informations pour prédire le token suivant immédiat, et non l'ensemble de la séquence d'entrée. Par conséquent, lorsque ces modèles sont utilisés pour compresser des séquences d'entrée en vecteurs de mémoire pour des tâches en aval (telles que la récupération ou des fenêtres de contexte étendues), ils échouent à soutenir l'« accès arbitraire à l'information » requis pour les tâches générales de langage.

Méthodologie

1. Mesure de la rétention d'information

Pour quantifier la quantité d'informations d'entrée retenue dans les vecteurs du modèle, les auteurs proposent un cadre de mesure basé sur l'inversion :

• Architecture : Un encodeur figé (le modèle testé) est couplé à un décodeur entraînable. Le décodeur tente de reconstruire la séquence d'entrée originale à partir du vecteur de la dernière couche cachée de l'encodeur.
• Métriques :
  • Ratio d'entropie ($H_r$) : Mesure la fraction d'informations d'entrée retenue en comparant la perte d'entropie croisée de la reconstruction par rapport aux bornes théoriques (distribution uniforme vs reconstruction parfaite).
  • Précision des tokens : La proportion de tokens d'entrée correctement identifiés par le décodeur.
• Faisabilité de l'inversion : Les auteurs définissent un modèle comme ayant une « bonne mémoire » si une fonction d'inversion peut être pratiquement trouvée (dans un budget de calcul d'environ 96 heures V100) pour reconstruire l'entrée.

2. Architectures de modèles de mémoire

L'article introduit des architectures de modèles de mémoire parallélisables conçues pour découpler la rétention d'information de l'extraction d'information :

• Structure Encodeur-Décodeur : Un encodeur traite des blocs de tokens d'entrée en vecteurs de mémoire de taille fixe. Un décodeur traite ces vecteurs (ainsi que de nouveaux tokens optionnels) pour prédire les tokens suivants.
• Parallélisation : Contrairement aux modèles de mémoire récurrents (par exemple, Bulatov et al., 2022), cette architecture permet à l'encodeur de traiter tous les blocs d'entrée en parallèle, tant pendant l'entraînement que pendant l'inférence.
• Stratégies d'entraînement :
  • Objectifs combinés : Entraînement avec une combinaison linéaire de la perte de prédiction du token suivant et d'une perte de « copie » (où le modèle doit reproduire une séquence à partir de la mémoire).
  • Apprentissage par curriculum : Une approche en deux étapes où les décodeurs apprennent d'abord à traiter des vecteurs riches en informations (sans tokens d'entrée) via une tâche de « copie vide », suivie d'un entraînement sur la prédiction standard du token suivant.
  • Encodeurs figés : Entraîner d'abord un encodeur d'autoencodeur haute fidélité, le figer, puis entraîner le décodeur à utiliser ces vecteurs.

3. Configuration expérimentale

• Ensembles de données : FineWeb-edu, FineMath et divers benchmarks standards (LongBench, IFEval, etc.).
• Modèles : Testés sur des Mixers et des Transformers de tailles variées (de petits modèles personnalisés à Llama 3.1, Qwen et BERT Large).
• Tâches : Prédiction causale du token suivant, régénération d'entrée (autoencodage), tâches de copie et tâches de copie vide (reconstruction de tokens uniquement à partir des vecteurs de mémoire).

Résultats clés

1. Les modèles causaux ont une mauvaise mémoire

• Faible rétention d'information : Les modèles de langage causaux et les modèles de récupération présentent une rétention d'information très faible. Même avec une augmentation de la taille du modèle et de la puissance de calcul, le ratio d'entropie ($H_r$) et la précision des tokens pour reconstruire les entrées à partir des vecteurs causaux restent faibles (par exemple, ~5 % de précision des tokens pour $n_{ctx}=512$).
• Non-inversibilité : Les vecteurs des modèles causaux sont insuffisants pour reconstruire la séquence d'entrée, confirmant que l'objectif de prédiction du token suivant n'encourage pas la formation d'une mémoire complète.

2. Les autoencodeurs atteignent une mémoire quasi parfaite

• Haute fidélité : Les autoencodeurs entraînés spécifiquement pour la régénération d'entrée atteignent une formation de mémoire quasi parfaite, avec des précisions de tokens dépassant 90 % et des ratios d'entropie proches de 1,0.
• Dépendance aux données : Ces mémoires haute fidélité ne sont pas « triviales » (c'est-à-dire qu'elles ne font pas simplement correspondre des tokens à des emplacements vectoriels) ; elles se généralisent à des données hors distribution et échouent sur du bruit aléatoire uniforme, prouvant qu'elles apprennent des représentations dépendantes des données et significatives.

3. Succès des objectifs combinés et de l'apprentissage par curriculum

• Entraînement combiné : Les modèles entraînés avec une combinaison d'objectifs causaux et de copie apprennent avec succès à former et à décoder des mémoires riches en informations tout en maintenant leurs capacités de prédiction causale.
• Apprentissage par curriculum : Le curriculum de « copie vide » (apprendre à décoder des mémoires sans tokens d'entrée en premier) est critique. Il permet aux décodeurs d'utiliser efficacement des vecteurs d'autoencodeur figés et haute fidélité. Sans ce curriculum, les décodeurs échouent souvent à exploiter les vecteurs de mémoire riches, revenant par défaut à la prédiction basée sur les tokens.
• Performance : Les modèles à objectif combiné montrent des améliorations significatives dans les benchmarks liés à l'information (par exemple, LongBench, QA multi-documents) par rapport aux modèles purement causaux, bien que les benchmarks de compréhension générale du langage (par exemple, HellaSwag) montrent des résultats mitigés ou neutres.

4. Efficacité architecturale

• Parallélisation : L'architecture encodeur-décodeur parallèle proposée réduit la complexité computationnelle de l'inférence de $O(n^2)$ à environ $O(n^{4/3})$ (ou $O(n^{2/3})$ pour l'encodeur avec parallélisation complète), offrant des gains substantiels en temps jusqu'au premier token et en utilisation de la mémoire.
• Incompatibilité Encodeur-Décodeur : L'article note que les encodeurs « Mixer » et les décodeurs « Transformer » ont souvent des variétés apprises incompatibles, entraînant une baisse d'efficacité lorsqu'ils sont appariés, ce qui suggère qu'un alignement spécifique à l'architecture est nécessaire.

Importance et revendications

L'article revendique fournir une base théorique et empirique expliquant pourquoi les modèles de mémoire ont historiquement sous-performé par rapport aux Transformers à contexte complet : l'entraînement causal est fondamentalement non inversible et donc inadapté à la formation de mémoires haute fidélité.

Les contributions principales sont :

1. Caractérisation : Démontrer que les vecteurs standards des LLM contiennent relativement peu d'informations d'entrée, indépendamment de l'échelle.
2. Changement de paradigme : Proposer que la formation précise de mémoire nécessite des objectifs au-delà de la prédiction du token suivant, spécifiquement des objectifs combinés ou un apprentissage par curriculum qui exigent explicitement la rétention d'information.
3. Architecture : Introduire un modèle de mémoire encodeur-décodeur parallélisable qui équilibre l'efficacité computationnelle avec la capacité de stocker et de récupérer des informations d'entrée arbitraires.
4. Implications pour la récupération : Les résultats suggèrent que les systèmes de récupération basés sur les vecteurs actuels sont limités par la faible densité d'information des vecteurs eux-mêmes, expliquant pourquoi ils échouent souvent aux tâches de recherche d'informations arbitraires par rapport aux limites théoriques.

Les auteurs concluent que si la prédiction du token suivant seule est insuffisante pour la mémoire, les modèles peuvent être entraînés à former et à décoder des mémoires « riches en informations », permettant une nouvelle classe de modèles de mémoire offrant à la fois une efficacité computationnelle et la capacité d'accès arbitraire à l'information.