title: ATLAS: Learning to Optimally Memorize the Context at Test Time
summary: 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.
paper_authors._id: 6839158b56bcc85d9f92199b
paper_authors.name: Ali Behrouz
paper_authors._id: 6839158b56bcc85d9f92199c
paper_authors.name: Zeman Li
paper_authors._id: 6839158b56bcc85d9f92199d
paper_authors.name: Praneeth Kacham
paper_authors._id: 6839158b56bcc85d9f92199e
paper_authors.name: Majid Daliri
paper_authors._id: 6839158b56bcc85d9f92199f
paper_authors.name: Yuan Deng
paper_authors._id: 6839158b56bcc85d9f9219a0
paper_authors.name: Peilin Zhong
paper_authors._id: 6839158b56bcc85d9f9219a1
paper_authors.name: Meisam Razaviyayn
paper_authors._id: 6839158b56bcc85d9f9219a2
paper_authors.name: Vahab Mirrokni
paper_title: ATLAS: Learning to Optimally Memorize the Context at Test Time
paper_summary: 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.
paper_ai_summary: 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.
paper_ai_keywords: 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