Un regard approfondi sur les Titans : la voie de la convergence de la mémoire à long terme et de la modélisation efficace des séquences

Titans : Apprendre à mémoriser à l'heure du test Original : https://arxiv.org/pdf/2501.00663v1

Architecture des Titans Mise en œuvre non officielle : https://github.com/lucidrains/titans-pytorch

Un, Contexte et motivation : les limites du transformateur et l'inspiration de la mémoire humaine

1. Limites du transformateur : un goulot d'étranglement dans le traitement des longues séquences

Transformateur Depuis son introduction, le modèle a été en mesure d'apporter une contribution significative au développement du modèle grâce à sa puissance.Mécanisme d'auto-attentionqui a permis des progrès révolutionnaires dans des domaines tels que le traitement du langage naturel et la vision par ordinateur. Cependant, à mesure que la complexité de la tâche augmente, Transformer expose progressivement certains problèmes critiques lorsqu'il s'agit de traiter de longues séquences :

• Une grande complexité de calcul limite l'évolutivité du modèle :.
  • Le mécanisme d'auto-attention nécessite le calcul de chaque jeton la similarité avec tous les autres jetons de la séquence, avec une complexité en temps et en espace de O(N²)qui N est la longueur de la séquence.
  • Cela signifie que le calcul et la consommation de mémoire augmentent par pas carrés lorsque la longueur de la séquence augmente, ce qui limite considérablement la capacité du modèle à gérer de longues séquences. Par exemple, Transformer est souvent dépassé lorsqu'il s'agit de tâches telles que la compréhension de textes longs ou de vidéos, ou encore la prédiction de séries temporelles à long terme.
  

  Figure 1 : Le processus de calcul du mécanisme d'auto-attention.

• Une fenêtre contextuelle limitée pour saisir les dépendances à longue distance :.
  • Pour alléger la charge de calcul, les transformateurs utilisent généralement une fenêtre contextuelle de longueur fixe (par exemple, 512 ou 1024), ce qui signifie que le modèle ne peut se concentrer que sur les informations contenues dans la fenêtre actuelle.
  • Cependant, de nombreuses tâches réelles nécessitent des modèles capables de saisir les dépendances sur des échelles de temps plus longues, comme la compréhension des informations contextuelles dans des textes longs ou des dialogues, l'intégration d'informations provenant de différents moments dans des vidéos et l'élaboration de prédictions à partir de tendances et de modèles à long terme dans des données historiques.

2. Le compromis des transformateurs linéaires : efficacité contre performance

Afin de résoudre le goulot d'étranglement informatique du Transformer, les chercheurs ont proposé l'outilTransformateur linéaireLes principales améliorations sont les suivantes :

• Remplacer softmax par une fonction noyau. Le remplacement du calcul du softmax dans le mécanisme d'auto-attention par une fonction noyau réduit la complexité du calcul à O(N).
• Raisonnement parallélisable. Le processus de calcul du transformateur linéaire peut être représenté sous une forme cyclique, ce qui permet un raisonnement plus efficace.

Cependant, le transformateur linéaire présente également certaines limites :

• Dégradation des performances.
  • L'astuce du noyau dégrade le modèle en un réseau linéaire récurrent où les données sont comprimées en un état à valeur matricielle, ce qui se traduit par des performances inférieures à celles du transformateur standard.
  • Ce type de compression permet difficilement de saisir efficacement les dépendances non linéaires complexes.
• Problèmes de gestion de la mémoire.
  • Le transformateur linéaire compresse les données historiques dans une matrice de taille fixe, mais lorsqu'il s'agit de contextes très longs, cette compression se traduit par une augmentation de la taille de la matrice.dépassement de mémoirece qui affecte la performance du modèle.
  

  Figure 2 : Processus de mise à jour de la mémoire pour le transformateur linéaire.

3. Inspiration du système de mémoire humaine : renforcer les mécanismes de la mémoire à long terme

Afin de surmonter les difficultés susmentionnées, les auteurs ont commencé àsystème de mémoire humaineS'inspirant de la

• La relation entre la mémoire et l'apprentissage : la
  • La thèse s'appuie sur les définitions de la mémoire et de l'apprentissage tirées de la littérature neuropsychologique, considérant la mémoire comme une mise à jour neuronale induite par des données, et définissant l'apprentissage comme le processus d'acquisition de souvenirs efficaces et utiles en fonction d'un objectif.
  • Cela signifie queIl n'y a pas d'apprentissage efficace sans un mécanisme de mémoire solide.
• La nature multiniveaux de la mémoire humaine : la
  • Le système de mémoire humain n'est pas une structure unique mais se compose de multiples sous-systèmes, tels que la mémoire à court terme, la mémoire de travail et la mémoire à long terme, chacun ayant des fonctions et des structures organisationnelles différentes et étant capable de fonctionner de manière indépendante.
  • Cette nature multicouche permet aux humains de stocker, d'extraire et de gérer l'information de manière efficace.
• Les lacunes des modèles existants.
  • Les architectures de réseaux neuronaux existantes (des réseaux Hopfield aux LSTM et aux transformateurs) ont des difficultés à gérer la généralisation, l'extrapolation de la longueur et l'inférence, des capacités qui sont essentielles pour de nombreuses tâches complexes dans le monde réel.
  • Ces architectures, tout en s'inspirant du cerveau humain, n'avaient pas la capacité demémoire éphémèreune modélisation efficace de laNature multi-niveaux du système de mémoireLa simulation.

Deux, Innovation fondamentale : module de mémoire neuronale à long terme et architecture des titans

Sur la base de ces réflexions, les auteurs proposent les innovations suivantes :

1. Module de mémoire neuronale à long terme

(1) Concept de conception.

• Mécanismes d'apprentissage méta-contextuel.
  • Le module est conçu comme unmétamodèleIl s'agit d'apprendre à mémoriser/stocker des données dans ses paramètres au cours des essais.
  • ce type deApprentissage en lignepermet au modèle d'adapter dynamiquement sa mémoire à l'entrée actuelle, plutôt que de s'appuyer sur des mémoires qui ont été fixées lors de la préformation.
• Mises à jour de la mémoire basées sur la surprise : le
  • Les auteurs s'appuient sur le mécanisme de la mémoire humaine selon lequel "les événements surprenants sont plus susceptibles d'être mémorisés", et proposent une méthode basée sur les éléments suivantsdegré de surpriseLe mécanisme de mise à jour de la mémoire.
  • degré de surpriseEn calculant le réseau neuronal par rapport à l'entréegradientPlus le gradient est important, plus les données d'entrée diffèrent des données historiques et plus elles méritent d'être retenues.
    • Cette méthode est efficace pour capturer les informations clés des données et les stocker dans la mémoire à long terme.
    • En revanche, le transformateur linéaire ne peut effectuer qu'une transformation linéaire basée sur les données d'entrée actuelles, ce qui rend difficile la capture efficace des dépendances à long terme.
  

  Figure 3 : Mécanisme de mise à jour de la mémoire basé sur la surprise.

(2) Technologies clés.

• Mécanismes d'impulsion.
  • Pour éviter que le modèle ne soit trop influencé par un seul événement surprenant, les auteurs ont introduit le paramètremécanisme d'impulsionen tenant compte de la surprise des moments passés.
  • Cela signifie que le modèle prend en compte la surprise des entrées actuelles et historiques, ce qui conduit à des mises à jour de la mémoire plus souples.
• Mécanismes d'atténuation.
  • Pour éviter les débordements de mémoire, les auteurs ont également introduit la fonctionmécanisme de désintégrationAu moyen de l'outildiminution du poidsmanière d'oublier progressivement des informations sans importance.
  • Le mécanisme peut être considéré comme unmécanisme de déclenchementqui permet d'effacer sélectivement les mémoires en fonction des besoins.
  • Les auteurs soulignent que ce mécanisme de décroissance est une généralisation du mécanisme d'oubli dans les modèles récurrents modernes et qu'il est équivalent à l'optimisation des réseaux méta-nuraux dans le cadre de la descente de gradient par petits lots, de l'élan et de la décroissance des poids.

(3) Structures de mémoire.

• Contrairement au modèle de mémoire linéaire traditionnel, les auteurs ont utilisé un modèle de mémoire linéaire.Perceptron multicouche (MLP) comme un module de mémoire.
  • Le MLP a une représentation non linéaire plus forte et peut stocker et récupérer des informations complexes plus efficacement.
  • En revanche, les transformateurs linéaires ne peuvent utiliser que des états à valeur matricielle pour stocker des informations, ce qui rend difficile la saisie de relations non linéaires complexes.

2. Architecture des Titans : intégration de la mémoire à long terme et de la mémoire à court terme

Après avoir conçu le module de mémoire neuronale à long terme, les auteurs ont réfléchi à la manière de l'intégrer efficacement dans une architecture d'apprentissage profond et ont proposé l'architecture Titans avec les caractéristiques clés suivantes :

(1) Trois super-têtes travaillant en tandem.

• Cœur de métier.
  • Constituée de la mémoire à court terme, elle est responsable du processus principal de traitement des données.
  • utiliserMécanismes d'attention pour des fenêtres de taille finieL'attention de la fenêtre coulissante (SWA) ou l'attention entièrement connectée (FCA).
  • La mémoire à court terme peut être considérée commemémoire à court termequi est utilisé pour capturer les dépendances dans le contexte actuel.
• Mémoire à long terme.
  • Responsable du stockage et de la mémorisation des informations du passé lointain.
  • Le module de mémoire neuronale à long terme décrit ci-dessus est utilisé.
  • La mémoire à long terme peut être considérée commemémoire éphémèrequi sert à stocker et à récupérer des informations sur une période plus longue.
• Mémoire persistante.
  • est un ensemble de paramètres pouvant être appris, mais indépendants des données, qui encodent des connaissances a priori sur la tâche.
  • Similaire aux paramètres de la couche entièrement connectée dans Transformer, mais avec une fonctionnalité différente.
  • La mémoire persistante peut être considérée commemétamémoirequi sert à stocker les connaissances liées à la tâche, telles que les règles de grammaire, les connaissances de bon sens, etc.
  

  Figure 4 : Schéma de l'architecture des Titans (variante MAC).

(2) Trois types d'intégration différents.

• La mémoire en tant que contexte (MAC).
  • Connecter des mémoires longues et persistantes à des séquences d'entrée en tant qu'informations complémentaires au contexte actuel.
  • Les mécanismes attentionnels déterminent quelles informations doivent être stockées dans la mémoire à long terme.
  • Au moment du test, les paramètres de la mémoire persistante sont restés fixes et le module d'attention a été pondéré pour l'apprentissage contextuel, tandis que le module de mémoire à long terme a continué à apprendre/souvenir des informations.
  • Cette conception permet au modèle d'utiliser de manière flexible les informations de la mémoire à long terme en fonction de l'entrée actuelle.
• La mémoire en tant que porte d'entrée (MAG).
  • Utilisation de l'attention par fenêtre coulissante comme mémoire à court terme et de modules de mémoire neuronale comme mémoire à long terme.
  • Combiner les deux par le biais de mécanismes de contrôle, par exemple en normalisant les deux à l'aide de poids à valeur vectorielle pouvant être appris, puis en appliquant une fonction d'activation non linéaire.
  • Cette conception peut être considérée comme une architecture à têtes multiples dans laquelle les différentes têtes sont structurées différemment.

  
  Figure 5 : Différentes variantes de l'architecture Titans (MAC et MAG).

• La mémoire en couche (MAL).
  • L'utilisation du module de mémoire neuronale comme couche d'un réseau neuronal profond comprime les informations contextuelles passées et actuelles avant le module d'attention.
  • Ce modèle est plus courant dans la littérature, par exemple le modèle H3.

(3) Points forts.

• Gestion plus souple de la mémoire.
  • En utilisant les modules de mémoire comme des branches contextuelles, l'architecture des Titans est capable d'utiliser dynamiquement les informations de la mémoire à long terme en fonction de l'entrée actuelle.
  • Cela offre plus de flexibilité que l'approche traditionnelle consistant à utiliser des modules de mémoire en tant que couches.
• Expression plus forte.
  • La synergie des trois super-têtes permet à l'architecture Titans de traiter plus efficacement les données séquentielles longues et d'intégrer les avantages de la mémoire à court terme, de la mémoire à long terme et de la mémoire persistante.
• Évolutivité.
  • Par rapport à Transformer, l'architecture Titans est plus évolutive lorsqu'elle traite de longues séquences et est capable de maintenir des performances élevées sur une fenêtre contextuelle plus large.

Trois, Résultats expérimentaux et analyse : validation de l'architecture des Titans

Les auteurs ont mené des expériences approfondies sur plusieurs tâches afin d'évaluer les performances de l'architecture Titans et de ses variantes :

1. Modélisation linguistique et raisonnement de bon sens.

• Dispositif expérimental.
  • Trois modèles Titans de tailles différentes (340M, 400M, 760M paramètres) ont été utilisés ainsi que plusieurs modèles de base dont Transformer++, RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT et Gated DeltaNet.
  • L'ensemble de données FineWeb-Edu a été utilisé comme données de formation.
• Principales conclusions.
  • Dans le modèle non mixte, le module de mémoire neuronale à long terme a obtenu les meilleures performances en termes de perplexité et de précision.
  • Les trois variantes de Titans (MAC, MAG, MAL) sont plus performantes que Samba (Mamba + attention) et Gated DeltaNet-H2 (Gated DeltaNet + attention).
  • MAC est plus performant dans le traitement des dépendances à longue distance, tandis que MAG et MAC sont plus performants que la variante MAL.
  

  Figure 6 : Comparaison des performances des Titans avec le modèle de base pour les tâches de modélisation linguistique et de raisonnement par le bon sens.

2. Mission "Aiguille dans la botte de foin".

• Dispositif expérimental.
  • En utilisant la tâche Single NIAH (S-NIAH) dans l'évaluation comparative RULER, la capacité de récupération du modèle a été évaluée sur des séquences de 2K, 4K, 8K et 16K de longueur.
• Principales conclusions.
  • Le module de mémoire neuronale à long terme a obtenu les meilleurs résultats dans les trois tâches.
  • Les variantes Titans ont également obtenu de bons résultats, la variante MAC étant la plus performante.

3. BABILong Benchmarking.

• Dispositif expérimental.
  • Cette tâche exige du modèle qu'il raisonne sur des informations factuelles réparties dans des documents extrêmement longs.
  • Il existe des paramètres de réduction de l'échantillon et des paramètres de réglage fin.
• Principales conclusions.
  • Dans le cadre d'un échantillonnage réduit, Titans surpasse toutes les lignes de base, y compris les modèles comportant un plus grand nombre de paramètres, tels que GPT-4 et GPT4o-mini.
  • Titans surpasse également tous les modèles dans les réglages fins, même les très grands modèles comme le GPT-4.
  • Par rapport aux modèles de mémoire à base de transformateurs (RMT), les Titans affichent de meilleures performances, principalement en raison de leurs fortes capacités de mémoire.
  

  Figure 7 : Comparaison des performances de Titans avec le modèle de base sur le benchmark BABILong.

4. Prévision des séries temporelles.

• Dispositif expérimental.
  • En utilisant le cadre Simba, le module Mamba a été remplacé par un module neuronal de mémoire à long terme.
  • Évalué sur les ensembles de données de référence ETT, ECL, Traffic et Weather.
• Principales conclusions.
  • Le module de mémoire neuronale à long terme surpasse toutes les lignes de base, y compris Mamba, le modèle linéaire et les architectures basées sur les transformateurs.

5. Modélisation de l'ADN.

• Dispositif expérimental.
  • Évaluer les performances des tâches en aval des modèles pré-entraînés sur GenomicsBenchmarks.
• Principales conclusions.
  • Titans (LMM) est compétitif dans différentes tâches de génomique en aval et se situe au même niveau que les méthodes de pointe.

6. Analyse de l'efficacité.

• Principales conclusions.
  • Le module de mémoire neuronale à long terme est légèrement plus lent à former que les autres modèles récurrents, principalement en raison de sa mémoire plus profonde et de son processus de transformation plus complexe, ainsi que du noyau hautement optimisé mis en œuvre dans Mamba2.
  • Titans (MAL) est plus rapide que la ligne de base et que le module de mémoire, principalement en raison de l'utilisation du noyau hautement optimisé de FlashAttention.

7. Études d'ablation.

• Principales conclusions.
  • Tous les composants de la conception de la mémoire neuronale contribuent positivement aux performances, la décroissance du poids, l'élan, la convolution et la mémoire persistante apportant la contribution la plus importante.
  • La conception architecturale a également un impact significatif sur les performances, MAC et MAG obtenant des résultats proches l'un de l'autre dans les tâches de modélisation linguistique et de raisonnement par le bon sens, tandis que MAC obtient de meilleurs résultats dans les tâches à contexte long.

Quatre, Points innovants et points forts de la thèse

1. Un nouveau module neuronal de mémoire à long terme est proposé : le
   • Il s'appuie sur des éléments clés des mécanismes de la mémoire humaine, tels que la surprise, l'élan et l'oubli, pour permettre une mise à jour et un stockage plus efficaces de la mémoire.
   • Un réseau neuronal profond est utilisé comme module de mémoire pour donner au modèle un plus grand pouvoir d'expression.
2. L'architecture des Titans a été conçue pour combiner la mémoire à long terme et la mémoire à court terme.
   • Trois approches d'intégration différentes sont proposées, offrant des options flexibles pour différents scénarios d'application.
   • La synergie des trois super-têtes, le cœur, la mémoire à long terme et la mémoire persistante, permet au modèle de traiter plus efficacement les données séquentielles longues.
3. Exceller dans l'accomplissement de tâches multiples.
   • Qu'il s'agisse de modélisation linguistique, de raisonnement de bon sens, de prédiction de séries temporelles ou de modélisation de l'ADN, l'architecture Titans fait preuve d'une grande performance qui surpasse les modèles existants de transformateurs et de boucles linéaires.
4. Évolutif.
   • La possibilité de maintenir des performances élevées sur une fenêtre contextuelle plus large ouvre la voie au traitement de très longues séquences.

Cinq, perspectives d'avenir

Bien que l'architecture des Titans ait donné des résultats impressionnants sur un certain nombre de fronts, les directions suivantes méritent encore d'être explorées.

1. Explorer des architectures de modules de mémoire plus complexes : la
   • Par exemple, l'introduction de structures de mémoire hiérarchiques ou la combinaison de modules de mémoire avec d'autres modèles tels que les réseaux neuronaux graphiques.
2. Développement de mécanismes plus efficaces pour la mise à jour et le stockage de la mémoire.
   • Par exemple, des techniques de sparification ou de quantification sont utilisées pour réduire la consommation de mémoire et les coûts de calcul.
3. Appliquer l'architecture des Titans à un plus grand nombre de domaines.
   • Les exemples incluent la compréhension des vidéos, le contrôle des robots, les systèmes de recommandation, etc.
4. Explorer des stratégies de formation plus efficaces.
   • Par exemple, l'introduction d'algorithmes d'optimisation plus avancés ou l'utilisation du méta-apprentissage pour accélérer la formation des modèles.
5. Enquête sur l'interprétabilité de l'architecture des Titans.
   • Une meilleure compréhension de la manière dont les Titans stockent et utilisent les informations de la mémoire à long terme pourrait fournir de nouvelles idées pour construire des systèmes d'intelligence artificielle plus puissants.

Six, résumés

La contribution centrale de ce document est la suivante :

• Un nouveau module de mémoire neuronale à long terme est proposéLa conception s'inspire du système de mémoire humaine et intègre des concepts clés de l'apprentissage profond tels que la descente de gradient, l'élan et la décroissance du poids.
• Construire l'architecture des Titansqui combine organiquement la mémoire à long terme et la mémoire à court terme, et explore trois méthodes d'intégration différentes, offrant des options flexibles pour différents scénarios d'application.
• La supériorité des performances de Titans a été vérifiée par des expériences rigoureuses.qui donne de bons résultats sur plusieurs tâches, en particulier lorsqu'il s'agit de tâches à contexte long, et qui démontre une forte évolutivité et une plus grande précision.