ATLAS: Optimale Kontextspeicherung zur Testzeit erlernen

ATLAS: Learning to Optimally Memorize the Context at Test Time

May 29, 2025

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

cs.AI

Zusammenfassung

Transformers haben sich als die beliebtesten Backbones in der Sequenzmodellierung etabliert, vor allem aufgrund ihrer Effektivität bei In-Context-Retrieval-Aufgaben und ihrer Fähigkeit, im großen Maßstab zu lernen. Ihre quadratische Speicher- und Zeitkomplexität begrenzt jedoch ihre Anwendbarkeit bei längeren Sequenzen und hat Forscher dazu motiviert, effektive alternative Architekturen wie moderne rekurrente neuronale Netze (auch bekannt als Long-Term Recurrent Memory Module) zu erforschen. Trotz ihres jüngsten Erfolgs in verschiedenen Downstream-Aufgaben, haben sie Schwierigkeiten bei Aufgaben, die ein langes Kontextverständnis und eine Extrapolation auf längere Sequenzen erfordern. Wir beobachten, dass diese Schwächen aus drei getrennten Aspekten in ihrem Design resultieren: (1) begrenzte Speicherkapazität, die durch die Architektur des Speichers und die Feature-Mapping des Eingabesignals beschränkt ist; (2) der Online-Charakter der Aktualisierung, d.h. die Optimierung des Speichers nur in Bezug auf die letzte Eingabe; und (3) weniger ausdrucksstarke Verwaltung ihres festen Speichers. Um alle drei Aspekte zu verbessern, präsentieren wir ATLAS, ein Long-Term Memory Module mit hoher Kapazität, das lernt, den Kontext zu speichern, indem es den Speicher basierend auf den aktuellen und vergangenen Tokens optimiert und so den Online-Charakter von Long-Term Memory Modellen überwindet. Aufbauend auf dieser Erkenntnis stellen wir eine neue Familie von Transformer-ähnlichen Architekturen vor, genannt DeepTransformers, die strikte Verallgemeinerungen der ursprünglichen Transformer-Architektur darstellen. Unsere experimentellen Ergebnisse in den Bereichen Sprachmodellierung, Common-Sense Reasoning, Recall-intensive und Long-Context-Verständnis-Aufgaben zeigen, dass ATLAS die Leistung von Transformers und aktuellen linearen rekurrenten Modellen übertrifft. ATLAS verbessert weiterhin die Long-Context-Leistung von Titans und erreicht eine +80\% Genauigkeit bei einer Kontextlänge von 10M im BABILong Benchmark.

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.