ATLAS: Обучение оптимальному запоминанию контекста во время тестирования

ATLAS: Learning to Optimally Memorize the Context at Test Time

May 29, 2025

Авторы: Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, Vahab Mirrokni

cs.AI

Аннотация

Трансформеры утвердились в качестве наиболее популярной архитектуры для моделирования последовательностей, главным образом благодаря их эффективности в задачах извлечения контекста и способности к масштабируемому обучению. Однако их квадратичная сложность по памяти и времени ограничивает применимость в более длинных последовательностях, что побудило исследователей изучить эффективные альтернативные архитектуры, такие как современные рекуррентные нейронные сети (также известные как модули долгосрочной рекуррентной памяти). Несмотря на их недавний успех в разнообразных задачах, они сталкиваются с трудностями в задачах, требующих понимания длинного контекста и экстраполяции на более длинные последовательности. Мы отмечаем, что эти недостатки обусловлены тремя разрозненными аспектами их дизайна: (1) ограниченная емкость памяти, обусловленная архитектурой памяти и отображением признаков входных данных; (2) онлайн-характер обновления, то есть оптимизация памяти только на основе последнего входного сигнала; и (3) менее выразительное управление их памятью фиксированного размера. Для улучшения всех трех аспектов мы представляем ATLAS — модуль долгосрочной памяти с высокой емкостью, который обучается запоминать контекст, оптимизируя память на основе текущих и прошлых токенов, преодолевая онлайн-характер моделей долгосрочной памяти. На основе этого подхода мы представляем новое семейство архитектур, подобных трансформерам, под названием DeepTransformers, которые являются строгими обобщениями оригинальной архитектуры трансформера. Наши экспериментальные результаты в задачах языкового моделирования, здравого смысла, интенсивного запоминания и понимания длинного контекста показывают, что ATLAS превосходит производительность трансформеров и современных линейных рекуррентных моделей. ATLAS также улучшает производительность Titans в задачах с длинным контекстом, достигая +80\% точности на 10 миллионах токенов в тесте 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.