2025-11-20T08:25:14.880374

Titans: Learning to Memorize at Test Time

Behrouz, Zhong, Mirrokni
Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.
academic

Titans: Learning to Memorize at Test Time

Grundinformationen

  • Paper-ID: 2501.00663
  • Titel: Titans: Learning to Memorize at Test Time
  • Autoren: Ali Behrouz, Peilin Zhong, Vahab Mirrokni (Google Research)
  • Klassifizierung: cs.LG cs.AI cs.CL
  • Veröffentlichungsdatum: 31. Dezember 2024
  • Paper-Link: https://arxiv.org/abs/2501.00663

Zusammenfassung

In diesem Artikel wird ein neues neuronales Langzeitgedächtnismodul vorgestellt, das lernen kann, historische Kontexte zu speichern und dem Aufmerksamkeitsmechanismus hilft, sich auf den aktuellen Kontext zu konzentrieren, während gleichzeitig Informationen aus der fernen Vergangenheit genutzt werden. Die Autoren argumentieren aus gedächtnistheoretischer Perspektive, dass Aufmerksamkeitsmechanismen aufgrund ihres begrenzten Kontexts, aber präzisen Abhängigkeitsmodellierung als Kurzzeitgedächtnis fungieren, während neuronale Gedächtnisse aufgrund ihrer Fähigkeit, Daten zu speichern, als langfristiges und dauerhafteres Gedächtnis fungieren. Basierend auf diesen beiden Modulen führen die Autoren eine neue Architekturfamilie namens Titans ein und präsentieren drei Varianten zur effizienten Integration des Gedächtnisses in die Architektur. Die experimentellen Ergebnisse zeigen, dass Titans bei Sprachmodellierung, Commonsense-Reasoning, Genomik und Zeitreihenaufgaben effektiver ist als Transformers und moderne lineare rekurrente Modelle und effektiv auf Kontextfenstergröße von über 2M skaliert werden kann.

Forschungshintergrund und Motivation

Kernprobleme

Bestehende Sequenzmodellierungsarchitekturen sehen sich mit einem Kompromiss zwischen Effizienz und Leistung konfrontiert:

  1. Transformers: Obwohl sie Abhängigkeiten präzise modellieren können, beträgt die Rechenkomplexität O(n²), was die Kontextlänge begrenzt
  2. Lineare Transformer/RNNs: Obwohl effizient, komprimieren sie Informationen in einen festen Zustand, was zu schlechterer Leistung bei langen Sequenzen führt
  3. Fehlende Gedächtnissysteme: Bestehende Architekturen ermangeln eines mehrstufigen Gedächtnissystems ähnlich dem menschlichen Gehirn (Kurzzeitgedächtnis, Langzeitgedächtnis, Metagedächtnis usw.)

Forschungsmotivation

Inspiriert durch das menschliche Gedächtnissystem argumentieren die Autoren, dass effektive Lernparadigmen Folgendes erfordern:

  • Unterschiedliche, aber miteinander verbundene Module, von denen jedes eine Schlüsselkomponente im Lernprozess übernimmt
  • Die Fähigkeit, aktiv aus Daten zu lernen und abstrakte Darstellungen der Vergangenheit zu speichern
  • Mechanismen, die kontinuierliches Lernen und Anpassung zur Testzeit ermöglichen

Kernbeiträge

  1. Neuronales Langzeitgedächtnismodul: Vorschlag eines tiefen neuronalen Netzes als Metamodell, das zur Testzeit lernt, wie man Daten in seinen Parametern speichert/speichert
  2. Gedächtnismanagemechanismus: Entwurf eines auf "Überraschung" basierenden Gedächtnisaktualisierungsmechanismus und eines adaptiven Vergessungsmechanismus
  3. Titans-Architekturfamilie: Präsentation von drei Methoden zur Integration von Gedächtnis in Deep-Learning-Architekturen: Memory as Context (MAC), Memory as Gate (MAG), Memory as Layer (MAL)
  4. Parallelisierter Trainingsalgorithmus: Bereitstellung eines schnellen parallelisierten Trainingsalgorithmus, der effizientes Training tieferer Gedächtnismodule ermöglicht
  5. Umfassende experimentelle Validierung: Validierung der Effektivität von Titans auf mehreren Aufgaben, einschließlich Sprachmodellierung, Commonsense-Reasoning, Genomik und Zeitreihenvorhersage

Methodische Details

Aufgabendefinition

Dieser Artikel untersucht Sequenzmodellierungsaufgaben, bei denen die Eingabe eine Sequenz x∈RN×dinx \in \mathbb{R}^{N \times d_{in}} ist und das Ziel darin besteht, ein Modell zu lernen, das lange Sequenzen effektiv verarbeiten kann. Das Modell muss:

  • Zur Testzeit kontinuierlich lernen und speichern können
  • Die Nutzung von Kurz- und Langzeitgedächtnis ausbalancieren
  • Lineare Komplexität aufweisen, aber hohe Ausdruckskraft bewahren

Neuronales Langzeitgedächtnismodul

Kerndesignidee

Inspiriert durch menschliches Langzeitgedächtnis werden Ereignisse, die Erwartungen verletzen (überraschend sind), leichter erinnert. Die Autoren verwenden den Gradienten des neuronalen Netzes relativ zur Eingabe, um "Überraschung" zu messen.

Gedächtnisaktualisierungsmechanismus

Grundlegende Aktualisierungsregel:

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

Verbesserte Aktualisierungsregel (mit Momentum):

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

Dabei gilt:

  • S_t: Überraschungs-Momentum, enthält vergangene und momentane Überraschung
  • η_t: Datenabhängiger Überraschungs-Abklingparameter
  • θ_t: Parameter, der den Grad der Einbeziehung momentaner Überraschung steuert

Vergessungsmechanismus

Zur Verarbeitung langer Sequenzen wird ein adaptiver Vergessungsmechanismus eingeführt:

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

wobei α_t ∈ [0,1] ein Gating-Mechanismus ist, der den Grad des Vergessens steuert.

Zielfunktion

Verwendung einer assoziativen Gedächtnisverlustfunktion:

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

wobei k_t = x_t W_K, v_t = x_t W_V

Titans-Architekturvarianten

1. Memory as Context (MAC)

  • Gedächtnis als Kontext für aktuelle Informationen
  • Sequenzverarbeitung in Blöcken, wobei jeder Block das Langzeitgedächtnis abfragt, um relevante historische Informationen zu erhalten
  • Aufmerksamkeitsmechanismus entscheidet, ob Langzeitgedächtnisinformationen erforderlich sind

2. Memory as Gate (MAG)

  • Ein Zweig verwendet gleitende Fensteraufmerksamkeit als Kurzzeitgedächtnis
  • Ein anderer Zweig verwendet das neuronale Gedächtnismodul als Langzeitgedächtnis
  • Kombination der Ausgaben beider Zweige durch einen Gating-Mechanismus

3. Memory as Layer (MAL)

  • Neuronales Gedächtnis als Schicht in einem tiefen Netzwerk
  • Sequenzielle Verarbeitung: zuerst durch die Gedächtnisschicht, dann durch die Aufmerksamkeitsschicht
  • Ähnlich dem Schichtdesign bestehender Hybridmodelle

Parallelisiertes Training

Durch Umformulierung des Trainingsprozesses zur Verwendung von Matrixmultiplikation und Summationsoperationen wurde effizientes parallelisiertes Training erreicht:

  • Sequenzaufteilung in Blöcke der Größe b
  • Verwendung paralleler assoziativer Scans zur Berechnung von Momentum-Termen
  • Implementierung schnellen Trainings durch Tensorisierung von Mini-Batch-Gradientenabstieg

Experimentelle Einrichtung

Datensätze

  • Sprachmodellierung: FineWeb-Edu-Datensatz, 15B/30B Tokens
  • Commonsense-Reasoning: PIQA, HellaSwag, WinoGrande, ARC-easy/challenge, SIQA, BoolQ
  • Langkontextaufgaben: RULER-Benchmark (S-NIAH), BABILong-Benchmark
  • Zeitreihen: ETT, ECL, Traffic, Weather-Datensätze
  • Genomik: GenomicsBenchmarks-Datensatz

Modellgröße

  • Modelle mit 170M, 340M, 400M, 760M Parametern
  • Trainierungslänge: 4K Tokens
  • Kontextfenster: skalierbar auf 2M+ Tokens

Vergleichsmethoden

  • Transformers: Transformer++
  • Lineare rekurrente Modelle: RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT, Gated DeltaNet
  • Hybridmodelle: Samba, Gated DeltaNet-H2
  • Große Modelle: GPT-4, Llama3, RecurrentGemma, Mistral

Experimentelle Ergebnisse

Sprachmodellierungsleistung

Bei 340M-Parameter-Modellen:

  • Titans (LMM): Perplexität 26,18 (Wiki), 29,97 (LMB)
  • Beste Baseline TTT: Perplexität 27,44 (Wiki), 34,19 (LMB)
  • Titans (MAG) unter Hybridmodellen: beste Leistung mit Perplexität 25,07 (Wiki), 28,72 (LMB)

Langkontextaufgaben

Bei S-NIAH-Aufgaben (16K Sequenzlänge):

  • 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%

BABILong-Benchmark

  • Titans übertrifft alle Baselines in Few-Shot-Einstellungen, einschließlich GPT-4
  • In Fine-Tuning-Einstellungen übertrifft kleines Titans GPT-4 mit 70-fach mehr Parametern

Zeitreihenvorhersage

Das neuronale Gedächtnismodul übertrifft Baselines auf allen Datensätzen, einschließlich Mamba-, Transformer- und linearer modellbasierter Methoden.

Ablationsstudien

Beitrag der Komponenten (nach Wichtigkeit geordnet):

  1. Gewichtszerfall (Vergessungsmechanismus)
  2. Momentum-Mechanismus
  3. Faltungsschicht
  4. Persistentes Gedächtnis
  5. Tiefes Gedächtnis vs. lineares Gedächtnis

Verwandte Arbeiten

Lineare rekurrente Modelle

  • Erste Generation: RetNet, LRU, RWKV, S4/S5 - Verwendung datenunabhängiger Übergänge
  • Zweite Generation: Griffin, Mamba-Serie - Einführung von Gating-Mechanismen
  • Dritte Generation: DeltaNet, TTT, Longhorn - Basierend auf Meta-Learning/Online-Learning-Aktualisierungsregeln

Transformer-Varianten

  • Effizienzoptimierung: Sparse Attention, lineare Aufmerksamkeit, I/O-bewusste Implementierung
  • Segmentierte Transformer: RMT und ähnliche verwenden einfache Vektoren zur Informationsübertragung zwischen Blöcken

Training zur Testzeit

  • Inspiriert durch frühe lokale Lernalgorithmen
  • Am meisten verwandt mit MNM und TTT-layer, aber Titans verfügt über Vergessungs- und Momentum-Aktualisierungsmechanismen

Schlussfolgerungen und Diskussion

Hauptschlussfolgerungen

  1. Bedeutung von Gedächtnissystemen: Mehrstufige Gedächtnissysteme (Kurz- + Langzeit- + persistentes Gedächtnis) sind für die Sequenzmodellierung entscheidend
  2. Effektivität des Trainings zur Testzeit: Kontinuierliches Lernen und Speichern zur Testzeit kann die Fähigkeit zur Verarbeitung langer Sequenzen erheblich verbessern
  3. Auswirkungen des Architekturdesigns: MAC- und MAG-Architekturen übertreffen traditionelles MAL-Schicht-Design
  4. Skalierungsvalidierung: Titans kann effektiv auf Kontextfenster von 2M+ skaliert werden

Einschränkungen

  1. Rechnerische Kosten: Tiefe Gedächtnismodule erfordern mehr Rechenressourcen als einfache Matrixzustände
  2. Gedächtnistiefe-Kompromiss: Tiefere Gedächtnismodule funktionieren besser, aber das Training ist langsamer
  3. Parametersensitivität: Erfordert sorgfältige Abstimmung überraschungsbezogener Parameter
  4. Unzureichende theoretische Analyse: Mangel an theoretischen Garantien für Gedächtniskapazität und Vergessensstrategien

Zukünftige Richtungen

  1. Gedächtnisarchitektur-Optimierung: Erforschung effizienterer neuronaler Gedächtnisarchitektur-Designs
  2. Theoretische Analyse: Bereitstellung theoretischer Analyse von Gedächtniskapazität und Vergessensstrategien
  3. Großflächige Validierung: Validierung der Methodeneffektivität auf größeren Modellen
  4. Anwendungserweiterung: Erforschung des Anwendungspotenzials in mehr Bereichen

Tiefgreifende Bewertung

Stärken

  1. Starke konzeptionelle Innovation: Neubewertung der Sequenzmodellierung aus der Perspektive des menschlichen Gedächtnissystems mit neuartiger mehrstufiger Gedächtnisarchitektur
  2. Umfassende technische Beiträge: Nicht nur Vorschlag des neuronalen Gedächtnismoduls, sondern auch Design von drei Integrationsmethoden und effizientem Parallelisierungsalgorithmus
  3. Umfassende experimentelle Validierung: Umfassende Experimente in mehreren Bereichen (NLP, Zeitreihen, Genomik) mit überzeugenden Ergebnissen
  4. Solide theoretische Grundlagen: Verbindung von Gedächtnisaktualisierungen mit Gradientenabstieg, Momentum und Gewichtszerfall mit theoretischen Erklärungen

Mängel

  1. Unzureichende Analyse der Rechenkomplexität: Obwohl lineare Komplexität behauptet wird, ist die Analyse der tatsächlichen Rechenkosten tieferer Gedächtnismodule nicht ausreichend detailliert
  2. Hyperparameter-Sensitivität: Mehrere datenabhängige Parameter (α_t, θ_t, η_t) könnten komplex einzustellen sein
  3. Begrenzte Analogie zum menschlichen Gedächtnis: Obwohl von menschlichem Gedächtnis inspiriert, ist die Definition von "Überraschung" relativ einfach
  4. Unzureichende Validierung bei großen Modellen: Größtes Modell nur 760M Parameter, fehlende Validierung bei Milliarden-Parameter-Modellen

Auswirkungen

  1. Akademischer Wert: Bietet neue Gedächtnisperspektive für Sequenzmodellierung, könnte weitere verwandte Forschung inspirieren
  2. Praktischer Wert: Zeigt hervorragende Leistung bei Langsequenz-Verarbeitungsaufgaben mit praktischem Anwendungspotenzial
  3. Reproduzierbarkeit: Autoren versprechen Open-Source-Code, unterstützt Verbreitung und Validierung der Methode

Anwendungsszenarien

  1. Langdokumentverarbeitung: Geeignet für NLP-Aufgaben, die Langdokumente verarbeiten müssen
  2. Zeitreihenanalyse: Besonders geeignet für Vorhersageaufgaben, die langfristige historische Informationen benötigen
  3. Online-Learning-Szenarien: Geeignet für Anwendungen, die kontinuierliche Anpassung zur Testzeit erfordern
  4. Gedächtnisintensive Aufgaben: Wie Frage-Antwort-Systeme, Dialogsysteme und andere Aufgaben, die große Informationsmengen speichern müssen

Literaturverzeichnis

Das Papier zitiert 138 verwandte Arbeiten, die Transformer, rekurrente neuronale Netze, Aufmerksamkeitsmechanismen, Gedächtnisnetze, Training zur Testzeit und andere relevante Bereiche abdecken und eine solide theoretische Grundlage für diese Forschung bieten.