PaperlayerView exact revision on arXiv

Source-linked AI summary

ATLAS: Learning to Optimally Memorize the Context at Test Time

Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, Vahab Mirrokni

arXiv:2505.23735v1cs.CLcs.AI

TL;DR

Transformers face quadratic cost for long sequences, while recurrent alternatives remain limited by memory capacity, online updates, and memory management. ATLAS addresses these limitations with context memorization and improved memory mechanisms, and its models show consistent gains across diverse benchmarks.

  • Problem

    Transformers have quadratic space and time complexity, while recurrent alternatives often struggle with long-context understanding, retrieval, and extrapolation.

  • Method

    ATLAS combines context memorization, higher-order feature mappings, and advanced memory optimization, alongside OmegaNet and DeepTransformers architectures.

  • Results

    ATLAS models achieve consistent improvements over Transformers and recent recurrent neural network variants across diverse benchmarks.

  • Takeaways & Limitations

    The paper presents ATLAS as a scalable long-term memory approach with theoretical and empirical support for improved context handling.

  • Takeaways & Limitations

    Linear recurrent memories can overflow as context length increases, and their expressivity remains bounded by linear structure; optimizing over all past tokens can also reduce efficiency and increase memory consumption.

Abstract

from arXiv · show

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.

1 Introduction

Transformers support scalable in-context retrieval but their quadratic attention costs limit long-context use, while recurrent alternatives face capacity, update, and memory-management constraints. ATLAS addresses these issues by memorizing context with higher-capacity memory and sliding-window updates, extending the approach to DeepTransformers and evaluating the resulting models across diverse tasks.

  • Motivation: Attention enables scalable in-context retrieval by directly storing and retrieving key-value mappings, but its pairwise computation has quadratic space and time complexity.These costs limit attention-based models in long-context understanding, memorization, and modeling.
  • Motivation: Modern recurrent alternatives use fixed-size associative memories, yet often struggle with long-context understanding, in-context retrieval, and extrapolation to longer sequences.Their designs differ in learning rules, forget gates, and memory architectures.
  • Motivation: The paper identifies three design limitations: online updates based on the current token, restricted memory capacity, and insufficiently expressive memory management.These limitations can cause models to memorize individual tokens without considering broader context.
  • ATLAS: ATLAS measures surprise over a local or global context window, learning to memorize token context at test time rather than relying only on the current input.The paper distinguishes test-time memorization from test-time training because core model parameters and pretrained initial states remain unchanged.
  • ATLAS: Higher-order feature mappings improve memory capacity, while the Omega rule updates memory using all past tokens in a context window instead of only the last token.This combination is intended to manage fixed-size memory and memorize local context rather than individual tokens.
  • Architectures and evaluation: DeepTransformers strictly generalize Transformers, and OmegaNet and ATLAS provide memory modules with polynomial features, improved optimization, and parallelizable training.Experiments evaluate these models on language modeling, common-sense reasoning, recall-intensive, and needle-in-a-haystack tasks against recurrent, local-attention, and Transformer baselines.

2 Preliminaries

This section frames attention and recurrent alternatives as associative-memory sequence models, emphasizing the trade-off between Transformer retrieval quality and long-context efficiency. It introduces fixed-size memory, deep memory modules, attentional bias, and inner-loop memory optimization.

  • Attention: Attention acts as associative memory by computing pairwise token dependencies, but its long-context cost grows with sequence length.The section describes attention as effective for in-context retrieval while requiring at least N×d operations per token.
  • (Linear) Recurrent Models: Linear recurrent models offer parallelizable, linear-time training but compress information into vector- or matrix-valued memory.These models are presented as efficient Transformer alternatives with comparable performance in some settings.
  • (Linear) Recurrent Models: Recurrent memory can overflow as context length increases, while linear structure bounds its expressivity despite improved forget gates.The passage identifies memory overflow and bounded expressivity as limitations of efficient recurrent models.
  • Deep Memory Module: Deep memory modules address limited memory expressivity by updating deep MLP memories with gradient descent or momentum.These architectures use a meta-learning perspective in which the memory itself is optimized.
  • Associative Memory: Associative memory maps keys to values through an attentional-bias objective whose optimization determines the memory update rule.The framework treats sequence modeling as meta in-context learning with separate inner and outer optimization levels.
  • Model Terminology: The paper characterizes models by memory architecture, internal objective, and memory-learning optimizer.This terminology distinguishes the memory mechanism from the model’s other optimized parameters.

3 Learning to Memorize the Context at Test Time

This section develops associative memory with higher capacity and context-aware updates, then introduces Omega rule and sliding-window optimization to balance context memorization with recurrent efficiency.

  • 3.1 Associative Memory with Super Linear Capacity: Matrix-valued memory with delta updates stores at most O(d_k) independent key-value pairs, giving sub-linear capacity relative to its parameter count.The result applies to linearly independent keys under an ℓ2 attentional bias.
  • 3.1 Associative Memory with Super Linear Capacity: Deep MLP memory increases representational power and capacity, with benefits growing with depth, although its upper bound remains subquadratic in key and value dimensions.The theorem gives lower and upper capacity bounds for mappings with linearly independent keys.
  • 3.1 Associative Memory with Super Linear Capacity: Polynomial feature mappings increase the number of key-value mappings that matrix memory can store to O(d_k^p), while approximating more expressive key-query similarities.The mapping uses polynomial degree p and is motivated by approximating Softmax and by input feature gating.
  • 3.2 Long-term Memory with Context Memorization: Global optimization over the entire context can increase computation and memory use, while irrelevant or changing context can cause sub-optimal memorization.Input-dependent gates can also increase parameters and memory usage, reducing recurrent-model advantages.
  • 3.2 Long-term Memory with Context Memorization: Omega rule optimizes memory over a sliding window of past tokens, allowing context memorization rather than greedy memorization of individual tokens.The update uses a decay term for each token’s contribution and is presented as more powerful than the Delta rule.
  • 3.2 Long-term Memory with Context Memorization: OmegaNet applies Omega rule with polynomial kernels, and its formulation is connected to global and local softmax attention through DeepTransformers and sliding-window variants.The paper presents DeepTransformers as strictly generalizing Transformers and introduces DLA to study deep memory.

4 DeepTransformers: Transformers with Deep Memory

The paper reframes attention and recurrent models as associative-memory systems, showing that attention globally optimizes memory while recurrent models typically update online. DeepTransformers generalize Transformers by combining deep memory with expressive feature mappings, while Deep Omega Transformers further replace the Hebbian rule with Omega learning.

  • Connection to Sliding Window Attention: Softmax attention can be viewed as a non-parametric associative-memory solution that globally optimizes its internal objective over the sequence.Restricting optimization to the past c tokens yields sliding window attention, whereas most recurrent models optimize online.
  • Deep Linear Attention: Deep Linear Attention replaces matrix-valued memory with a deep neural network, using nonlinear memory to extend linear attention.Its training can be parallelized through hybrid linear and nonlinear chunk-wise computation.
  • Memory Capacity: Softmax attention uses an effectively unbounded feature-space memory, motivating DeepTransformers to replace polynomial kernels with the 𝜙∗ kernel.This construction is intended to improve memorization of larger contexts.
  • DeepTransformers: DeepTransformers are strict generalizations of Transformers with softmax attention because their linear-memory closed form matches the unnormalized Transformer output.The equivalence is established for the special case of linear memory.
  • Deep Omega Transformer: Deep Omega Transformers replace the original Hebbian learning rule with the Omega learning rule to provide more expressive memory management.In the online c = 1 case, the formulation generalizes Transformers with the Delta Rule while retaining unbounded memory.

5 Atlas: A Locally Optimal Memory with High Capacity

Atlas improves long-term memory by addressing optimization quality and architectural capacity, while using chunk-wise recurrences that support parallel training. Its memory architecture can also be extended with a gated MLP to form Atlas++.

  • Memory Optimization: Muon replaces simple gradient descent for internal memory optimization, approximating second-order information while remaining parallelizable across the sequence.The change is motivated by the risk that gradient descent reaches a low-quality local optimum in longer contexts.
  • Parallel Training: Chunk-wise recurrence computes gradients relative to the last state of the previous chunk, enabling parallel calculation within each chunk.The recurrence is generalized from c = 1 to arbitrary local context lengths using the procedure described for chunked training.
  • Architectural Backbone: Atlas uses projected keys, values, and queries, short convolutions, normalization, and hybrid MAL or MAG variants as its architectural backbone.The architectures are illustrated in Figure 3.
  • Atlas++: Atlas++ extends the memory architecture to a gated MLP using three learnable linear matrices.The paper names the resulting model Atlas++.

6 Experiments

Experiments evaluate Atlas, OmegaNet, DeepTransformers, and related baselines across language modeling, reasoning, recall, memorization, and long-context tasks. Results support gains from context-based memory updates, higher-capacity mappings, deep memory, and improved memory management.

  • Language Modeling and Common-Sense Reasoning: Atlas and OmegaNet achieve the best perplexity and accuracy among non-hybrid models in language modeling and common-sense reasoning.The authors attribute this to memorizing context rather than individual tokens.
  • Language Modeling and Common-Sense Reasoning: DeepTransformers and Dot consistently outperform Transformer++, with gains attributed to deep memory and the Omega rule.Deep memory improves token-dependency modeling, while the Omega rule improves memory management.
  • Long Context: Needle In a Haystack: Atlas outperforms Titans and DeltaNet on needle-in-a-haystack tasks, while hybrid variants extrapolate to sequences 4× their training context size.The authors attribute these results to improved memory capacity; exponential feature mapping is important for longer sequences.
  • Long Context: BABILong Benchmark: +80% accuracy is achieved by Atlas at 10M context length on BABILong, while Titans’ performance drops at that length.Atlas is competitive with Titans through 1M context length and maintains performance at 10M.
  • In-context Recall, MAD Synthetic Benchmark, and Associative Recall: Atlas achieves the best results across recall, memorization, compression, and copying on the MAD benchmark, especially for memorization.Additional evaluations report strong performance on in-context recall and multi-query associative recall.
  • Ablation Study and Scaling Patterns: Ablations find that gated MLPs, hybrid variants, polynomial mappings, and deep memory improve Atlas, while larger local context length improves OmegaNet performance.The reported scaling results also show lower perplexity than baselines across model sizes and good scaling with context length.

7 Conclusion

ATLAS addresses limited capacity, online-only updates, and weak memory management in modern recurrent models through sliding-window learning, higher-order mappings, and advanced optimizers. Across diverse benchmarks, the resulting models consistently improve over Transformers and recent RNN variants.

  • ATLAS targets limited memory capacity, online-only updates, and weak memory management in modern recurrent models.
  • Sliding window learning, higher-order feature mappings, and advanced memory optimizers form ATLAS’s proposed solution.
  • OmegaNet, Atlas, DeepTransformers, and Dot consistently improve over Transformers and recent RNN variants across diverse benchmarks.
  • The paper provides theoretical insight into memory capacity and optimization dynamics, explaining context-length limitations observed in prior work.

A Additional Related Work

Related work develops recurrent and associative-memory alternatives to Transformers to reduce quadratic costs in long-context processing. These approaches draw on fast weight programs, Hebbian or delta learning, and Hopfield-style key-value storage.

  • Modern recurrent alternatives target the quadratic computational complexity and limitations of Transformers for long-context sequences.
  • Fast weight programs use dynamic fast programs integrated into recurrent networks as writable memory stores.
  • Associative-memory formulations learn mappings between keys and values, following the tradition of Hopfield networks.

B Miras Framework

The Miras framework formulates sequence models as meta in-context learners with an inner optimization loop for memory and an outer loop for the remaining model parameters. Different attentional-bias objectives recover linear attention and DeltaNet-style updates.

  • Associative memory maps keys K to values V through an attentional-bias objective that determines memory type and priorities.
  • Iteratively optimizing the attentional-bias objective produces the memory update rule.
  • The inner loop optimizes memory-module parameters while treating the rest of the model as fixed hyperparameters.
  • The outer loop optimizes the model’s remaining parameters, including projections, MLP layers, and convolutions.
  • Dot-product similarity with gradient descent yields linear attention, whereas regression loss yields DeltaNet.

C Supporting Proofs

The supporting analysis characterizes memory capacity through rank and feature-map dimension, showing how polynomial mappings and deep MLP memories expand storage beyond linear attention. It also defines several architectural variants through their internal objectives and optimization procedures.

  • Capacity bounds: A linear matrix memory with dot-product similarity stores at most O(d_k) pairs with linearly independent keys.
  • Optimization: Full-batch gradient descent converges to the minimum-norm interpolating solution under the stated step-size condition.
  • Capacity bounds: Linear and multi-head attention remain limited by an output-context rank of at most d_k, creating a low-rank bottleneck.
  • Architectures: Deep MLP memories can store between O(d_k d_v) and an upper bound determined by hidden dimension and architecture depth, while DLA, SWLA, OmegaNet, and Atlas differ in objectives and optimizers.
  • Capacity bounds: Polynomial feature mappings raise the capacity bound to the dimension of the lifted feature space, scaling as O(d_k^p) for fixed degree p.
  • Capacity bounds: Exact memorization requires the feature matrix to have sufficient rank; when the number of pairs exceeds that rank, the system is unsolvable.

E Experimental Details

The experiments use established language-modeling and commonsense benchmarks, with baseline results reported across model scales. The memory architecture uses a two-layer MLP with residual connections and layer normalization.

  • Experiments cover Wikitext, LMB, PIQA, HellaSwag, WinoGrande, ARC-easy, ARC-challenge, SIQA, and BoolQ.
  • Baseline results are reported for 1.3B and 760M model scales.
  • The memory architecture uses a two-layer MLP with expansion factor 4 and GELU activation.Residual connections and layer normalization are applied at the end of each chunk.

F Additional Experimental Results

Additional experiments evaluate model design choices and performance across long-context, in-context recall, and MAD tasks.

  • Additional experiments evaluate the models in long-context tasks.
  • The experiments include in-context recall evaluation.
  • The evaluation also includes MAD tasks and component analyses.

F.1 Language Modeling and Common-sense Reasoning (Small Scale)

At small scale, Atlas and OmegaNet perform favorably against modern linear recurrent models and Transformers on language modeling and commonsense reasoning tasks. The results also associate context memorization and the attentional-bias objective with improved performance.

  • Atlas and OmegaNet achieve lower perplexity and better downstream-task accuracy than modern linear recurrent models and Transformers.Their fully recurrent variants also outperform hybrid models including Samba and Gated DeltaNet-H2.
  • Hybrid MAG and MAL variants further improve Atlas performance, highlighting complementary roles for recurrent long-term memory and attention.
  • The comparisons associate ℓ2 regression loss as attentional bias with model performance and context memorization with better results than memorizing individual tokens.SWLA outperforms GLA and RetNet in the reported comparison.
Loading 2505.23735v1…