Titans: Learning to Memorize at Test Time — Mathematical Foundations of Neural Long-Term Memory, Surprise-Driven Gradient Updates, Adaptive Forgetting, and Context Scaling

Sequence modeling in large language models has long been caught in a fundamental trade-off between attention mechanisms and recurrent architectures. Softmax attention provides precise associative recall over all tokens in the active context window, but incurs quadratic computational and memory complexity, $O(N^2)$, with respect to sequence length $N$. Conversely, linear recurrent models, such as State Space Models (Mamba, Mamba-2) and linear Transformers (GLA, RetNet, DeltaNet), compress sequence history into fixed-size vector or matrix hidden states. While achieving linear $O(N)$ computational complexity and constant inference memory, their finite state capacity creates an information bottleneck that degrades long-range retrieval and multi-hop reasoning over extended horizons.

In January 2025, researchers from Google Research (Ali Behrouz, Peilin Zhong, and Vahab Mirrokni) introduced Titans, a family of neural architectures that resolves this trade-off by treating memory through an online meta-learning framework. Rather than compressing history into passive recurrent vectors or maintaining unbounded key-value caches, Titans deploy a deep neural long-term memory module that actively learns to memorize, retrieve, and forget information during test-time inference.

The Three-Tier Memory Hierarchy

Human cognition relies on a confederation of distinct memory systems operating over varying temporal scales and representations. Titans operationalize this principle by decomposing sequence modeling into three interconnected memory modules:

1. Short-Term Memory (Attention): High-precision associative recall over local context windows using standard causal or sliding-window attention. It captures exact token dependencies within immediate token neighborhoods.
2. Long-Term Memory (Neural Memory Module): A parameter-based associative memory implemented as a multi-layer neural network (such as a multi-layer perceptron). Its weights are dynamically updated at test time through gradient descent on incoming data.
3. Persistent Memory (Meta-Task Inductive Bias): A set of learnable, data-independent prefix vectors prepended to the sequence. These vectors store immutable task-level knowledge, emulate feed-forward network key-value priors, and eliminate attention sink artifacts at sequence beginnings.

By delegating long-range historical context to neural parameter updates while retaining local attention for exact short-range dependencies, Titans decouple context length from attention compute overhead.

Mathematical Foundations of Neural Long-Term Memory

The Long-Term Memory Module (LMM), denoted as $\mathcal{M}$, acts as an associative memory function parameterized by weights $W$. Unlike classical fast weight programmers or linear RNNs that store representations in linear matrices ($y_t = Q_t \mathcal{M}_t$), $\mathcal{M}$ can be a deep non-linear network whose parameter trajectory encodes sequence history.

1. Associative Memory Loss Formulation

Given an input sequence $x \in \mathbb{R}^{N \times d_{\text{in}}}$, the model projects each token $x_t$ into key and value representations via learnable linear projections $W_K, W_V \in \mathbb{R}^{d_{\text{in}} \times d_{\text{in}}}$:

$$\mathbf{k}_t = x_t W_K, \quad \mathbf{v}_t = x_t W_V$$

The objective of the neural memory at step $t$ is to bind the key $\mathbf{k}_t$ to its corresponding value $\mathbf{v}_t$. The associative reconstruction error is defined by the squared Euclidean loss:

$$\ell(\mathcal{M}_{t-1}; x_t) = \frac{1}{2} \|\mathcal{M}_{t-1}(\mathbf{k}_t) - \mathbf{v}_t\|_2^2$$

In the meta-learning hierarchy, the outer loop optimizes global parameters (including projection matrices $W_Q, W_K, W_V$) across training corpora, while the inner loop optimizes the memory weights $\mathcal{M}$ on-the-fly during inference.

2. Surprise-Driven Gradient Updates with Momentum

In biological memory systems, events that deviate significantly from prior expectations produce heightened neural plasticity and stronger encoding. Titans formalize this concept by defining momentary surprise as the gradient of the associative loss with respect to current memory parameters:

$$\nabla \ell(\mathcal{M}_{t-1}; x_t) = (\mathcal{M}_{t-1}(\mathbf{k}_t) - \mathbf{v}_t) \nabla_{\mathcal{M}} \mathcal{M}_{t-1}(\mathbf{k}_t)$$

If memory parameters were updated solely via standard stochastic gradient descent ($\mathcal{M}_t = \mathcal{M}_{t-1} - \theta_t \nabla \ell$), two pathologies arise:

• Momentary Volatility: A single aberrant token could corrupt the entire memory state.
• Gradient Vanishing in Plateaus: After several high-surprise steps, gradients may vanish in flat optimization regions, causing the model to miss subsequent context shifts.

To preserve temporal coherence across token streams, Titans introduce a momentum accumulator $S_t$, which maintains a memory of surprise:

$$S_t = \eta_t S_{t-1} - \theta_t \nabla \ell(\mathcal{M}_{t-1}; x_t)$$

Here, $\theta_t \in (0, 1)$ is a data-dependent learning rate (controlling momentary surprise absorption), and $\eta_t \in [0, 1]$ is a data-dependent momentum decay factor (controlling historical surprise persistence). When a prompt switches topics or contexts, $\eta_t \to 0$ flushes past surprise; when processing cohesive passages, $\eta_t \to 1$ accumulates gradient trajectories.

3. Adaptive Forgetting Gate

Even deep neural networks possess finite parameter capacity. To prevent memory saturation over multi-million-token contexts, Titans integrate an adaptive forgetting mechanism that functions as generalized dynamic weight decay:

$$\mathcal{M}_t = (1 - \alpha_t) \mathcal{M}_{t-1} + S_t$$

The scalar gating parameter $\alpha_t \in [0, 1]$ is dynamically predicted from input token features $x_t$:

• When $\alpha_t \to 0$, memory weights are preserved without decay.
• When $\alpha_t \to 1$, the historical memory state is cleared, enabling clean context resets.

Combining the momentum update and forgetting gate yields the complete inner-loop recurrence:

$$\begin{aligned} S_t &= \eta_t S_{t-1} - \theta_t \nabla \ell(\mathcal{M}_{t-1}; x_t) \\ \mathcal{M}_t &= (1 - \alpha_t) \mathcal{M}_{t-1} + S_t \end{aligned}$$

4. Memory Retrieval

To retrieve context relevant to the current query without perturbing the learned parameters, the model computes a query projection $\mathbf{q}_t = x_t W_Q$ and evaluates a pure forward pass through the updated memory state $\mathcal{M}_t^*$:

$$y_t = \mathcal{M}_t^*(\mathbf{q}_t)$$

The asterisk notation $\mathcal{M}^*$ denotes forward evaluation without triggering backward weight updates.

Persistent Memory and Inductive Bias

Contextual memory modules are strictly data-dependent. However, sequence modeling also requires invariant, task-level abstractions (such as syntactic rules, reasoning protocols, or output formatting constraints).

Titans incorporate Persistent Memory, defined as a matrix of $N_p$ learnable, data-independent parameter vectors $P = [p_1, p_2, \dots, p_{N_p}]$ prepended to the sequence representation:

$$x_{\text{new}} = [P \,\|\, x]$$

Persistent memory fulfills three key architectural functions:

• Task Knowledge Storage: It stores static procedural priors independent of input sequence dynamics, behaving analogously to key-value pairs in feed-forward layers ($FFN(x) = W_V \text{Softmax}(W_K x)$).
• Attention Sink Mitigation: Causal attention mechanisms naturally concentrate disproportionate probability mass on initial sequence tokens (the attention sink phenomenon). Learnable prefix parameters absorb this structural bias without degrading real token embeddings.
• Warm-Start Condition for Neural Memory: Prefix vectors provide stable initial hidden states and projection baselines before long-range sequence context accumulates.

Architectural Integration Variants

Titans explore three architectural topologies for combining neural long-term memory with attention:

1. Memory as Context (MAC)

In the MAC topology, the sequence is partitioned into segments $S^{(t)}$ of chunk size $C$. Before processing segment $S^{(t)}$, the previous memory state $\mathcal{M}_{t-1}$ is queried to retrieve historical context summary $h_t = \mathcal{M}_{t-1}^*(S^{(t)} W_Q)$.

The retrieved vector is concatenated with persistent memory $P$ and the active segment $S^{(t)}$:

$$\tilde{S}^{(t)} = [P \,\|\, h_t \,\|\, S^{(t)}]$$

The composite representation is passed through full self-attention ($y_t = \text{Attn}(\tilde{S}^{(t)})$), and the resulting contextualized representation is used to update the neural memory for subsequent segments:

$$\mathcal{M}_t = \mathcal{M}_{t-1}(y_t), \quad o_t = y_t \otimes \mathcal{M}_t^*(y_t)$$

MAC enables standard attention to filter which historical information is relevant, preventing redundant token memorization.

2. Memory as Gating (MAG)

MAG processes the input sequence through two parallel branches:

• Branch A executes prefix-augmented sliding-window attention ($\text{SW-Attn}^*(\tilde{x})$).
• Branch B executes the neural long-term memory recursion ($\mathcal{M}(\tilde{x})$).

The two representations are combined via learned non-linear gating:

$$o = \text{SW-Attn}^*(\tilde{x}) \otimes \mathcal{M}(\tilde{x})$$

This formulation creates a multi-head hybrid structure where attention heads capture high-frequency local syntax while memory heads maintain broad sequence continuity.

3. Memory as a Layer (MAL)

MAL stacks the neural memory module sequentially within the transformer layer pipeline, alternating between LMM layers and sliding-window attention blocks:

$$y = \mathcal{M}(\tilde{x}), \quad o = \text{SW-Attn}(y)$$

This configuration mirrors hybrid models such as H3 and Griffin, serving as a direct drop-in replacement for standard linear recurrent or convolution blocks.

Hardware-Efficient Parallel Training

A recurring bottleneck for recurrent architectures with test-time updates is execution latency during training. While naive recurrent formulation requires sequential step-by-step weight updates ($O(N)$ sequential operations), Titans leverage chunkwise tensorization and parallel associative scans.

Chunkwise Mini-Batch Formulation

Dividing a sequence of length $N$ into sub-chunks of size $b$, the unrolled gradient descent update across a chunk can be expressed in closed algebraic form:

$$\mathcal{M}_t = \beta_t \mathcal{M}_0 - \sum_{i=1}^t \theta_i \frac{\beta_t}{\beta_i} \nabla \ell(\mathcal{M}_{t'}; x_i)$$

where $t' = t - (t \bmod b)$ represents the chunk boundary index, and $\beta_i = \prod_{j=1}^i (1 - \alpha_j)$ represents cumulative forgetting decay factors.

For linear memory weights $W_0$, the intra-chunk gradient sum collapses into matrix multiplications over chunk matrices $X$:

$$\sum_{i=1}^b \theta_i \frac{\beta_b}{\beta_i} \nabla \ell(W_0; x_i) = \Theta_b \mathbf{B}_b (W_0 X - X) X^\top$$

where $\Theta_b = \text{diag}([\theta_1, \dots, \theta_b])$ and $\mathbf{B}_b$ encodes decay ratios.

Because the momentum recurrence $S_t = \eta_t S_{t-1} - \theta_t u_t$ is a first-order linear recurrence over pre-computed gradient terms $u_t = \nabla \ell(\mathcal{M}_{t'}; x_t)$, it satisfies associative properties. It is computed across GPU/TPU hardware using a parallel associative scan, maintaining $O(\log b)$ span complexity per chunk and saturating tensor core matrix multiplication units.

Theoretical Expressivity and Empirical Evaluation

Beyond $TC^0$ Circuit Complexity

Standard Transformers (without chain-of-thought) and diagonal linear recurrent models (such as Mamba, S4, and DeltaNet) are theoretically constrained to the circuit complexity class $TC^0$ for constant-depth configurations. Consequently, they cannot solve certain multi-step state tracking problems (such as parity checks or graph connectivity) without intermediate reasoning tokens.

• Theorem 4.1: Titans equipped with deep neural memory modules transcend the $TC^0$ complexity barrier, providing strictly higher computational expressivity for dynamic state tracking and algorithmic sequence modeling.

Empirical Benchmarks

Experimental evaluations across language modeling, long-context retrieval, genomics, and time-series benchmarks demonstrate:

• Needle-In-A-Haystack and BABILong: Titans scale effectively beyond 2,000,000 (2M+) tokens in context window length, achieving near-perfect retrieval accuracy where standard sliding-window Transformers, Mamba-2, and linear recurrent baselines experience catastrophic recall decay.
• Language Modeling Perplexity: Across standard pre-training benchmarks (such as The Pile and SlimPajama), Titans achieve lower validation perplexity than LLaMA-style Transformers and Mamba architectures at equivalent parameter and FLOP budgets.
• Ablation on Deep Memory: Multi-layer non-linear memory modules ($L_{\mathcal{M}} \ge 2$) systematically outperform single-layer linear memory matrices ($L_{\mathcal{M}} = 1$), validating the necessity of deep non-linear parameter spaces for complex historical context compression.

Conclusion and Future Directions

Titans demonstrate that the dichotomy between attention mechanisms and recurrent compression is not insurmountable. By framing long-term memory as an online optimization process driven by surprise gradients, momentum accumulators, and adaptive forgetting, Titans offer a principled foundation for million-token context processing with constant inference memory footprints and parallel training throughput.

Sources

• Behrouz, Zhong, & Mirrokni (2025). Titans: Learning to Memorize at Test Time. arXiv:2501.00663
• Sun et al. (2024). Learning to (Learn at Test Time): RNNs with Expressive Hidden States. arXiv:2407.04620
• Dao & Gu (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. arXiv:2405.21060
• Yang et al. (2024). Gated Linear Attention Transformers with Hardware-Efficient Training. arXiv:2312.06635
• Vaswani et al. (2017). Attention Is All You Need. NeurIPS 2017
• Sukhbaatar et al. (2019). Augmenting Self-Attention with Persistent Memory. arXiv:1907.01470