ATLAS: Learning to Optimally Memorize the Context at Test Time
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.
6839158b56bcc85d9f92199b
Ali Behrouz
6839158b56bcc85d9f92199c
Zeman Li
6839158b56bcc85d9f92199d
Praneeth Kacham
6839158b56bcc85d9f92199e
Majid Daliri
6839158b56bcc85d9f92199f
Yuan Deng
6839158b56bcc85d9f9219a0
Peilin Zhong
6839158b56bcc85d9f9219a1
Meisam Razaviyayn
6839158b56bcc85d9f9219a2
Vahab Mirrokni
ATLAS: Learning to Optimally Memorize the Context at Test Time
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.
A new long-term memory module called ATLAS addresses limitations of Transformers in handling long contexts by optimizing memory based on current and past inputs, leading to improved performance in long-context understanding tasks.
Transformers sequence modeling in-context retrieval memory capacity online nature fixed-size memory long-term memory module DeepTransformers language modeling common-sense reasoning recall-intensive tasks long-context understanding ATLAS Titans BABILong benchmark