ATLAS: Aprendendo a Memorizar de Forma Ótima o Contexto no Momento do Teste

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

Resumo

Os Transformers foram estabelecidos como as arquiteturas mais populares para modelagem de sequências, principalmente devido à sua eficácia em tarefas de recuperação em contexto e à capacidade de aprendizado em larga escala. No entanto, sua complexidade quadrática em memória e tempo limita sua aplicabilidade em sequências mais longas, o que motivou pesquisadores a explorar arquiteturas alternativas eficazes, como redes neurais recorrentes modernas (também conhecidas como módulos de memória recorrente de longo prazo). Apesar de seu sucesso recente em diversas tarefas subsequentes, elas enfrentam dificuldades em tarefas que exigem compreensão de contexto longo e extrapolação para sequências mais extensas. Observamos que essas limitações decorrem de três aspectos distintos em seu design: (1) capacidade de memória limitada, restrita pela arquitetura da memória e mapeamento de características da entrada; (2) natureza online da atualização, ou seja, otimização da memória apenas em relação à última entrada; e (3) gerenciamento menos expressivo de sua memória de tamanho fixo. Para aprimorar esses três aspectos, apresentamos o ATLAS, um módulo de memória de longo prazo com alta capacidade que aprende a memorizar o contexto otimizando a memória com base nos tokens atuais e passados, superando a natureza online dos modelos de memória de longo prazo. Com base nessa ideia, apresentamos uma nova família de arquiteturas semelhantes aos Transformers, chamada DeepTransformers, que são generalizações estritas da arquitetura Transformer original. Nossos resultados experimentais em tarefas de modelagem de linguagem, raciocínio de senso comum, tarefas intensivas em recuperação e compreensão de contexto longo mostram que o ATLAS supera o desempenho dos Transformers e dos modelos recorrentes lineares recentes. O ATLAS ainda melhora o desempenho de contexto longo dos Titans, alcançando +80% de precisão no benchmark BABILong com comprimento de contexto de 10M.

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.