ATLAS: Apprendimento per Memorizzare Ottimamente il Contesto al Momento del Test

Abstract

I Transformer si sono affermati come le architetture più popolari per la modellazione di sequenze, principalmente grazie alla loro efficacia nei compiti di recupero in contesto e alla capacità di apprendere su larga scala. Tuttavia, la loro complessità quadratica in termini di memoria e tempo limita la loro applicabilità a sequenze più lunghe, motivando così i ricercatori a esplorare architetture alternative efficaci, come le moderne reti neurali ricorrenti (note anche come moduli di memoria ricorrente a lungo termine). Nonostante il loro recente successo in vari compiti downstream, queste architetture faticano in attività che richiedono una comprensione del contesto a lungo termine e un'estrapolazione a sequenze più lunghe. Osserviamo che queste carenze derivano da tre aspetti disgiunti nel loro design: (1) una capacità di memoria limitata, vincolata dall'architettura della memoria e dalla mappatura delle caratteristiche dell'input; (2) la natura online dell'aggiornamento, ovvero l'ottimizzazione della memoria solo rispetto all'ultimo input; e (3) una gestione meno espressiva della loro memoria a dimensione fissa. Per migliorare tutti e tre questi aspetti, presentiamo ATLAS, un modulo di memoria a lungo termine ad alta capacità che impara a memorizzare il contesto ottimizzando la memoria in base ai token correnti e passati, superando la natura online dei modelli di memoria a lungo termine. Basandoci su questa intuizione, presentiamo una nuova famiglia di architetture simili ai Transformer, chiamate DeepTransformers, che sono generalizzazioni rigorose dell'architettura Transformer originale. I nostri risultati sperimentali su compiti di modellazione del linguaggio, ragionamento di senso comune, richiamo intensivo e comprensione del contesto a lungo termine mostrano che ATLAS supera le prestazioni dei Transformer e dei recenti modelli ricorrenti lineari. ATLAS migliora ulteriormente le prestazioni nel contesto lungo dei Titans, raggiungendo un'accuratezza del +80\% su una lunghezza di contesto di 10M nel benchmark BABILong.

English

Transformers have been established as the most popular backbones in sequence modeling, mainly due to their effectiveness in in-context retrieval tasks and the ability to learn at scale. Their quadratic memory and time complexity, however, bound their applicability in longer sequences and so has motivated researchers to explore effective alternative architectures such as modern recurrent neural networks (a.k.a long-term recurrent memory module). Despite their recent success in diverse downstream tasks, they struggle in tasks that requires long context understanding and extrapolation to longer sequences. We observe that these shortcomings come from three disjoint aspects in their design: (1) limited memory capacity that is bounded by the architecture of memory and feature mapping of the input; (2) online nature of update, i.e., optimizing the memory only with respect to the last input; and (3) less expressive management of their fixed-size memory. To enhance all these three aspects, we present ATLAS, a long-term memory module with high capacity that learns to memorize the context by optimizing the memory based on the current and past tokens, overcoming the online nature of long-term memory models. Building on this insight, we present a new family of Transformer-like architectures, called DeepTransformers, that are strict generalizations of the original Transformer architecture. Our experimental results on language modeling, common-sense reasoning, recall-intensive, and long-context understanding tasks show that ATLAS surpasses the performance of Transformers and recent linear recurrent models. ATLAS further improves the long context performance of Titans, achieving +80\% accuracy in 10M context length of BABILong benchmark.

ATLAS: Apprendimento per Memorizzare Ottimamente il Contesto al Momento del Test

ATLAS: Learning to Optimally Memorize the Context at Test Time

Abstract

Support