Neural Long-Term Memory Module: A memory component that learns to memorize at test time via gradient descent, treating memorization as an optimization problem. The memory update rule is surprise-based — new information is stored proportional to how unexpected it is.
Three Architectural Variants: MAC (Memory as Context), MAG (Memory as Gate), MAL (Memory as Layer) — each integrating long-term memory with attention and MLP in different ways.
Scalable Long-Context Processing: Handles contexts up to 2M tokens, outperforming Transformers and linear recurrent models on long-context benchmarks while remaining competitive on short-context tasks.
Methodology
Architecture: Three-branch design combining (1) attention for short-term in-context memory, (2) a neural long-term memory module, and (3) MLP for persistent parametric memory.
Memory as Optimization: The long-term memory module is updated during inference via gradient descent on a surprise-based objective. Surprise is measured as the prediction error of the memory on new input, encouraging storage of unexpected information.
Forgetting Mechanism: Exponential decay on memory weights allows the module to forget stale information, analogous to biological memory consolidation.
Variants:
MAC: Memory output concatenated with attention context
MAG: Memory output used as a gating signal for attention
MAL: Memory replaces or supplements the attention layer
Results
Outperforms Transformers on long-context language modeling tasks (e.g., 2M-token context)
Competitive with state-of-the-art recurrent models (Mamba, RWKV) on short-context benchmarks
Memory module learns to selectively retain task-relevant information across very long sequences
Scales favorably: memory module size can be adjusted independently of attention heads
Limitations
Neural memory module adds additional parameters and compute overhead at inference time
Gradient-based memory updates increase training and inference complexity
Memory capacity is still bounded; very long contexts may require hierarchical memory structures
Surprise-based update heuristic may not generalize optimally across all task types