Deep Test-Time Memorization Modules

Updated 12 November 2025

Deep test-time memorization modules are advanced neural architectures that adapt memory during inference by dynamically updating stored representations based on contextual input.
They leverage a range of designs including deep MLPs, matrix memories, and key-value stores to efficiently manage long-sequence dependencies and optimize resource use.
These modules have practical applications in language modeling, video synthesis, and symbolic regression, achieving improved accuracy and compute efficiency.

Deep test-time memorization modules are a class of neural architectures and mechanisms designed to adapt or optimize portions of a model’s memory or parameters at inference, with the explicit goal of leveraging contextual input and memorizing relevant information online. Unlike classical models with static parameters at inference, these modules enable models to store, recall, and update information dynamically in response to data encountered at test time, achieving a form of associative, context-aware, or even “lifelong” memory. This computational paradigm underpins a wide range of new architectures—including fast-weight Transformers, deep recurrent memories, and parametric adapter stores—with core applications in sequence modeling, reasoning, long-context language modeling, video synthesis, symbolic regression, and compute-efficient inference.

1. Foundational Principles and Taxonomy

Deep test-time memorization (TTM) modules are characterized by a separation between persistent parameters (trained offline) and stateful components (adapted online during inference). The fundamental components and choices are as follows (Behrouz et al., 17 Apr 2025):

Associative Memory Architecture: Typically instantiated as a matrix memory (e.g., Hebbian or delta-rule), a deep MLP, or a non-parametric key-value cache.
Attentional Bias Objective: The loss used to guide memory adaptation—commonly dot-product similarity or ℓ₂ regression, but extended to ℓ_p, Huber, or f-divergence-based objectives.
Retention Gate (Forgetting Mechanism): Local (stepwise) or global (across-step) regularization to counteract memory overflow or catastrophic forgetting; includes weight decay, q-norms, KL/entropy, and elastic-net penalties.
Inner-loop Memory Optimizer: The algorithm for updating the memory (e.g., gradient descent, momentum, Muon/second-order methods, mirror descent, online Newton–Schulz updates).

These modules are instantiated per inference scenario, with architectures varying from large-chunk parallel adaptation in Transformers (Zhang et al., 29 May 2025), deep RNN memory modules in Titans (Behrouz et al., 31 Dec 2024), parametric key-value stores for adapters in video diffusion (Qu et al., 9 Oct 2025), to memorization-efficient feed-forward updates in LLMs (Zhu et al., 23 Jun 2024). Miras (Behrouz et al., 17 Apr 2025) formalizes the design space by treating the memory architecture, bias objective, retention, and learning rules as axes for architecture search.

2. Memory Architectures and Update Mechanisms

Architectures span a spectrum of complexity and adaptability:

Matrix/Linear Memory: These maintain and update a matrix $M_t$ , with updates of the form $M_t = \alpha_t(I - \eta_t k_t k_t^\top) M_{t-1} + v_t k_t^\top$ . Used in linear attention, DeltaNet, and RWKV. Capacity is $O(d)$ but suffers early saturation on long contexts.
Deep MLP Memory: A multi-layer perceptron $M_W(k)$ parameterized by $W$ , which is adapted at each time step via a surprise or regression objective. Deeper MLPs substantially increase the functional/memory capacity, enabling Titans, ATLAS, and Moneta/Yaad/Memora models to scale to millions of tokens and beyond (Behrouz et al., 31 Dec 2024 Behrouz et al., 29 May 2025 Behrouz et al., 17 Apr 2025).
Chunkwise and Large-Chunk Parallelization: Early TTT methods used small chunks (e.g., 16 or 64 tokens), which severely limited hardware utilization (FLOPs/bytes ratio $r \ll 1$ ) and scalability. LaCT (Zhang et al., 29 May 2025) demonstrates that raising the chunk to $2 \,\mathrm{K}$ – $1\,\mathrm{M}$ tokens brings GPU utilization to $40$– $70\%$ of peak and enables state sizes up to $40\%$ of model parameters.
Optimizers: Momentum-based updates, second-order Muon updates (utilizing Newton–Schulz orthogonalization), and mirror descent on the simplex (e.g., via row-wise softmax) are employed for stable, locally-optimal adaptation.

A generic per-step update is:

$W_t = W_{t-1} - \eta_t \nabla_W\,\ell(W_{t-1};k_t,v_t)$

or, for chunkwise updates,

$W \leftarrow \text{L}_2\text{-Normalize}( W - \sum_{i=1}^b \eta_i\,\nabla_W \ell(f_W(k_i), v_i) )$

with higher-order updates replacing the gradient step.

3. Applications and Empirical Advances

TTM modules have been validated across modalities and tasks:

Domain / Task	Memory Module	Capacity	Context Length	Notable Results
Language Modeling	Deep MLP (Titans, ATLAS, LaCT)	Up to $40\%$ params	$10^7$ tokens	ATLAS achieves $80\%$ accuracy @ $10$M on BABILong (Behrouz et al., 29 May 2025)
Video Diffusion/Synthesis	Adapter Key-Value Store (TTOM)	Per-prompt LoRA head	$56$K tokens	TTOM + $34.5\%$ (CogVideo-5B) on T2V-CompBench (Qu et al., 9 Oct 2025)
Recall/Needle-in-a-Haystack	Deep MLP (MAC/MAG–Titans)	Linear/MLP	$2$M tokens	Titans-MAC reaches $98–99\%$ needle rate at $16$K (Behrouz et al., 31 Dec 2024)
Genomics, Time Series	Deep MLP	Modal-specific	$>$ 1M tokens	Titans surpass Transformer++ on genomics (Behrouz et al., 31 Dec 2024)
Memorization-based Inference	Non-parametric table (MBI)	Table lookup	N/A	$2.7\times$ more energy-eff. than MLP-CIM (Giacomini et al., 2023)

In all cases, scaling memory depth/width and optimizing the update mechanism improves performance on long-context or recall-intensive tasks.

4. Parallelization and Resource Efficiency

The shift to large-chunk updates (Zhang et al., 29 May 2025) and hierarchical memory (Li et al., 10 Nov 2025) has enabled near-linear resource scaling:

Parallelism: LaCT and TNT demonstrate that chunkwise or hierarchical strategies (global + local memories, periodic resets) break sequential dependencies, enabling high-throughput hardware utilization and $O(1)$ steps per chunk or shard.
Chunk Size vs. Performance Trade-off: Small chunks yield fine-grained adaptation but low throughput; large chunks accelerate computation but may degrade per-token adaptability. Two-stage processes (e.g., TNT) resolve this via efficiency-focused pretraining (large chunk) and accuracy-focused fine-tuning (small chunk).
Resource Scaling: State size per block can reach $40\%$ of total parameters, with empirical scaling demonstrating steady gains as capacity and chunk size increase (Zhang et al., 29 May 2025). Memory growth is controlled via gating and low-rank/sparse parameterizations.

5. Design Variants and Theoretical Foundations

The design space encompasses multiple axes:

Variant Spectrum: Miras (Behrouz et al., 17 Apr 2025) systematically studies models ranging from dot-product similarity (Hebbian), $\ell_2$ regression (Delta), to $\ell_p$ , Huber, robust, and KL/entropy retention. Moneta achieves best perplexity via $\ell_3/\ell_4$ ; Yaad stabilizes via Huber; Memora utilizes KL+entropy for robust retention on the simplex.
Expressivity and Capacity: Deep MLPs (Titans, ATLAS, Moneta) provide $O(d^2)$ or even super-linear capacity (via polynomial feature maps), surpassing the $O(d)$ bottleneck of linear recurrent models. The inclusion of learned retention gates (momentum, step size, forget factors) is essential to avoid memory overflow and to enable attentive gating to prune context.
Theoretical Equivalence: ATLAS demonstrates that, with appropriate choices of feature map and optimizer (Omega rule + Muon), deep test-time memorization strictly generalizes full and local-window softmax attention.
Limitations: Linear attention and SwiGLU fast-weights lack $Q$ – $K$ rotation invariance, a property of softmax attention, with unexplored implications for certain tasks (Zhang et al., 29 May 2025).

6. Practical Applications and Algorithmic Implementations

Canonical application patterns include:

LLMs: At every new sequence or prompt, memory modules (fast weights) are adapted in-place via gradient or momentum-based updates, with chunk size determined by budget-latency trade-offs. FastMem (Zhu et al., 23 Jun 2024) demonstrates prompt-specific fine-tuning by updating only the last-layer FFN, reducing perplexity and error rates on context-sensitive QA and summarization.
Video Generation: TTOM (Qu et al., 9 Oct 2025) maintains a per-prompt key-value store, storing and retrieving LoRA adapters (values) via hashed prompt keys. At each prompt, the alignment loss between attention maps and prescribed spatial layouts is minimized by updating only the adapter, allowing compositional generalization and precise attention steering.
Symbolic Regression: Test-time strategies (e.g., prompting with verified subtrees or using MCTS) reduce memorization bias of Transformers but reveal that lowering memorization bias does not guarantee improved numerical fit (Sato et al., 28 May 2025).

Example pseudocode (LaCT layer, per-chunk):

for i in range(b):
    o[i] = fast_weight_net.forward(q[i])
grad = sum([eta[i] * grad_loss(fast_weight_net(fast_weight_net.k[i]), v[i]) for i in range(b)])
W = l2_normalize(W - grad)

7. Limitations and Open Challenges

Key unresolved issues and research directions:

Capacity Envelope: Empirical gains flatten beyond $\sim$ 40% parameter state; developing compression or sparse-state memories may extend this further (Zhang et al., 29 May 2025).
Kernel Optimizations: Most implementations use native PyTorch; maximal hardware utilization will require custom CUDA or Triton kernels, especially for extreme state sizes or batch parallelism (Zhang et al., 29 May 2025 Li et al., 10 Nov 2025).
Rotation Invariance and Robustness: The lack of $Q$ – $K$ rotation invariance remains an open concern for the class of fast-weight modules. Its practical impact has yet to be quantified.
Task Scope: While validated across LM, video, image, and recall tasks, reasoning and 3D unposed tasks are less explored. The adaptability of these modules to broader reasoning or symbolic tasks is an ongoing area of paper.
Optimizer Innovation: Advances such as adaptive chunking, meta-learned optimizers, and further exploration of second-order or attention-compatible mechanisms (e.g., Muon) are expected to yield further improvements.
Bias and Generalization: In symbolic regression, reduction in memorization bias does not necessarily imply improved test accuracy, highlighting a gap between compositional generalization and numerical fit (Sato et al., 28 May 2025).

References

Test-Time Training Done Right (Zhang et al., 29 May 2025)
TTOM: Test-Time Optimization and Memorization for Compositional Video Generation (Qu et al., 9 Oct 2025)
TNT: Improving Chunkwise Training for Test-Time Memorization (Li et al., 10 Nov 2025)
Titans: Learning to Memorize at Test Time (Behrouz et al., 31 Dec 2024)
Towards Model-Size Agnostic, Compute-Free, Memorization-based Inference of Deep Learning (Giacomini et al., 2023)
ATLAS: Learning to Optimally Memorize the Context at Test Time (Behrouz et al., 29 May 2025)
FastMem: Fast Memorization of Prompt Improves Context Awareness of LLMs (Zhu et al., 23 Jun 2024)
It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization (Behrouz et al., 17 Apr 2025)
Can Test-time Computation Mitigate Memorization Bias in Neural Symbolic Regression? (Sato et al., 28 May 2025)