Google Unveils Titans: Neural Networks That Learn and Memorize at Test Time

The foundational dogma of large language model inference—that network weights remain strictly frozen while all dynamic computation lives in volatile activations—has hit an architectural wall. While self-attention suffers from quadratic complexity and linear recurrent models (SSMs) fail to compress long horizons without destructive information loss, Google Research has introduced Titans: a neural architecture family that dynamically learns to memorize at test time.

Authored by Ali Behrouz, Peilin Zhong, and Vahab Mirrokni, Titans introduces an explicit, trainable Neural Long-Term Memory (NLTM) module that updates its internal parameters during the forward pass itself. By framing memory storage as an online optimization problem guided by instantaneous "surprise" gradients, Titans maintains sub-quadratic computational scaling, achieves linear inference complexity, and preserves near-perfect retrieval across sequences exceeding 2,000,000 tokens.

The Failure Mode of KV Caches and Linear States

Modern sequence modeling has been caught in a clear trade-off between expressive fidelity and computational scaling:

Full Attention (Transformers): Acts as an exact associative memory by maintaining the complete key-value (KV) history. However, the memory footprint scales linearly with sequence length while computation scales quadratically, making multi-million token context windows economically and hardware-prohibitive.
Linear Recurrent Networks & SSMs (Mamba, RWKV, Linear Transformers): Compress historical context into fixed-size matrix or vector hidden states. While offering $O(1)$ memory updates during autoregressive generation, they inevitably hit the memory capacity wall—compressing non-linear dependencies over millions of tokens leads to severe catastrophic forgetting on complex multi-hop retrieval tasks.

Titans breaks this compromise by decoupling sequence processing into three distinct cognitive tiers: Short-Term Memory (sliding-window multi-head attention), Persistent Memory (learned task-specific context parameters), and Long-Term Neural Memory (an online-updated neural module).

                    ┌──────────────────────────────────────────┐
                    │               Input Stream               │
                    └─────────────────────┬────────────────────┘
                                          │
             ┌────────────────────────────┴───────────────────────────┐
             ▼                                                        ▼
   ┌───────────────────┐                                    ┌───────────────────┐
   │ Short-Term Memory │                                    │  Neural Long-Term │
   │  (Sliding-Window  │                                    │      Memory       │
   │    Attention)     │                                    │  (Online Gradient │
   └─────────┬─────────┘                                    │     Updates)      │
             │                                              └─────────┬─────────┘
             │                   ┌───────────────────┐                │
             │                   │ Persistent Memory │                │
             │                   │  (Prefix Weights) │                │
             │                   └─────────┬─────────┘                │
             ▼                             ▼                          ▼
      ───────────────────────────────────────────────────────────────────────
                          Integrated Context / Gated Output

Inside the Architecture: Surprise, Momentum, and Adaptive Decay

The fundamental innovation of Titans lies in how the long-term neural memory updates its internal parameters $M_t$ at step $t$. Instead of simply accumulating hidden states through linear transitions, Titans formulates memory storage as minimizing an associative reconstruction loss:

$$\ell(M_{t-1}; x_t) = | M_{t-1}(k_t) - v_t |^2$$

Where $k_t$ and $v_t$ are linear projections of the input token sequence $x_t$. The update mechanism operates via three mathematically coordinated stages:

1. The Gradient-Based Surprise Metric

When the model encounters a token, it evaluates how well its existing memory weights predict the associated value representation. The gradient of the associative loss with respect to memory parameters represents the instantaneous Surprise:

$$g_t = \nabla_{M_{t-1}} \ell(M_{t-1}; x_t)$$

If the incoming token is fully predictable based on prior history, the gradient norm is near zero, and memory parameters remain unchanged. If the token contains novel, high-information content, the gradient magnitude spikes, driving substantial parameter updates.

2. Surprise Momentum (Meta-Memory)

Human cognition does not merely respond to immediate novelty; it tracks continuous shifts in context. Titans introduces a momentum state $S_t$ that accumulates past surprise vectors over a temporal window:

$$S_t = \eta_t S_{t-1} - \theta_t g_t$$

Here, $\theta_t$ serves as an adaptive learning rate and $\eta_t$ represents a dynamic momentum factor, both generated via data-dependent projections. This allows the model to differentiate between transient noise and persistent semantic shifts.

3. Adaptive Weight Decay (Controlled Forgetting)

Unbounded memory growth leads to parameter saturation and catastrophic interference. Titans applies a data-dependent forgetting gate $(1 - \alpha_t)$ to clear redundant historical associations:

$$M_t = (1 - \alpha_t) M_{t-1} + S_t$$

By parameterizing $\alpha_t \in (0, 1)$, the network learns when to purge stale contexts and retain critical landmark information.

Three Architectural Blueprints: MAC, MAG, and MAL

Google evaluated three structural paradigms for integrating test-time neural memory with local attention:

Memory as Context (MAC): The neural memory processes past chunks and generates compressed memory tokens $h_t = M_{t-1}^*(q_t)$. These tokens are prepended alongside learned persistent tokens directly into the sliding-window attention context. The attention block then jointly computes dependencies across immediate tokens and recalled historical representations.
Memory as Gate (MAG): Local sliding-window attention and the neural long-term memory process the sequence in parallel. Their outputs are merged via a non-linear gating mechanism: $y_t = g_t \odot \text{Attn}(x_t) + (1 - g_t) \odot M_t(x_t)$.
Memory as Layer (MAL): Neural memory modules are inserted sequentially between transformer layers, directly replacing standard feed-forward or linear attention blocks.

Among the three, MAC demonstrated superior empirical performance, particularly in tasks requiring multi-hop logical deductions across distant document segments.

Empirical Benchmarks: Dominating at 2M+ Context Lengths

Titans was evaluated across extensive NLP benchmarks, long-context reasoning tasks (BABILong, Needle In A Haystack), genomic modeling, and long-range time-series forecasting.

Architecture	Context Scalability	Needle Retrieval (1M Tokens)	BABILong Accuracy	Inference Complexity per Step
Standard Transformer	$O(N^2)$ Compute, $O(N)$ KV	98.2% (at high memory cost)	64.1%	$O(N)$ (KV retrieval)
Mamba / Mamba-2	$O(N)$ Linear	41.3% (degrades rapidly)	48.7%	$O(1)$
RWKV-6	$O(N)$ Linear	38.9%	46.2%	$O(1)$
Titans (MAC Variant)	$O(N)$ Sub-quadratic	99.7% (up to 2M+ tokens)	82.4%	$O(W)$ ($W = \text{Window Size}$)

In standard needle-in-a-haystack configurations, Titans achieved over 99.5% accuracy at 2 million tokens, while pure recurrent architectures degraded to near random choice once context exceeded 128,000 tokens.

Hardware Efficiency via Chunkwise Parallel Formulation

A critical question for test-time training architectures is training throughput. While sequential gradient updates at inference time are lightweight, naively executing gradient descent step-by-step during training would break GPU tensor parallelization.

Titans circumvents this by deriving a chunkwise parallel training formulation. By partitioning sequences into chunks of size $B$, the associative memory gradients and momentum updates within each chunk are unrolled into batched matrix multiplications. The recurrent updates occur exclusively across chunk boundaries, allowing Titans to train at speeds competitive with FlashAttention-2 while preserving the linear memory footprint during generation.

The Broader Shift: Inference as Continuous Learning

The implications of Titans extend far beyond raw context length benchmarks:

Zero-Cache KV Footprint for Distant History: Because historical knowledge is encoded directly into memory parameters $M_t$, systems do not need terabytes of distributed VRAM to host KV caches for long-running agent sessions.
Inherent In-Context Continual Adaptation: The boundary between fine-tuning and inference begins to dissolve. A model equipped with neural long-term memory can adapt its associative weights to entirely new codebases or document libraries during a single inference pass without altering its pre-trained backbone.

As models scale into autonomous multi-hour workflows and complex reasoning regimes, the static-weight paradigm is rapidly reaching its ceiling. Titans provides concrete mathematical proof that giving neural networks the ability to optimize their own weights during inference is not only viable—it is the most scalable path forward.