Titans : Apprendre à Mémoriser au Moment du Test

Informations Fondamentales

• ID de l'article : 2501.00663
• Titre : Titans: Learning to Memorize at Test Time
• Auteurs : Ali Behrouz, Peilin Zhong, Vahab Mirrokni (Google Research)
• Classification : cs.LG cs.AI cs.CL
• Date de publication : 31 décembre 2024
• Lien de l'article : https://arxiv.org/abs/2501.00663

Résumé

Cet article propose un nouveau module de mémoire à long terme neuronal capable d'apprendre à mémoriser le contexte historique et d'aider le mécanisme d'attention à se concentrer sur le contexte actuel tout en exploitant les informations du passé lointain. Les auteurs argumentent d'un point de vue mémoriel que les mécanismes d'attention, en raison d'un contexte limité mais d'une modélisation précise des dépendances, fonctionnent comme une mémoire à court terme, tandis que la mémoire neuronale, en raison de sa capacité à mémoriser les données, agit comme une mémoire à long terme plus durable. Sur la base de ces deux modules, les auteurs introduisent une nouvelle famille d'architectures appelée Titans et proposent trois variantes pour intégrer efficacement la mémoire dans l'architecture. Les résultats expérimentaux montrent que Titans est plus efficace que les Transformers et les modèles linéaires récurrents modernes sur les tâches de modélisation du langage, de raisonnement de sens commun, de génomique et de séries temporelles, et peut s'étendre efficacement à des fenêtres de contexte dépassant 2M tokens.

Contexte de Recherche et Motivation

Problèmes Fondamentaux

Les architectures de modélisation de séquences existantes font face à un compromis entre efficacité et performance :

1. Transformers : Bien qu'ils modélisent précisément les dépendances, leur complexité de calcul est O(n²), ce qui limite la longueur du contexte
2. Transformers linéaires/RNNs : Bien qu'efficaces, ils compressent les informations dans un état de taille fixe, ce qui entraîne une dégradation des performances sur les longues séquences
3. Absence de systèmes de mémoire : Les architectures existantes manquent d'un système de mémoire multicouche similaire au cerveau humain (mémoire à court terme, mémoire à long terme, métamémoire, etc.)

Motivation de la Recherche

Inspirés par le système de mémoire humain, les auteurs estiment qu'un paradigme d'apprentissage efficace nécessite :

• Des modules distincts mais interconnectés, chacun responsable d'un composant clé du processus d'apprentissage
• La capacité d'apprendre activement à partir des données et de mémoriser les abstractions de l'historique passé
• Un mécanisme permettant un apprentissage et une adaptation continus au moment du test

Contributions Principales

1. Module de mémoire à long terme neuronal : Propose un réseau de neurones profond comme métamodèle qui apprend comment mémoriser/stocker les données dans ses paramètres au moment du test
2. Mécanisme de gestion de la mémoire : Conçoit un mécanisme de mise à jour de la mémoire basé sur la « surprise » et un mécanisme d'oubli adaptatif
3. Famille d'architectures Titans : Propose trois façons d'intégrer la mémoire dans les architectures d'apprentissage profond : Memory as Context (MAC), Memory as Gate (MAG), Memory as Layer (MAL)
4. Algorithme d'entraînement parallélisé : Fournit un algorithme d'entraînement parallélisé rapide rendant l'entraînement des modules de mémoire profonde efficace
5. Validation expérimentale étendue : Valide l'efficacité de Titans sur plusieurs tâches, notamment la modélisation du langage, le raisonnement de sens commun, la génomique et la prédiction de séries temporelles

Détails de la Méthode

Définition de la Tâche

Cet article étudie les tâches de modélisation de séquences, où l'entrée est une séquence $x \in \mathbb{R}^{N \times d_{in}}$, et l'objectif est d'apprendre un modèle capable de traiter efficacement les longues séquences, qui doit :

• Permettre un apprentissage et une mémorisation continus au moment du test
• Équilibrer l'utilisation de la mémoire à court terme et à long terme
• Avoir une complexité linéaire tout en maintenant une grande expressivité

Module de Mémoire à Long Terme Neuronal

Concept de Conception Fondamental

Inspiré par la mémoire à long terme humaine, les événements qui violent les attentes (surprenants) sont plus faciles à retenir. Les auteurs utilisent le gradient du réseau de neurones par rapport à l'entrée pour mesurer le degré de « surprise ».

Mécanisme de Mise à Jour de la Mémoire

Règle de mise à jour de base :

M_t = M_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

Règle de mise à jour améliorée (introduction de la dynamique) :

M_t = M_{t-1} + S_t
S_t = η_t S_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

Où :

• S_t : Dynamique de surprise, contenant la surprise passée et instantanée
• η_t : Paramètre de décroissance de surprise dépendant des données
• θ_t : Paramètre contrôlant le degré d'intégration de la surprise instantanée

Mécanisme d'Oubli

Pour traiter les longues séquences, un mécanisme d'oubli adaptatif est introduit :

M_t = (1 - α_t)M_{t-1} + S_t

Où α_t ∈ [0,1] est un mécanisme de porte contrôlant le degré d'oubli.

Fonction Objectif

Utilise une fonction de perte de mémoire associative :

ℓ(M_{t-1}; x_t) = ||M_{t-1}(k_t) - v_t||²₂

Où k_t = x_t W_K, v_t = x_t W_V

Variantes d'Architecture Titans

1. Memory as Context (MAC)

• Utilise la mémoire comme contexte pour les informations actuelles
• Traite la séquence par blocs, chaque bloc interrogeant la mémoire à long terme pour obtenir les informations historiques pertinentes
• Le mécanisme d'attention décide si les informations de mémoire à long terme sont nécessaires

2. Memory as Gate (MAG)

• Une branche utilise l'attention à fenêtre glissante comme mémoire à court terme
• Une autre branche utilise le module de mémoire neuronale comme mémoire à long terme
• Combine les sorties des deux branches via un mécanisme de porte

3. Memory as Layer (MAL)

• Utilise la mémoire neuronale comme couche d'un réseau profond
• Traitement séquentiel : d'abord la couche de mémoire, puis la couche d'attention
• Conception hiérarchique similaire aux modèles hybrides existants

Entraînement Parallélisé

Réalise un entraînement parallélisé efficace en reformulant le processus d'entraînement pour utiliser des opérations de multiplication matricielle et de sommation :

• Divise la séquence en blocs de taille b
• Utilise un balayage associatif parallèle pour calculer les termes de dynamique
• Réalise un entraînement rapide via la descente de gradient en mini-lots tensorisée

Configuration Expérimentale

Ensembles de Données

• Modélisation du langage : Ensemble de données FineWeb-Edu, 15B/30B tokens
• Raisonnement de sens commun : PIQA, HellaSwag, WinoGrande, ARC-easy/challenge, SIQA, BoolQ
• Tâches de contexte long : Benchmark RULER (S-NIAH), Benchmark BABILong
• Séries temporelles : Ensembles de données ETT, ECL, Traffic, Weather
• Génomique : Ensemble de données GenomicsBenchmarks

Tailles de Modèles

• Modèles de 170M, 340M, 400M, 760M paramètres
• Longueur d'entraînement : 4K tokens
• Fenêtre de contexte : Extensible jusqu'à 2M+ tokens

Méthodes de Comparaison

• Transformers : Transformer++
• Modèles linéaires récurrents : RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT, Gated DeltaNet
• Modèles hybrides : Samba, Gated DeltaNet-H2
• Grands modèles : GPT-4, Llama3, RecurrentGemma, Mistral

Résultats Expérimentaux

Performance de Modélisation du Langage

Sur un modèle de 340M paramètres :

• Titans (LMM) : Perplexité 26.18 (Wiki), 29.97 (LMB)
• Meilleure ligne de base TTT : Perplexité 27.44 (Wiki), 34.19 (LMB)
• Titans (MAG) parmi les modèles hybrides : Perplexité 25.07 (Wiki), 28.72 (LMB)

Tâches de Contexte Long

Sur la tâche S-NIAH (longueur de séquence 16K) :

• Titans (MAC) : S-NIAH-PK 98.4%, S-NIAH-N 97.4%, S-NIAH-W 95.2%
• Mamba2 : S-NIAH-PK 5.4%, S-NIAH-N 0.0%, S-NIAH-W 0.0%
• TTT : S-NIAH-PK 88.4%, S-NIAH-N 4.4%, S-NIAH-W 0.0%

Benchmark BABILong

• Titans dépasse toutes les lignes de base en configuration few-shot, y compris GPT-4
• En configuration fine-tuning, Titans à petite échelle dépasse GPT-4 avec 70 fois plus de paramètres

Prédiction de Séries Temporelles

Le module de mémoire neuronale surpasse les méthodes de base sur tous les ensembles de données, y compris les méthodes basées sur Mamba, Transformer et les modèles linéaires.

Études d'Ablation

Contribution des composants (par ordre d'importance) :

1. Décroissance des poids (mécanisme d'oubli)
2. Mécanisme de dynamique
3. Couche convolutive
4. Mémoire persistante
5. Mémoire profonde vs mémoire linéaire

Travaux Connexes

Modèles Linéaires Récurrents

• Première génération : RetNet, LRU, RWKV, S4/S5 - Utilisant des matrices de transition indépendantes des données
• Deuxième génération : Griffin, Série Mamba - Introduisant des mécanismes de porte
• Troisième génération : DeltaNet, TTT, Longhorn - Basés sur des règles de mise à jour d'apprentissage méta/en ligne

Variantes de Transformer

• Optimisation d'efficacité : Attention clairsemée, attention linéaire, implémentations conscientes des E/S
• Transformers segmentés : RMT et autres utilisant une simple mémoire vectorielle pour transférer les informations entre blocs

Entraînement au Moment du Test

• Inspiré par les premiers algorithmes d'apprentissage local
• Plus étroitement lié à MNM et TTT-layer, mais Titans possède un mécanisme d'oubli et des mises à jour avec dynamique

Conclusion et Discussion

Conclusions Principales

1. Importance des systèmes de mémoire : Un système de mémoire multicouche (mémoire à court terme + mémoire à long terme + mémoire persistante) est crucial pour la modélisation de séquences
2. Efficacité de l'apprentissage au moment du test : L'apprentissage et la mémorisation continus au moment du test peuvent améliorer significativement la capacité de traitement des longues séquences
3. Impact de la conception d'architecture : Les architectures MAC et MAG surpassent la conception traditionnelle de couche MAL
4. Vérification de l'extensibilité : Titans peut s'étendre efficacement à des fenêtres de contexte de 2M+ tokens

Limitations

1. Surcharge de calcul : Le module de mémoire profonde nécessite plus de ressources de calcul qu'un simple état matriciel
2. Compromis de profondeur de mémoire : Les modules de mémoire plus profonds donnent de meilleurs résultats mais s'entraînent plus lentement
3. Sensibilité aux paramètres : Nécessite un réglage minutieux des paramètres liés à la surprise
4. Analyse théorique insuffisante : Manque de garanties théoriques sur la capacité de mémoire et les stratégies d'oubli

Directions Futures

1. Optimisation de l'architecture de mémoire : Explorer des conceptions d'architecture de mémoire neuronale plus efficaces
2. Analyse théorique : Fournir une analyse théorique de la capacité de mémoire et des stratégies d'oubli
3. Validation à grande échelle : Vérifier l'efficacité de la méthode sur des modèles plus grands
4. Extension d'application : Explorer le potentiel d'application dans plus de domaines

Évaluation Approfondie

Points Forts

1. Forte innovativité conceptuelle : Repense la modélisation de séquences du point de vue du système de mémoire humain, proposant une architecture de mémoire multicouche novatrice
2. Contributions techniques complètes : Non seulement propose le module de mémoire neuronale, mais conçoit également trois façons de l'intégrer dans l'architecture et un algorithme de parallélisation efficace
3. Validation expérimentale suffisante : Conduit des expériences complètes sur plusieurs domaines (NLP, séries temporelles, génomique) avec des résultats convaincants
4. Fondation théorique solide : Relie la mise à jour de mémoire à la descente de gradient, à la dynamique et à la décroissance des poids, fournissant une explication théorique

Insuffisances

1. Analyse de complexité de calcul insuffisante : Bien que prétendant à une complexité linéaire, l'analyse des surcharges de calcul réelles du module de mémoire profonde manque de détails
2. Sensibilité aux hyperparamètres : Plusieurs paramètres dépendant des données (α_t, θ_t, η_t) peuvent être complexes à configurer
3. Analogie avec la mémoire humaine limitée : Bien qu'inspirée par la mémoire humaine, la définition de la « surprise » est relativement simple
4. Validation insuffisante sur grands modèles : Le plus grand modèle n'a que 760M paramètres, manquant de validation sur des modèles de milliards de paramètres

Impact

1. Valeur académique : Fournit une nouvelle perspective de mémoire pour la modélisation de séquences, pouvant inspirer davantage de recherches connexes
2. Valeur pratique : Montre d'excellentes performances sur les tâches de traitement de longues séquences, avec un potentiel d'application pratique
3. Reproductibilité : Les auteurs s'engagent à ouvrir le code source, facilitant la promotion et la vérification de la méthode

Scénarios d'Application

1. Traitement de documents longs : Applicable aux tâches NLP nécessitant le traitement de documents longs
2. Analyse de séries temporelles : Particulièrement adaptée aux tâches de prédiction nécessitant des informations historiques à long terme
3. Scénarios d'apprentissage en ligne : Applicable aux applications nécessitant une adaptation continue au moment du test
4. Tâches intensives en mémoire : Comme les systèmes de questions-réponses, les systèmes de dialogue et autres tâches nécessitant de mémoriser de grandes quantités d'informations

Références

L'article cite 138 références connexes, couvrant plusieurs domaines pertinents tels que Transformer, réseaux de neurones récurrents, mécanismes d'attention, réseaux de mémoire et entraînement au moment du test, fournissant une base théorique solide pour cette recherche.