# ATLAS: Learning to Optimally Memorize the Context at Test Time

arXiv: https://arxiv.org/abs/2505.23735 · HTML: https://mlsanity.com/paper/2505.23735/

## Key facts

- Cited by 66 papers in the map (data through 14 Oct 2026).
- Builds on 67 of its 93 references that are in the map.
- First posted on arXiv 29 May 2025.

## Builds on (most cited first, 30 of 67)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Gemma 3 Technical Report](https://mlsanity.com/paper/2503.19786/) (2025)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs](https://mlsanity.com/paper/1903.00161/) (2019)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [On the Number of Linear Regions of Deep Neural Networks](https://mlsanity.com/paper/1402.1869/) (2014)
- [Learning to (Learn at Test Time): RNNs with Expressive Hidden States](https://mlsanity.com/paper/2407.04620/) (2024)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [Dense Associative Memory for Pattern Recognition](https://mlsanity.com/paper/1606.01164/) (2016)
- [BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack](https://mlsanity.com/paper/2406.10149/) (2024)

## Built on by (most cited first, 30 of 66)

- [Kimi Linear: An Expressive, Efficient Attention Architecture](https://mlsanity.com/paper/2510.26692/) (2025)
- [Cartridges: Lightweight and general-purpose long context representations via self-study](https://mlsanity.com/paper/2506.06266/) (2025)
- [A Survey on Latent Reasoning](https://mlsanity.com/paper/2507.06203/) (2025)
- [Nested Learning: The Illusion of Deep Learning Architectures](https://mlsanity.com/paper/2512.24695/) (2025)
- [MesaNet: Sequence Modeling by Locally Optimal Test-Time Training](https://mlsanity.com/paper/2506.05233/) (2025)
- [ZipMap: Linear-Time Stateful 3D Reconstruction via Test-Time Training](https://mlsanity.com/paper/2603.04385/) (2026)
- [Efficient Attention Mechanisms for Large Language Models: A Survey](https://mlsanity.com/paper/2507.19595/) (2025)
- [ViT3: Unlocking Test-Time Training in Vision](https://mlsanity.com/paper/2512.01643/) (2025)
- [Memory Caching: RNNs with Growing Memory](https://mlsanity.com/paper/2602.24281/) (2026)
- [Test-Time Training with KV Binding Is Secretly Linear Attention](https://mlsanity.com/paper/2602.21204/) (2026)
- [Understanding LoRA as Knowledge Memory: An Empirical Analysis](https://mlsanity.com/paper/2603.01097/) (2026)
- [Native Hybrid Attention for Efficient Sequence Modeling](https://mlsanity.com/paper/2510.07019/) (2025)
- [Memory-Augmented Transformers: A Systematic Review from Neuroscience Principles to Enhanced Model Architectures](https://mlsanity.com/paper/2508.10824/) (2025)
- [GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent](https://mlsanity.com/paper/2603.13875/) (2026)
- [Fast weight programming and linear transformers: from machine learning to neurobiology](https://mlsanity.com/paper/2508.08435/) (2025)
- [TNT: Improving Chunkwise Training for Test-Time Memorization](https://mlsanity.com/paper/2511.07343/) (2025)
- [Preconditioned DeltaNet: Curvature-aware Sequence Modeling for Linear Recurrences](https://mlsanity.com/paper/2604.21100/) (2026)
- [On the Expressiveness of Softmax Attention: A Recurrent Neural Network Perspective](https://mlsanity.com/paper/2507.23632/) (2025)
- [TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers](https://mlsanity.com/paper/2602.13498/) (2026)
- [Mem3R: Streaming 3D Reconstruction with Hybrid Memory via Test-Time Training](https://mlsanity.com/paper/2604.07279/) (2026)
- [FLARE: Diffusion for Hybrid Language Model](https://mlsanity.com/paper/2606.01774/) (2026)
- [Trellis: Learning to Compress Key-Value Memory in Attention Models](https://mlsanity.com/paper/2512.23852/) (2025)
- [DexWorldModel: Causal Latent World Modeling towards Automated Learning of Embodied Tasks](https://mlsanity.com/paper/2604.16484/) (2026)
- [PERK: Long-Context Reasoning as Test-Time Learning](https://mlsanity.com/paper/2507.06415/) (2025)
- [Q-RAG: Long Context Multi-step Retrieval via Value-based Embedder Training](https://mlsanity.com/paper/2511.07328/) (2025)
- [MDN: Parallelizing Stepwise Momentum for Delta Linear Attention](https://mlsanity.com/paper/2605.05838/) (2026)
- [Fast Spatial Memory with Elastic Test-Time Training](https://mlsanity.com/paper/2604.07350/) (2026)
- [Training LLM Agents for Spontaneous, Reward-Free Self-Evolution via World Knowledge Exploration](https://mlsanity.com/paper/2604.18131/) (2026)
- [Linearizing Vision Transformer with Test-Time Training](https://mlsanity.com/paper/2605.02772/) (2026)
- [Modular TTT: Rethinking Test-Time Training as Composable Modules](https://mlsanity.com/paper/2608.07110/) (2026)

## Abstract (arXiv)

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.