ATLAS: Aprendizaje para Memorizar Óptimamente el Contexto en el Momento de la Prueba

ATLAS: Learning to Optimally Memorize the Context at Test Time

May 29, 2025

Autores: Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, Vahab Mirrokni

cs.AI

Resumen

Los Transformers se han establecido como los modelos más populares en el modelado de secuencias, principalmente debido a su eficacia en tareas de recuperación en contexto y su capacidad para aprender a gran escala. Sin embargo, su complejidad cuadrática en memoria y tiempo limita su aplicabilidad en secuencias más largas, lo que ha motivado a los investigadores a explorar arquitecturas alternativas efectivas, como las redes neuronales recurrentes modernas (también conocidas como módulos de memoria recurrente a largo plazo). A pesar de su reciente éxito en diversas tareas posteriores, estas luchan en tareas que requieren comprensión de contexto largo y extrapolación a secuencias más extensas. Observamos que estas limitaciones provienen de tres aspectos disjuntos en su diseño: (1) capacidad de memoria limitada, restringida por la arquitectura de la memoria y el mapeo de características de la entrada; (2) naturaleza en línea de la actualización, es decir, optimizar la memoria solo con respecto a la última entrada; y (3) gestión menos expresiva de su memoria de tamaño fijo. Para mejorar estos tres aspectos, presentamos ATLAS, un módulo de memoria a largo plazo de alta capacidad que aprende a memorizar el contexto optimizando la memoria en función de los tokens actuales y pasados, superando la naturaleza en línea de los modelos de memoria a largo plazo. Basándonos en esta idea, presentamos una nueva familia de arquitecturas similares a los Transformers, llamada DeepTransformers, que son generalizaciones estrictas de la arquitectura original del Transformer. Nuestros resultados experimentales en tareas de modelado de lenguaje, razonamiento de sentido común, tareas intensivas en recuperación y comprensión de contexto largo muestran que ATLAS supera el rendimiento de los Transformers y los modelos recurrentes lineales recientes. ATLAS además mejora el rendimiento en contexto largo de los Titans, logrando un +80\% de precisión en una longitud de contexto de 10M en el 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.