ATLAS: Learning to Optimally Memorize the Context at Test Time | Summary

This retrospective review covers arXiv v1, dated May 29, 2025, of ATLAS: Learning to Optimally Memorize the Context at Test Time, by Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni.

This article explains the key points of ATLAS: Learning to Optimally Memorize the Context at Test Time.

• 2025-05-29 (arXiv)
• Behrouz, Ali, Li, Zeman, Kacham, Praneeth, Daliri, Majid, Deng, Yuan, Zhong, Peilin, Razaviyayn, Meisam, Mirrokni, Vahab.
• Paper

Read this article in Korean

Summary

• ATLAS treats a context-dependent memory as a small model updated as a sequence is processed. It combines a recent window of key–value associations, a richer key representation, and an optimizer inspired by Muon. The authors report improvements over several recurrent-memory baselines on language modeling and retrieval tasks. Full attention nevertheless remains stronger on some recall measurements. These are comparisons under the paper’s experimental settings, not a guarantee for every long-context application.
• Test-time memorization means updating a memory state with the current input. It is distinct from permanently retraining shared model weights: the state carries information through a sequence.

1. Introduction

Full attention can retrieve a previous token directly, but its cost grows with sequence length. A fixed-size recurrent state offers a different tradeoff: its size need not grow with every token, but information must be compressed and can be overwritten. The paper asks how memory representation and optimization can improve retention and retrieval.

2. Preliminaries

2.1 Backgrounds

Let a token produce a key, a value, and a query. A memory model maps a key toward its associated value; later, a query retrieves a prediction. In a linear memory, the state is a matrix. A loss on the prediction error supplies an update. The outer model learns projections and other parameters during training; the memory state is updated as a context is processed. The paper distinguishes the expressiveness of the memory from the optimization rule used to update it.

3. Learning to Memorize the Context at Test Time

3.1 Associative Memory with Super Linear Capacity

For linear memory trained with squared error, the authors analyze the number of independent key–value associations it can store. Under the stated independence and model assumptions, the bound scales as \(O(d_k)\), where \(d_k\) is key dimension. A degree-p polynomial feature map yields an \(O(d_k^p)\) upper bound in the corresponding analysis. This is a capacity bound under the paper’s conditions, not a promise that arbitrary real sequences will achieve it. Feature expansion changes the representational and computational tradeoffs.

3.2 Long-term Memory with Context Memorization

A single-token update fits the current association. Omega instead combines the losses of a window of recent associations with input-dependent weights. A window of one recovers the single-association case. OmegaNet combines this objective with polynomial key and query representations. The recent window used for an update is not the total history available in recurrent memory: information from earlier tokens can remain in the carried state.

Figure 2 shows sliding-window attention in its first panel. The remaining panels show ATLAS with context windows of 1, 4, and 7; the diagrams compare dependency patterns rather than measured accuracy.

The first panel shows sliding-window attention; the other panels show ATLAS dependencies for the indicated window sizes.

3.3 Parallelizing Omega Rule

The authors derive a chunk-wise formulation that computes updates relative to the memory state at the preceding chunk boundary and uses a causal window mask. Computations within a chunk can be parallelized while memory still recurs across chunks. This is a particular approximation and training formulation; it does not eliminate sequential dependence or establish a deployment speedup.

4. DeepTransformers: Transformers with Deep Memory

4.1 Online and Local Context Optimization of Memory

The paper connects attention to weighted regression: a query weights associations, and a fitted mapping produces an output. This view motivates Deep Transformers and Dot, which uses an exponential feature-space construction with the Omega objective.

4.2 Memory Capacity and Exponential Kernels

A special case recovers the unnormalized Transformer formulation discussed in the paper. That qualification matters: the derivation does not make every practical softmax-attention implementation interchangeable with a constant-cost recurrent memory.

4.3 Deep Omega Transformer (Dot): Transformers with Omega learning rule

Dot combines the exponential feature construction with the Omega objective.

5. Atlas: A Locally Optimal Memory with High Capacity

ATLAS applies Muon-inspired updates to the memory. Momentum is approximately made semi-orthogonal by Newton–Schulz iterations, which the authors interpret as an approximation related to second-order optimization. Their analysis describes a locally optimal update, not a general proof of global optimality for a trained language model. More elaborate updates can require additional computation.

5.1 Parallel Training

The chunk-wise training construction preserves inter-chunk recurrence, as described above.

The default memory is a two-layer residual MLP. ATLAS++ uses a gated MLP, and the paper also evaluates hybrids that combine memory with local attention. These variants should be distinguished when comparing results.

Figure 3 shows, from left to right, the ATLAS architecture, ATLAS MAG, ATLAS MAL, and the ATLAS layer block design. The diagram labels these components; the text separately describes the default residual memory and gated ATLAS++ variant.

The left three panels distinguish ATLAS, MAG, and MAL; the right panel expands the ATLAS layer.

6. Experiments

6.1 Language Modeling and Common-Sense Reasoning

The paper reports perplexity, where lower is better, and accuracy, where higher is better. In Table 2, the 760M ATLAS model trained on 30B tokens reaches WikiText perplexity 18.92 and mean reasoning accuracy 52.77, compared with 19.16 and 52.56 for OmegaNet and 25.21 and 48.69 for Transformer++. At 1.3B parameters and 100B tokens, ATLAS reports 14.97 and 57.62; ATLAS++ reports 14.40 and 58.03. Individual datasets and variants differ, so the table does not establish a win on every metric.

Compare perplexity and accuracy within the stated model-size and token-budget rows; their favorable directions differ.

The experimental setup (PDF page 14) gives model sizes of 340M, 400M, 790M, and 1.3B, whereas Table 7 (page 30) gives 170M, 340M, 760M, and 1.3B. In particular, the setup says 790M while the results in Table 2 (page 15) are labeled 760M. This discrepancy is unresolved; it is not established to be a typo. The comparison above uses 760M because that is the label attached to the reported Table 2 results, not because the training configuration has been independently resolved. Section 6.1 also describes ATLAS and OmegaNet as the best non-hybrid models in both measures, but at 1.3B Table 2 gives OmegaNet better WikiText perplexity, 14.91 versus 14.97 for ATLAS. The table is the basis for the comparison here.

6.2 Long Context: Needle In a Haystack

On RULER S-NIAH-N at 16K tokens, Table 3 reports 84.0 for ATLAS and 80.2 for Titans; ATLAS MAG reaches 97.0. The models were trained with a 4K context. This synthetic retrieval task does not measure document understanding in general.

The 16K column tests retrieval beyond the 4K training context, under this synthetic benchmark.

6.3 Long Context: BABILong Benchmark

The BABILong experiment reports accuracy above 80% at 10M tokens for the evaluated ATLAS hybrid. The plotted setting includes benchmark fine-tuning and the specified MAC architecture without persistent memory tokens. This is neither zero-shot performance nor a universal long-context accuracy claim.

These curves reflect benchmark fine-tuning and the specified MAC architecture; they do not establish zero-shot accuracy.

6.4 Learnability Experiments

Controlled synthetic experiments compare memory architectures and expansion sizes across five mappings: low-rank linear mapping, an MLP mapping, attention followed by an MLP, an MLP whose inputs are attention outputs, and sliding-window attention followed by an MLP. In the attention-output-input setting the inputs are correlated, unlike the independently sampled inputs in the basic MLP setting. The authors report that this correlated-input setting learns the mapping faster than the basic MLP setting in their experiment. Their attention uses randomly initialized vectors, so its behavior does not reproduce all the structure of trained language-model attention.

6.5 Additional Experiments: In-context Recall, MAD Synthetic Benchmark, and Associative Recall

On MAD, Table 4 reports an average of 79.50 for ATLAS, 78.98 for OmegaNet, 76.44 for Titans, and 75.46 for Transformer. Several individual results are tied or differ in direction. On the separate in-context recall suite, Table 5 reports 43.70 for ATLAS, 43.13 for OmegaNet, and 53.55 for the full Transformer. Recurrent memory is not always a substitute for direct attention in these experiments.

The average score for the full Transformer remains higher than the recurrent baselines on this recall suite.

6.6 Ablation Study and Scaling Patterns

In Table 6, reducing the Omega window to one changes perplexity from 19.97 to 21.98 and average accuracy from 52.77 to 49.26; removing polynomial features gives 22.14 and 50.57. Removing Muon gives 19.65 perplexity and 52.56 accuracy: perplexity improves while mean accuracy falls slightly. This conflicts with the table caption’s claim that every component contributes positively; the tabulated values show a metric-dependent tradeoff.

Removing Muon lowers reported perplexity while slightly lowering mean reasoning accuracy in this setting.

The scaling and multi-query associative recall plots provide further evidence for the tested settings. Perplexity or FLOP plots alone do not establish production latency or peak memory.

7. Conclusion

ATLAS treats recurrent state as learned associative memory and studies what it can represent and how it is updated. The windowed objective and higher-capacity memory are associated with improvements on several tests. Results depend on architecture, data, metrics, and training setup; full attention remains stronger on reported average recall, and the fine-tuned 10M-token experiment does not establish general zero-shot performance.

The appendices contain related work, the meta-learning formulation, capacity arguments, architecture details, and experimental settings. The analysis depends on its assumptions, and the size-reporting inconsistency above deserves attention before replication.

Appendix

• A Additional Related Work: The authors place the proposed models alongside linear recurrent networks, input-dependent updates, delta-rule methods, and deeper neural memories. These are design families rather than controlled new comparisons in the appendix.
• B Miras Framework: Miras expresses sequence models through an associative memory, an attention bias, a retention mechanism, and a memory-learning rule. The framework supplies a common language for comparing different architectures.
• C Supporting Proofs: The supporting arguments analyze capacity under specified model and data assumptions. The polynomial-map result counts feature-space dimensions and limits the independent associations that can be fit exactly; it should not be interpreted as an empirical capacity guarantee.
• D Detailed Formulations of All Architectures: The appendix spells out the constructions for deep linear attention, sliding-window linear attention, OmegaNet, and Atlas. Their objectives and recurrence rules differ, even where they share key, query, and value projections.
• E Experimental Details: The authors describe FineWeb training, architectural settings, and baseline sources. Table 7 lists the model configurations and token budgets; its size list differs from the prose in the main experimental setup.
• F Additional Experimental Results: Table 8 presents smaller-scale language-modeling and reasoning results, including the 340M setting. These additional comparisons do not resolve the size-list discrepancy in the main text.

Brief Thoughts

The most useful distinction is between what memory can express and how an update preserves useful past information. The gap to full attention on recall and the mixed Muon ablation make a strong case for evaluating both accuracy and resource costs in the intended application.