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: Imparare a Memorizzare al Tempo di Test

Informazioni Fondamentali

  • ID Articolo: 2501.00663
  • Titolo: Titans: Learning to Memorize at Test Time
  • Autori: Ali Behrouz, Peilin Zhong, Vahab Mirrokni (Google Research)
  • Classificazione: cs.LG cs.AI cs.CL
  • Data di Pubblicazione: 31 dicembre 2024
  • Link Articolo: https://arxiv.org/abs/2501.00663

Riassunto

Questo articolo propone un nuovo modulo di memoria neurale a lungo termine in grado di imparare a memorizzare il contesto storico e aiutare il meccanismo di attenzione a concentrarsi sul contesto attuale mentre sfrutta le informazioni del passato remoto. Gli autori sostengono dal punto di vista della memoria che il meccanismo di attenzione, a causa del contesto limitato ma della modellazione accurata delle dipendenze, funziona come memoria a breve termine, mentre la memoria neurale, grazie alla sua capacità di memorizzare dati, agisce come memoria a lungo termine più persistente. Sulla base di questi due moduli, gli autori introducono una nuova famiglia di architetture denominata Titans e propongono tre varianti per integrare efficacemente la memoria nell'architettura. I risultati sperimentali dimostrano che Titans è più efficace dei Transformer e dei moderni modelli ricorrenti lineari nella modellazione del linguaggio, nel ragionamento di senso comune, nella genomica e nei compiti di serie temporali, con la capacità di scalare efficacemente a finestre di contesto superiori a 2M.

Contesto di Ricerca e Motivazione

Problemi Fondamentali

Le architetture di modellazione sequenziale esistenti affrontano un compromesso tra efficienza e prestazioni:

  1. Transformer: Sebbene in grado di modellare accuratamente le dipendenze, hanno una complessità computazionale O(n²), che limita la lunghezza del contesto
  2. Transformer Lineari/RNN: Sebbene efficienti, comprimono le informazioni in uno stato di dimensione fissa, causando un calo di prestazioni su sequenze lunghe
  3. Assenza di Sistemi di Memoria: Le architetture esistenti mancano di un sistema di memoria multilivello simile al cervello umano (memoria a breve termine, memoria a lungo termine, metamemoria, ecc.)

Motivazione della Ricerca

Ispirati dal sistema di memoria umano, gli autori sostengono che un paradigma di apprendimento efficace richiede:

  • Moduli diversi ma interconnessi, ciascuno responsabile di componenti chiave nel processo di apprendimento
  • La capacità di imparare attivamente dai dati e memorizzare astrazioni della storia passata
  • Meccanismi che consentono l'apprendimento e l'adattamento continui al momento del test

Contributi Fondamentali

  1. Modulo di Memoria Neurale a Lungo Termine: Propone una rete neurale profonda come metamodello che apprende come memorizzare/archiviare dati nei suoi parametri al momento del test
  2. Meccanismo di Gestione della Memoria: Progetta un meccanismo di aggiornamento della memoria basato sulla "sorpresa" e un meccanismo di dimenticanza adattivo
  3. Famiglia di Architetture Titans: Propone tre modi per integrare la memoria nelle architetture di apprendimento profondo: Memory as Context (MAC), Memory as Gate (MAG), Memory as Layer (MAL)
  4. Algoritmo di Addestramento Parallelizzato: Fornisce un algoritmo di addestramento parallelizzato veloce che rende efficiente l'addestramento dei moduli di memoria profonda
  5. Verifica Sperimentale Estesa: Verifica l'efficacia di Titans su più compiti, inclusa la modellazione del linguaggio, il ragionamento di senso comune, la genomica e la previsione di serie temporali

Dettagli del Metodo

Definizione del Compito

Questo articolo studia compiti di modellazione sequenziale, dove l'input è una sequenza x∈RN×dinx \in \mathbb{R}^{N \times d_{in}}, e l'obiettivo è imparare un modello in grado di elaborare efficacemente sequenze lunghe, che deve:

  • Continuare ad imparare e memorizzare al momento del test
  • Bilanciare l'uso della memoria a breve e lungo termine
  • Avere complessità lineare mantenendo elevata capacità espressiva

Modulo di Memoria Neurale a Lungo Termine

Idea di Progettazione Fondamentale

Ispirato dalla memoria a lungo termine umana, gli eventi che violano le aspettative (sorprendenti) sono più facili da ricordare. Gli autori utilizzano il gradiente della rete neurale rispetto all'input per misurare il livello di "sorpresa".

Meccanismo di Aggiornamento della Memoria

Regola di aggiornamento di base:

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

Regola di aggiornamento migliorata (con introduzione del momento):

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

Dove:

  • S_t: Momento della sorpresa, contenente sorpresa passata e istantanea
  • η_t: Parametro di decadimento della sorpresa dipendente dai dati
  • θ_t: Parametro che controlla il grado di integrazione della sorpresa istantanea

Meccanismo di Dimenticanza

Per gestire sequenze lunghe, viene introdotto un meccanismo di dimenticanza adattivo:

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

Dove α_t ∈ [0,1] è un meccanismo di gating che controlla il grado di dimenticanza.

Funzione Obiettivo

Utilizza una funzione di perdita di memoria associativa:

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

Dove k_t = x_t W_K, v_t = x_t W_V

Varianti dell'Architettura Titans

1. Memory as Context (MAC)

  • Utilizza la memoria come contesto per le informazioni attuali
  • Elabora la sequenza in blocchi, ogni blocco interroga la memoria a lungo termine per ottenere informazioni storiche rilevanti
  • Il meccanismo di attenzione decide se sono necessarie informazioni dalla memoria a lungo termine

2. Memory as Gate (MAG)

  • Un ramo utilizza l'attenzione a finestra scorrevole come memoria a breve termine
  • L'altro ramo utilizza il modulo di memoria neurale come memoria a lungo termine
  • Combina gli output dei due rami attraverso un meccanismo di gating

3. Memory as Layer (MAL)

  • Utilizza la memoria neurale come strato di una rete profonda
  • Elaborazione sequenziale: prima attraverso lo strato di memoria, poi attraverso lo strato di attenzione
  • Progettazione gerarchica simile ai modelli ibridi esistenti

Addestramento Parallelizzato

Implementa un addestramento parallelizzato efficiente riformulando il processo di addestramento per utilizzare operazioni di moltiplicazione matriciale e somma:

  • Divide la sequenza in blocchi di dimensione b
  • Utilizza scansione associativa parallela per calcolare i termini di momento
  • Implementa discesa del gradiente in mini-batch tensorizzata per addestramento veloce

Configurazione Sperimentale

Dataset

  • Modellazione del Linguaggio: Dataset FineWeb-Edu, 15B/30B token
  • Ragionamento di Senso Comune: PIQA, HellaSwag, WinoGrande, ARC-easy/challenge, SIQA, BoolQ
  • Compiti di Contesto Lungo: Benchmark RULER (S-NIAH), Benchmark BABILong
  • Serie Temporali: Dataset ETT, ECL, Traffic, Weather
  • Genomica: Dataset GenomicsBenchmarks

Dimensioni del Modello

  • Modelli con 170M, 340M, 400M, 760M parametri
  • Lunghezza di addestramento: 4K token
  • Finestra di contesto: scalabile fino a 2M+ token

Metodi di Confronto

  • Transformer: Transformer++
  • Modelli Ricorrenti Lineari: RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT, Gated DeltaNet
  • Modelli Ibridi: Samba, Gated DeltaNet-H2
  • Modelli di Grandi Dimensioni: GPT-4, Llama3, RecurrentGemma, Mistral

Risultati Sperimentali

Prestazioni di Modellazione del Linguaggio

Su modello con 340M parametri:

  • Titans (LMM): Perplessità 26.18 (Wiki), 29.97 (LMB)
  • Baseline Migliore TTT: Perplessità 27.44 (Wiki), 34.19 (LMB)
  • Titans (MAG) tra i Modelli Ibridi: Prestazioni Migliori con Perplessità 25.07 (Wiki), 28.72 (LMB)

Compiti di Contesto Lungo

Su compito S-NIAH (lunghezza sequenza 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 supera tutti i baseline in configurazione few-shot, incluso GPT-4
  • In configurazione fine-tuning, Titans di piccole dimensioni supera GPT-4 con 70 volte più parametri

Previsione di Serie Temporali

Il modulo di memoria neurale supera i metodi baseline su tutti i dataset, inclusi metodi basati su Mamba, Transformer e modelli lineari.

Esperimenti di Ablazione

Contributo dei componenti (ordinato per importanza):

  1. Decadimento dei pesi (meccanismo di dimenticanza)
  2. Meccanismo di momento
  3. Strato convoluzionale
  4. Memoria persistente
  5. Memoria profonda vs memoria lineare

Lavori Correlati

Modelli Ricorrenti Lineari

  • Prima Generazione: RetNet, LRU, RWKV, S4/S5 - Utilizzano matrici di transizione indipendenti dai dati
  • Seconda Generazione: Griffin, Serie Mamba - Introducono meccanismi di gating
  • Terza Generazione: DeltaNet, TTT, Longhorn - Basati su regole di aggiornamento derivate da meta-apprendimento/apprendimento online

Varianti di Transformer

  • Ottimizzazione dell'Efficienza: Attenzione sparsa, attenzione lineare, implementazioni consapevoli dell'I/O
  • Transformer Segmentati: RMT e altri utilizzano memoria vettoriale semplice per trasmettere informazioni tra blocchi

Addestramento al Momento del Test

  • Ispirato da algoritmi di apprendimento locale iniziali
  • Più correlato a MNM e TTT-layer, ma Titans include meccanismo di dimenticanza e aggiornamenti con momento

Conclusioni e Discussione

Conclusioni Principali

  1. Importanza dei Sistemi di Memoria: Un sistema di memoria multilivello (memoria a breve termine + memoria a lungo termine + memoria persistente) è cruciale per la modellazione sequenziale
  2. Efficacia dell'Apprendimento al Momento del Test: L'apprendimento e la memorizzazione continui al momento del test possono migliorare significativamente la capacità di elaborare sequenze lunghe
  3. Impatto della Progettazione dell'Architettura: Le architetture MAC e MAG superano il tradizionale design a strati MAL
  4. Verifica della Scalabilità: Titans scala efficacemente a finestre di contesto di 2M+

Limitazioni

  1. Costi Computazionali: Il modulo di memoria profonda richiede più risorse computazionali rispetto a stati matriciali semplici
  2. Compromesso della Profondità della Memoria: Moduli di memoria più profondi forniscono risultati migliori ma richiedono addestramento più lento
  3. Sensibilità dei Parametri: Richiede un'attenta regolazione dei parametri relativi alla sorpresa
  4. Analisi Teorica Insufficiente: Mancano garanzie teoriche sulla capacità di memoria e strategie di dimenticanza

Direzioni Future

  1. Ottimizzazione dell'Architettura di Memoria: Esplorare progettazioni di architetture di memoria neurale più efficienti
  2. Analisi Teorica: Fornire analisi teorica sulla capacità di memoria e strategie di dimenticanza
  3. Verifica su Larga Scala: Verificare l'efficacia del metodo su modelli di dimensioni maggiori
  4. Estensione dell'Applicazione: Esplorare il potenziale di applicazione in più domini

Valutazione Approfondita

Punti di Forza

  1. Forte Innovazione Concettuale: Ripensa la modellazione sequenziale dal punto di vista dei sistemi di memoria umana, proponendo un'architettura di memoria multilivello innovativa
  2. Contributi Tecnici Completi: Non solo propone il modulo di memoria neurale, ma progetta anche tre modi per integrarlo nell'architettura e un algoritmo di addestramento parallelizzato efficiente
  3. Verifica Sperimentale Completa: Conduce esperimenti completi su più domini (PNL, serie temporali, genomica) con risultati convincenti
  4. Fondamento Teorico Solido: Collega l'aggiornamento della memoria alla discesa del gradiente, al momento e al decadimento dei pesi, fornendo spiegazioni teoriche

Insufficienze

  1. Analisi della Complessità Computazionale Non Sufficientemente Approfondita: Sebbene affermi complessità lineare, l'analisi dei costi computazionali effettivi del modulo di memoria profonda manca di dettagli
  2. Sensibilità agli Iperparametri: Più parametri dipendenti dai dati (α_t, θ_t, η_t) potrebbero richiedere un'impostazione complessa
  3. Analogia Limitata con la Memoria Umana: Sebbene ispirato dalla memoria umana, la definizione di "sorpresa" è relativamente semplice
  4. Verifica Insufficiente su Modelli di Grandi Dimensioni: Il modello più grande ha solo 760M parametri, mancando di verifica su modelli con miliardi di parametri

Valore di Impatto

  1. Valore Accademico: Fornisce una nuova prospettiva di memoria per la modellazione sequenziale, potenzialmente ispirando ricerche correlate
  2. Valore Pratico: Mostra prestazioni eccellenti nei compiti di elaborazione di sequenze lunghe con potenziale applicativo pratico
  3. Riproducibilità: Gli autori si impegnano a rendere open-source il codice, facilitando la diffusione e la verifica del metodo

Scenari Applicabili

  1. Elaborazione di Documenti Lunghi: Applicabile a compiti PNL che richiedono l'elaborazione di documenti lunghi
  2. Analisi di Serie Temporali: Particolarmente adatto a compiti di previsione che richiedono informazioni storiche a lungo termine
  3. Scenari di Apprendimento Online: Applicabile ad applicazioni che richiedono adattamento continuo al momento del test
  4. Compiti Intensivi di Memoria: Come sistemi di domande e risposte, sistemi di dialogo e altri compiti che richiedono la memorizzazione di grandi quantità di informazioni

Bibliografia

L'articolo cita 138 articoli correlati, coprendo aree rilevanti come Transformer, reti neurali ricorrenti, meccanismi di attenzione, reti di memoria e addestramento al momento del test, fornendo una base teorica solida per questa ricerca.