Source-linked AI summary
Titans: Learning to Memorize at Test Time
Ali Behrouz, Peilin Zhong, Vahab Mirrokni
TL;DR
Transformers model direct token dependencies accurately but face quadratic costs and limited context, while recurrent memories compress information into fixed-size states. The paper introduces neural long-term memory and Titans architectures that combine persistent memorization with attention, achieving stronger performance across diverse tasks and scaling beyond 2M-token contexts. However, one Titans design is limited because each layer cannot exploit complementary attention and neural-memory processing.
Problem
Transformers provide accurate dependency modeling but have quadratic time and memory complexity, while linear recurrent memories can overflow during long-context processing.
Method
The paper introduces a neural long-term memory that learns to memorize historical context at test time and combines it with attention through three Titans variants.
Results
Titans outperform modern recurrent models and Transformers across diverse benchmarks, while scaling beyond 2M-token context windows.
Takeaways & Limitations
Titans provide a family of architectures for combining short-term attention with more persistent neural memory across long-context tasks.
Takeaways & Limitations
The memory-as-a-layer design is limited because each layer cannot exploit complementary processing from attention and neural memory.
Abstract
from arXiv · showhide
Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.
1 Introduction
Transformers model current-context dependencies accurately but incur quadratic costs, while recurrent and linear-memory approaches compress history with weaker long-context performance. Titans introduce neural long-term memory and three ways to combine it with attention, targeting efficient memorization and longer-context modeling.
- Motivation: Transformers retrieve key-value associations from the current context window, but their dependency modeling has quadratic time and memory complexity in context length.Their outputs are conditioned on direct token dependencies within the available window.
- Motivation: Linear Transformers reduce memory consumption and scale to longer contexts, but do not show competitive performance with Transformers because kernelization yields linear recurrence.The kernel trick changes attention into a recurrent network that compresses data.
- Memory Perspective: Existing memory architectures face challenges in generalization, length extrapolation, and reasoning, motivating architectures with interconnected short-term, long-term, and other memory systems.The paper frames memory as a collection of systems with different functions rather than a single process.
- Neural Memory: The proposed neural long-term memory learns to memorize at test time, emphasizing surprising or surprise-adjacent tokens and using decay for memory management.Its training can be parallelized through tensorized mini-batch gradient descent and matrix multiplications.
- Titans Architectures: Titans combine limited-window attention, neural long-term memory, and persistent task-independent parameters through context, layer, and gated-branch variants.The three branches serve short-term processing, long-past storage, and task knowledge, respectively.
- Experimental Results: Across language, reasoning, recall, long-context, forecasting, and DNA tasks, Titans outperform modern recurrent models and match or exceed Transformer comparisons while scaling beyond 2M tokens.The reported comparisons include better performance than Transformers at the same context size and competitive performance with full-context Transformers.
2 Preliminaries
This section presents attention, efficient attention, and recurrent models through a memory perspective. Standard attention retains context with high resource demands, whereas linear and recurrent formulations compress history into fixed-size states but may suffer memory overflow.
- Attention: Causal attention maps input sequences to outputs using input-dependent query, key, and value matrices with a softmax mechanism.The formulation operates over an input x in R^N×d_in.
- Efficient Attentions: Standard attention requires at least N × d operators for output computation, increasing memory consumption and lowering throughput on longer sequences.This cost motivates I/O-aware, sparse, approximate, and kernel-based attention alternatives.
- Efficient Attentions: Linear attention replaces softmax with a kernel function, and an identity kernel permits a recurrent formulation that supports efficient inference.The recurrent form reuses accumulated quantities across steps.
- Modern Linear Models and Their Memory Perspective: Recurrent neural networks treat the hidden state as a memory unit updated through write and read operations over time.The memory state M_t indexes the state at time t.
- Modern Linear Models and Their Memory Perspective: Linear Transformers additively compress keys and values into matrix-valued memory, which can overflow on long contexts and damage performance.Prior responses include adaptive forgetting mechanisms and improved write rules such as the Delta Rule.
- Memory Modules: The section situates these mechanisms within broader memory-module research, including associative memory, fast-weight programs, and related efficient-attention models.The cited related work spans recurrent, attention, kernel-based, and memory-network approaches.
3 Learning to Memorize at Test Time
This section introduces a neural long-term memory that learns to memorize and forget information at test time, then describes its training, retrieval, and architectural extensions. The design uses associative-memory objectives, adaptive forgetting, persistent task-memory parameters, and deeper memory modules to improve expressiveness.
- Neural Memory: The neural long-term memory is a meta-model that learns how to memorize and forget data at test time rather than overfitting to training data.It is trained as an online meta-learning process intended to encode abstractions of past history into its parameters.
- Learning Process: Surprise-based updates use input gradients to prioritize memorable events, while momentum-like past surprise preserves influence across time.The formulation combines past surprise and momentary surprise, with data-dependent decay and weighting controlling how surprise persists and how new surprise is incorporated.
- Learning Process: Associative-memory training stores past inputs as key-value pairs and optimizes the memory to learn their mappings in an inner loop.The key and value are produced by linear projections; memory weights are optimized in the inner loop while the architecture’s other parameters are optimized in the outer loop.
- Forgetting Mechanism: Adaptive forgetting manages limited memory capacity by allowing the module to discard information that is no longer needed.A gating value α_t controls forgetting: values near zero preserve past abstractions, whereas values near one can clear the memory.
- Memory Architecture: Deep memory modules are more expressive than matrix-valued memories, whose online linear-regression objective assumes linear historical dependencies.The paper reports that deep memory modules are more effective in practice and identifies richer neural memory architectures as future work.
- Architectural Extensions: Titans incorporate long-term neural memory with persistent, input-independent parameters that store task-level knowledge, while memory retrieval uses a forward pass without weight updates.The persistent parameters complement contextual memory by retaining task abstractions, and retrieval projects an input into a query before reading from memory.
4 How to Incorporate Memory?
Titans incorporate neural long-term memory with attention through three architectural variants, each assigning different roles to historical context, persistent memory, and current information. The designs expose trade-offs between efficiency, effectiveness, and the complementary processing capacity of attention and neural memory.
- Memory perspective: Titans treat attention as short-term memory and neural memory as long-term memory that continuously stores information in its weights.The proposed variants explore how to combine these two memory roles in deep learning architectures.
- Memory as a Context: Memory as a Context retrieves historical information using the current segment, supplies it with persistent memory to attention, and updates long-term memory from the resulting representation.The input sequence is divided into fixed-size segments, with past segments serving as historical information for the current segment.
- Memory as a Context: This context-based design lets attention decide whether long-term information is needed and helps memory retain only useful current-context information.Attention thereby helps manage memory capacity because memorizing every token could cause memory overflow.
- Memory as a Gate: Memory as a Gate uses separate memory and sliding-window-attention branches, combining their outputs through a nonlinear gate without segmenting the input.Sliding-window attention acts as precise short-term memory, while neural memory acts as a fading memory.
- Memory as a Layer: Memory as a Layer compresses past and current context before attention, but its layerwise structure limits the complementary processing available between attention and neural memory.The experiments instantiate this design by replacing the sequence model in an H3-like architecture with the neural memory module.
- Expressivity: Titans are theoretically more expressive than Transformers, diagonal linear recurrent models, and DeltaNet for state-tracking problems beyond TC0.The theorem contrasts Titans with models that are limited to TC0.
5 Experiments
The experiments evaluate Titans and their variants across language, reasoning, long-context retrieval, forecasting, and genomics tasks, comparing them with recurrent, Transformer, and hybrid baselines. Titans generally outperform these baselines, including on challenging long-context benchmarks.
- Experimental setup: Titans are evaluated on language modeling, commonsense reasoning, needle-in-a-haystack, BABILong, time-series forecasting, and DNA modeling tasks.The baselines include state-of-the-art linear recurrent models, Transformers, and recurrent-attention hybrids.
- Language Modeling: The neural memory module achieves the best performance among non-hybrid models on language-modeling perplexity and commonsense-reasoning accuracy.The authors associate this result with weight decay, momentum, surprise-based updates, and deep nonlinear memory.
- Language Modeling: All three Titans variants outperform Samba and Gated DeltaNet-H2, while MAC performs better than MAG on longer dependencies.MAC and MAG outperform MAL despite using the same modules, which the authors attribute to architectural design.
- Needle in a Haystack: Titans achieve the best results across the three S-NIAH tasks at sequence lengths of 2K, 4K, 8K, and 16K, with MAC generally strongest among variants.The authors attribute the advantage to momentum, forgetting through weight decay, and greater memory capacity.
- BABILong Benchmark: Titans outperform all listed BABILong baselines in both few-shot and fine-tuning settings, including GPT-4 and other extremely large models.In fine-tuning, Titans also outperform RMT and a RAG-augmented Llama3.1-8B model with about 70 times fewer parameters.
5.5 The Effect of Deep Memory
The deep-memory experiments examine how memory depth affects performance, throughput, forecasting, genomics, efficiency, and architectural components. Deeper memory improves perplexity but slows training, while the neural memory remains competitive or superior across several benchmarks.
- The Effect of Deep Memory: Increasing memory depth from L_M = 1 to 4 improves perplexity across sequence lengths and increases robustness when models have fewer parameters.Larger models also perform better on longer sequences.
- The Effect of Deep Memory: Training throughput scales linearly with context length, while increasing memory depth produces a linear slowdown in tokens per second.The experiment measures throughput as tokens per second across memory depths L_M = 1, 2, 3, 4.
- Time Series Forecasting: The neural memory module outperforms Mamba-based, linear-based, and Transformer-based architectures on ETT, ECL, Traffic, and Weather forecasting benchmarks.These results broaden the evaluation beyond language tasks.
- DNA Modeling: LMM is competitive with state-of-the-art architectures across downstream GenomicsBenchmarks tasks, measured by top-1 classification accuracy.The evaluation uses pretrained DNA models.
- Efficiency: Titans (MAL) are faster than the compared baselines and the neural memory module in the reported throughput comparison.The authors attribute this advantage mainly to FlashAttention kernels used for sliding-window and full attention.
- Architectural Design: Ablations indicate that deep memory, convolution, momentum, weight decay, and persistent memory each positively contribute to performance.The study changes one architectural component at a time from the neural-memory base model.
6 Conclusion
The paper introduces a neural long-term memory that learns to memorize at test time and uses it in three Titans architecture variants. Experiments across diverse tasks support Titans’ effectiveness for long-context modeling, including contexts larger than 2M tokens.
- 6 Conclusion: Titans incorporate a neural long-term memory as a context, gating mechanism, or layer.The memory is a recurrent model that adaptively memorizes surprising or near-surprising tokens.
- 6 Conclusion: Titans are reported as more effective than Transformers and recent modern linear recurrent models on diverse tasks, particularly for long context.
- 6 Conclusion: Titans scale to context windows larger than 2M tokens with better accuracy than baselines.
- 6 Conclusion: The implementations use PyTorch and JAX, with the training and evaluation code intended for future release.
A Related Work
The paper situates Titans within three categories of prior studies, reflecting diverse perspectives that can independently motivate the design of Titans or their components.
- Prior studies offer diverse perspectives that can independently lead to Titans or their components.
- The related-work discussion organizes these studies into three categories.
- The paper uses this categorization to situate its work in a broader research context.
A.1 Linear Recurrent Models
Linear recurrent models address Transformer efficiency concerns through fast training and inference, progressing from data-independent transitions toward more adaptive mechanisms.
- Linear recurrent models have attracted attention because they offer fast training and inference.
- Early models such as RetNet, LRU, RWKV, S5, and S4 use data-independent transition matrices or decay mechanisms.
- The literature includes multiple generations of linear recurrent models pursuing more efficient sequence processing.
A.2 Transformer-based Architectures
Transformer-based approaches improve long-sequence efficiency through implementation changes, segmentation strategies, or external memory, but each retains important context or integration constraints.
- Transformers: Transformers rely on attention but incur quadratic computational cost, limiting their ability to scale to long context windows.
- Transformers: I/O-aware attention implementations target lower memory consumption and higher throughput for longer sequences.
- Segment-based Transformers: Segment-based Transformers separate inputs into chunks, so their context window is limited by chunk length.
- Segment-based Transformers: Memory mechanisms can transfer information across chunks, addressing a limitation created by segment separation.
- Memory for Large Language Models: External memory modules are added to language models after training, whereas Titans incorporate memory into the initial architecture and train end-to-end.
A.3 Test Time Training and Fast Weight Programs
This section situates test-time training and fast weight programs as approaches for adapting models or providing writable memory. These lines of work motivate memory mechanisms that can respond to data-dependent context.
- Fast Weight Programs: Fast weight programs incorporate dynamic fast programs into recurrent networks as writable memory.They interpret linear layers as key-value associative memory systems.
- Fast Weight Programs: Hebbian and delta rules are the most popular learning rules for fast weight programs.
- Test Time Training: Test-time training updates a model using each test sample or its neighbors before prediction.The approach has also shown promising performance in vision tasks.
B Language Modeling and Common-sense Reasoning Datasets
The evaluation uses language-modeling and common-sense reasoning datasets, including Wikitext, LMB, PIQA, HellaSwag, and several additional reasoning benchmarks. Reported 400M-model baseline results are drawn from prior reports.
- Language Modeling: The language-modeling datasets are Wikitext and LMB.
- Common-sense Reasoning: The common-sense reasoning benchmarks include PIQA, HellaSwag, WinoGrande, ARC-easy, ARC-challenge, SIQA, and BoolQ.
- Baselines: Baseline results for 400M models come from previously reported results.
C Long-term Memory Module (LMM) as a Sequence Model
LMM frames memory updating as online optimization and generalizes recent linear recurrent models with momentum, deep memory, nonlinear recurrence, and forgetting. Its design targets richer memory management and expressive power than simpler recurrent formulations.
- LMM as Generalized Gated DeltaNet: LMM generalizes Gated DeltaNet with momentum-based updates, deep memory, and inter-chunk nonlinear recurrence.Gated DeltaNet is recovered when η_t = 0 in the compared formulations.
- Momentum-based Rule: LMM’s momentum rule incorporates both past and momentary surprise, unlike DeltaNet’s momentary-surprise update.
- Deep Memory: LMM supports deep memory through a gradient-based formulation, whereas Gated DeltaNet is limited to linear matrix-valued memory.The paper associates this design with higher expressive power.
- Non-Linear Recurrence: LMM uses inter-chunk nonlinear recurrence and intra-chunk linear recurrence, unlike the linear recurrence of DeltaNet and Gated DeltaNet.The paper presents this mixed design as increasing expressive power.
- Forgetting Mechanism: LMM adds a forget gate that clears very old information and improves memory management when memory size is fixed.The paper relates forgetting to weight decay and gives a parallel-training implementation.
- Novelty: The neural long-term memory module is presented as the first linear recurrent model with a momentum-based update rule.