Learning to Adapt in Test-Time (Titans/MIRAS)

Introduction: The Problem Titans/MIRAS Solves

Attention lets a Transformer look back at any earlier token and decide how much it matters. That is the whole trick, and it works. The price is that the cost grows with the square of the sequence length, which is fine at 4k tokens and ruinous at the lengths you'd want for a whole codebase, a genome, or a shelf of documents.

The usual escape is to give up on attention and compress the context into a fixed-size state instead. Linear RNNs and state space models like Mamba-2 do this, and they scale linearly. The trouble is that a fixed-size state has a fixed capacity. Ask it to hold a very long sequence and something has to be thrown away. It's the same problem as summarizing a novel in one sentence: you can do it, but you shouldn't expect to answer detailed questions afterwards.

Two papers from Google Research, Titans and MIRAS, try to keep the speed of the RNN route without paying the accuracy tax. Titans is a concrete architecture. MIRAS is the theory that sits underneath it and generalizes it. What ties them together is test-time memorization: the model keeps learning while it runs, updating a long-term memory based on how surprising each new input is, with no offline retraining involved.

That is a different posture from the usual one. Instead of squeezing information into a static state, the model treats its own memory as parameters to be optimized as data streams past. New, specific details go into the memory immediately, roughly the way a person can pick up a fact mid-conversation rather than only during study sessions.

Titans Architecture - Deep Neural Memory

The Core Innovation: Deep Memory Modules

A learning system needs more than one kind of memory. The brain separates short-term from long-term storage, and Titans does something similar. Attention is kept for precise, short-range recall. For long-range memory, Titans adds a separate neural module that works differently from anything in a standard RNN.

The difference is what the memory is. In a traditional RNN the memory is a vector or a matrix. In Titans it is a deep neural network, specifically a multi-layer perceptron. An MLP can represent far richer functions than a fixed-size array, so it can summarize a large amount of input without flattening it into mush. The model isn't taking notes so much as building an understanding of what it has read.

There are three parts to the architecture. The contextual memory is the deep MLP that learns and updates as tokens are processed. The core is the ordinary transformer-like component that handles the current context with attention. The persistent memory is a set of fixed, pre-trained weights that carry general knowledge and don't change at inference time.

The contextual memory compresses past tokens into a summary, and that summary is placed into the context for attention to see. Attention can then choose, token by token, whether to consult the summary of the past or concentrate on what just arrived. The result is a hierarchy: recent information gets exact attention, distant information gets a learned summary.

Test-Time Memorization

The unusual part of Titans is test-time memorization. A normal model is frozen once training ends; it can only apply what it already learned. Titans keeps updating its long-term memory module during inference, as tokens stream in.

What it stores is not raw tokens. The memory learns representations of the input, the relationships between pieces of it, and the recurring themes that connect distant parts of the sequence. In effect the model is still learning while it is being used.

This is what makes very long contexts workable. With a two-million-token document you cannot keep everything in active memory, but you also can't afford to lose the details that matter. Test-time memorization lets the model decide what to keep, producing a compressed representation of the whole sequence that is still rich enough to answer questions from.

The Surprise Metric

Titans decides what to memorize using what the authors call the surprise metric. The idea comes from how people remember things: routine, expected events fade quickly, while things that break the pattern stick. Titans has a mathematical version of this.

Surprise is the gap between what the memory currently expects and what the new input actually says. Formally it is measured with gradients, the same error signal used in training, which tells you how far the current memory state is from accommodating the new information.

When the new word is "cat" and the memory already expects an animal, the gradient is small. That is low surprise, and the model can skip writing "cat" into long-term memory because it fits what is already there. When the memory has been summarizing a serious financial report and the next input is a picture of a banana peel, the gradient is large. That is high surprise, and it means the input is anomalous or important and should be stored.

So the gradient acts as the model's way of saying "this is unexpected and worth keeping." Only the novel, pattern-breaking inputs make it into long-term memory, which keeps the memory focused and the update process cheap.

Momentum and Forgetting Mechanisms

Raw surprise on its own has two problems, and Titans adds a mechanism for each.

The first is that surprising events are followed by unsurprising context you still want. If something startling happens in a text, the next few tokens usually explain it, and they aren't individually surprising. Momentum handles this: the model tracks both the "momentary surprise" of the current token and the "past surprise" carried over from recent tokens, so the explanation that follows a surprise gets remembered along with it.

The second is capacity. Memory is finite and sequences aren't, so something has to be discarded eventually. Titans uses an adaptive weight decay as a forgetting gate. The decay is not uniform: information judged unimportant decays quickly, information judged important decays slowly.

Together these give the memory some balance. Momentum keeps continuity so that context around important events survives; forgetting keeps the memory from filling up with everything it has ever seen.

MIRAS Framework - Theoretical Unification

Four Key Design Choices

MIRAS (Memory-Informed Robust Associative Sequence modeling) is a framework rather than a model, and its central claim is a unifying one: every successful sequence architecture, from Transformers to the fastest linear RNNs, is at bottom an associative memory module. They differ in the details of how they store and update that memory, not in kind.

Seen that way, the question stops being "which architecture is right" and becomes "how should new information be combined with old memory without losing what matters." MIRAS breaks that question into four design choices.

The memory architecture is the structure that holds information: a vector, a matrix, or a deep MLP as in Titans. The attentional bias is the internal objective the memory optimizes, which determines what it prioritizes. The retention gate is the regularizer; MIRAS reinterprets the various "forgetting mechanisms" in the literature as forms of regularization that trade new learning against retaining the past. The memory algorithm is the optimizer used to apply updates.

These four axes define a design space, and existing architectures are just points in it. Transformers, Mamba, RWKV and the rest are specific combinations of the four choices.

Memory Architecture

The memory architecture is how information is stored. Traditional models use vectors, as in basic RNNs, or matrices, as in some attention variants.

Titans uses a deep MLP instead, and the gain in expressiveness is large. A vector holds $d$ values. A matrix holds $d^2$ relationships. An MLP can learn arbitrary functions from inputs to compressed representations, and can in principle encode $2^d$ or more distinct patterns.

Depth matters here. In the ablations, deeper memory modules consistently reach lower perplexity on language modeling and hold up better as sequence length grows.

Attentional Bias

The attentional bias is the learning objective the memory optimizes when it updates. It decides what the model treats as important.

Almost every existing model uses mean squared error or dot-product similarity for this. That works on average but has known weaknesses. A single typo or anomaly can move the memory disproportionately. MSE implicitly assumes Gaussian errors, which real data often violates. And every error counts the same regardless of how much it matters.

MIRAS treats the objective as a free choice. The Titans surprise metric is one option, one that favors unexpected information. Others include robust losses such as Huber or quantile loss, information-theoretic objectives, and objectives tuned to a particular downstream task.

Retention Gate

The retention gate balances learning new things against keeping old ones. In MIRAS, forgetting is regularization.

The simple version is exponential decay, $M_{new} = \alpha \cdot M_{old} + (1-\alpha) \cdot M_{update}$, which treats every piece of information identically. MIRAS allows more selective retention: different decay rates for different kinds of information, faster forgetting for whatever is judged unimportant, and stability constraints so that updates can't knock the memory into a bad state.

Titans sits in this space with an adaptive decay that depends on importance. Knowledge that matters decays slowly; routine information decays fast.

Memory Algorithm

The memory algorithm is the optimizer that applies updates. Most models use plain gradient descent or something close to it. MIRAS opens this up too, to momentum-based updates that look at past gradients as well as the current one, to adaptive learning rates that treat different parts of the memory differently, and to second-order methods that use curvature.

Titans uses gradient descent with momentum, which is what lets it capture the context around a surprising event rather than only the event itself.

Beyond Mean Squared Error

Nearly all existing sequence models rely on MSE or dot-product similarity for both bias and retention. That makes them sensitive to outliers and limits what the memory can express.

MIRAS treats the choice of objective and regularizer as an open design question and borrows from the optimization and statistics literature to fill it in. That means non-Euclidean objectives and regularizers become available: robust losses that don't overreact to outliers, information-theoretic objectives such as maximizing mutual information or minimizing entropy, task-specific biases, and different norms and distance metrics.

This is the flexibility that produces the three MIRAS variants below. Each one is a different point in the design space.

Part 3: MIRAS Variants

Using the framework, the authors built three attention-free models, each making a different set of choices.

YAAD: Robust to Outliers

YAAD (Yet Another Attention-free Architecture with Robustness) is built to shrug off outliers, such as a single typo in a long document. It swaps MSE for Huber loss, which penalizes big mistakes more gently so that one odd input doesn't dominate the memory update.

Key Innovation: Instead of MSE, YAAD uses Huber loss:

Inside the threshold $\delta$ this is just MSE. Outside it, the penalty grows linearly rather than quadratically, so outliers stop having outsized influence.

That's the behavior you want when the input is messy or inconsistent, which describes most real data at scale.

MONETA: Generalized Norms

MONETA asks what happens if you use stricter, more general penalties than MSE, and applies them to both sides of the problem: what the model attends to and what it forgets.

Key Innovation: MONETA uses $p$-norms and other generalized distance metrics:

Changing $p$ changes the geometry. At $p=2$ you have ordinary Euclidean distance, equivalent to MSE. At $p=1$ you have Manhattan distance, which is more tolerant of outliers. As $p \to \infty$ you approach Chebyshev distance, which only cares about the largest error.

MONETA applies these norms to the attentional bias and the retention gate together, which gives it a more disciplined memory than the MSE default.

MEMORA: Probability Map Constraints

MEMORA goes after stability. It forces the memory to behave like a probability distribution, so that every update is controlled and can't produce a nonsensical state.

Key Innovation: MEMORA constrains memory updates to maintain probability distribution properties. Memory values stay non-negative, the memory state stays normalized (summing or integrating to 1), and updates preserve ordering relationships between entries.

The payoff is a clean, predictable update rule. Because the memory is always a valid distribution over states, it's also easier to interpret.

Part 4: Experimental Results

Language Modeling Performance

Titans and the MIRAS variants were compared against Transformer++, Mamba-2 and Gated DeltaNet on standard language modeling datasets (C4, WikiText) and zero-shot reasoning tasks (HellaSwag, PIQA). Across the board they reached higher accuracy on the downstream tasks and lower perplexity on the language modeling ones, while keeping training parallelizable and inference linear in sequence length.

The three new variants, MONETA, YAAD and MEMORA, also beat the baselines, which is evidence that moving away from MSE is worth doing rather than just theoretically interesting.

The Power of Deep Memory

The ablations on memory depth are the clearest result in the paper. Holding memory size fixed and varying depth, deeper memory MLPs reach lower perplexity, hold their performance better as sequence length increases, and do so at both the 360M and 760M parameter scales.

That supports the core design decision. A deep network as memory really does hold more than a fixed-size vector or matrix of the same size.

Extreme Long-Context Recall

The headline result is on very long contexts, measured with BABILong, a benchmark that requires reasoning over facts scattered through extremely long documents.

Titans beats every baseline here, including models as large as GPT-4, with far fewer parameters. It scales past 2M tokens, and its accuracy doesn't fall off as the context gets longer.

Recall over a two-million-token context is what you need for reading an entire book, legal filing or codebase in one pass, for analyzing whole genomes, and for holding onto context across a very long conversation or analysis session.

Efficiency Comparisons

None of this costs the efficiency that motivated the design. Inference is $O(n)$ in sequence length, the same as an RNN. Training still parallelizes, unlike a sequential RNN. And the deep memory module is compact next to a full attention matrix.

In short: Transformer-level accuracy at RNN-level cost.

Part 5: Technical Deep Dive

Mathematical Formulations

The core of Titans' memory update mechanism can be formalized as follows. Let $M_t$ be the memory state at time $t$, and let $x_t$ be the new input token.

The surprise metric is computed as the gradient of the loss with respect to the memory:

where $L$ is the loss function comparing the model's prediction with the expected output.

The memory update incorporates surprise, momentum, and forgetting:

Here $\lambda_t$ is the adaptive learning rate driven by surprise, $f_\theta$ is the deep memory network with parameters $\theta$, and the $(1 - \lambda_t)$ term is the adaptive weight decay that does the forgetting.

The momentum mechanism considers recent context:

This is what keeps a surprising event and the tokens right after it together in memory.

Architecture Overview

The Titans architecture can be conceptualized as:

Input Sequence → [Contextual Memory (Learning)] → Summary
                                      ↓
                    [Core (In-Context Learning)] → Attention
                                      ↓
                    [Persistent Memory (Fixed)] → Output

The contextual memory compresses past tokens into a summary representation, which is then:

1. Incorporated into the current context
2. Passed to the attention mechanism
3. Used alongside recent tokens for prediction

At each step attention can weight recent tokens for local precision, the memory summary for compressed global context, or a mix of both. That is the hierarchy again: precision where it's cheap, summary where it isn't.

Code Examples

Here's a simplified implementation of the Titans memory update mechanism:

import torch
import torch.nn as nn
 
class TitansMemory(nn.Module):
    def __init__(self, d_model: int, memory_dim: int, num_layers: int = 3):
        super().__init__()
        self.d_model = d_model
        self.memory_dim = memory_dim
        
        # Deep neural network memory module
        layers = []
        layers.append(nn.Linear(d_model, memory_dim))
        for _ in range(num_layers - 2):
            layers.append(nn.Linear(memory_dim, memory_dim))
            layers.append(nn.ReLU())
        layers.append(nn.Linear(memory_dim, memory_dim))
        self.memory_net = nn.Sequential(*layers)
        
        # Surprise threshold
        self.surprise_threshold = 0.1
        self.momentum_alpha = 0.9
        
    def forward(self, x: torch.Tensor, memory_state: torch.Tensor):
        """
        x: [batch, d_model] - current token embedding
        memory_state: [batch, memory_dim] - current memory state
        """
        # Compute surprise (gradient magnitude)
        x_requires_grad = x.requires_grad_(True)
        memory_requires_grad = memory_state.requires_grad_(True)
        
        # Forward through memory network
        memory_update = self.memory_net(x_requires_grad)
        
        # Compute loss (simplified - in practice, this would be 
        # the actual model loss)
        loss = torch.mean((memory_update - memory_state) ** 2)
        
        # Compute surprise as gradient magnitude
        grad = torch.autograd.grad(
            loss, memory_requires_grad, 
            create_graph=True, retain_graph=True
        )[0]
        surprise = torch.norm(grad, dim=-1)
        
        # Adaptive learning rate based on surprise
        lambda_t = torch.sigmoid(surprise - self.surprise_threshold)
        
        # Update memory with momentum
        if not hasattr(self, 'prev_surprise'):
            self.prev_surprise = surprise
        surprise_momentum = (
            self.momentum_alpha * surprise + 
            (1 - self.momentum_alpha) * self.prev_surprise
        )
        self.prev_surprise = surprise
        
        # Apply adaptive forgetting
        lambda_adaptive = lambda_t * (1 + surprise_momentum)
        lambda_adaptive = torch.clamp(lambda_adaptive, 0, 1)
        
        # Update memory state
        new_memory = (
            (1 - lambda_adaptive) * memory_state + 
            lambda_adaptive * memory_update
        )
        
        return new_memory, surprise

The pieces to notice are the MLP used as memory, the surprise computed from a gradient norm, the learning rate that scales with surprise, the momentum term, and the adaptive decay that does the forgetting.

Putting the design together, Titans gets an expressive memory (a deep network can represent far more patterns than a fixed-size vector), a memory that keeps learning at test time, a filter that stores only what is surprising, linear-time inference with parallel training, and demonstrated recall past 2M tokens.

It pays for this in complexity. A deep memory module has more parameters than a vector, training has to account for the test-time updates, and the memory network itself takes more space than a fixed-size state. Given the long-context results, that seems like a good trade.

Implications and Future Directions

Titans and MIRAS replace the fixed recurrent state with a deep network that learns to memorize as data streams in, and in doing so get around the capacity ceiling that has limited linear models.

A few things follow from that. MIRAS gives a single lens for understanding sequence models: they are all associative memories, differing in four design choices, which makes new architectures easier to design and compare. Leaving MSE behind opens up robust, task-specific and information-theoretic objectives that were previously off the table. Test-time memorization means models that keep learning from their inputs without being retrained. And usable 2M-token contexts change what's practical in document understanding, genomics and long-horizon reasoning.

There is also a lot left to try. Titans-style memory could be combined with other efficient attention schemes such as Infini-Attention or Ring Attention. Attentional biases could be designed for specific tasks, code or scientific reasoning for instance. The memory could be extended to images, audio and other modalities. The expressivity and limits of deep memory modules deserve proper theoretical treatment. The update mechanism can probably be made faster still. And there is more to explore in robust losses and retention rules for noisy real-world data.

The broader point is that RNN-level efficiency and Transformer-level expressiveness may not be a trade-off after all. If models are going to read whole codebases and genomes and hold long conversations, something like Titans, built on something like MIRAS, is a plausible way to get there.

References

1. Behrouz, A., Razaviyayn, M., & Mirrokni, V. (2024). Titans: Test-Time Adaptation for Long-Context Language Models via In-Context Memorization. arXiv preprint. arXiv:2504.13173

2. Behrouz, A., Razaviyayn, M., & Mirrokni, V. (2024). MIRAS: Memory-Informed Robust Associative Sequence Modeling. arXiv preprint. arXiv:2501.00663

3. Google Research. (2024). Titans + MIRAS: Helping AI have long-term memory. Google Research Blog. https://research.google/blog/titans-miras-helping-ai-have-long-term-memory/

4. Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint. arXiv:2312.00752

5. Peng, B., Alcaide, E., Anthony, Q., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. EMNLP. arXiv:2305.13048

6. Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762