Chunkwise Test-Time Memorization
- Chunkwise test-time memorization is a technique that segments long input sequences into manageable chunks to enable efficient, memory-bound inference.
- It employs methods such as test-time optimization, explicit memory tokens, and compressor modules to store and adapt parameters on-the-fly.
- The approach provides linear compute scaling, supports online adaptation, and boosts performance in tasks like language modeling, video generation, and long-context reasoning.
Chunkwise test-time memorization refers to a set of algorithms and frameworks that enable neural sequence models to efficiently process long or streamed inputs by compressing, optimizing, or storing representations or parameters for individual, temporally contiguous "chunks" of data during inference. This process allows models—ranging from Transformers and RNNs to specialized compressive memory or meta-learning architectures—to extend context lengths, adapt online, and support fast recall for subsequent queries, even when direct access to the raw context is no longer feasible. Chunkwise test-time memorization is foundational in contemporary research on model scalability, adaptive computation, and memory-efficient inference.
1. Principles of Chunkwise Test-Time Memorization
Chunkwise test-time memorization divides long input sequences into discrete segments ("chunks"), allowing the model to process, store, and later recall information in fixed-size or streaming units. This approach enables several core objectives:
- Linear or sublinear inference cost: By batching computation over constant-length segments, models maintain memory and compute growth independent of overall context length.
- Online adaptation and memory updates: Models can conduct local, on-the-fly parameter or state optimization per chunk, enhancing performance for evolving or compositional queries.
- Cross-chunk memory and retrieval: Mechanisms to store, update, and retrieve per-chunk representations or adapted parameters are central to reducing redundancy and enabling generalization across similar queries.
Approaches differ in how they realize memorization: as parameter updates (e.g., test-time optimization), as explicit memory tokens, via compressor modules, or by chunkwise solutions of in-context objectives. The enabling assumption is that context segmentation and efficient per-chunk memory access can collectively approximate or surpass the fidelity of naive, full-sequence processing.
2. Architectures and Algorithms for Chunkwise Memorization
The design and instantiation of chunkwise test-time memorization architectures varies according to model family and domain:
Transformer-based Models
- ETT (Extend at Test-Time) (Zahirnia et al., 8 Jul 2025):
- Segments the input context into overlapping chunks of length and overlap .
- Each chunk is used to fine-tune a subset of model parameters (notably weights) during inference, with the optimization objective being log-likelihood over all tokens in each chunk.
- Overlap ensures information continuity near chunk borders and improves context retention.
- The algorithm achieves linear compute and constant memory with respect to total sequence length .
Compressive and Prefix Memory Models
- GradMem (Kuratov et al., 14 Mar 2026):
- Uses a small set of prefix memory vectors to summarize arbitrarily long contexts.
- At test time, for each context or chunk, is optimized by gradient descent on a self-supervised reconstruction loss that forces to encode all contextual information.
- Memory prefix can be read repeatedly for many queries, amortizing the cost of chunkwise optimization.
- LycheeMemory (Chen et al., 9 Feb 2026):
- Splits the input into fixed-size, non-overlapping chunks (0 in experiments).
- Each chunk is compressed into a set of learned "memory tokens" via a LoRA-augmented compressor; these are stored as a compressed memory bank.
- A gating module, attached via LoRA, selects relevant chunks for retrieval based on query and working memory state.
- The reasoning module iteratively updates a working memory buffer by combining retrieved chunk-memories and current hypotheses.
RNNs with Test-Time Memorization
- MesaNet (Oswald et al., 5 Jun 2025):
- Implements chunkwise optimization for fast-weight linear maps 1 that satisfy a least-squares fit to historical key-value pairs, regularized via a positive-definite matrix 2 and decayed by 3.
- Solves, per chunk, the linear system 4 using the conjugate gradient (CG) method, exploiting chunk parallelism for speed and constant memory.
- Learning and recall are online, with only summary states 5 carried across chunk boundaries.
- TNT (Li et al., 10 Nov 2025):
- Addresses the hardware–fidelity trade-off in chunkwise RNN memory by pre-training with large, hardware-efficient chunks (via global memories) and short, parallel local shards (via local memories periodically reset).
- Final accuracy is restored by a brief fine-tuning phase with high-resolution (small-chunk) updates, bridging the mismatch between hardware-efficient training and chunkwise test-time inference.
- TTOM for Video Generation (Qu et al., 9 Oct 2025):
- Maintains a fixed-capacity parametric memory for video prompts, enabling per-chunk optimization of lightweight parameters (such as LoRA weights).
- Memory supports insert, read, update, and delete operations, allowing adaptation and reuse of optimization contexts tied to similar prompts.
- The approach aligns generated video layout and attention with planned spatiotemporal layouts, using chunkwise prompt segmentation as a natural unit for memorization and update.
3. Memory Operations and Objectives
Across chunkwise memorization schemes, several core mechanisms are prominent:
| Operation | Description | Typical Use |
|---|---|---|
| Insert | Store a new chunk's representation or parameters | Add new memory after optimization |
| Read (Load) | Retrieve memory for a given chunk or prompt | Initialize state for inference |
| Update | Refine stored memory after further optimization | Improve upon retrieval |
| Delete | Evict least-used/oldest entries from fixed-capacity | Maintain bounded memory |
For test-time optimization-based methods, a local loss is minimized per chunk. For TTOM (Qu et al., 9 Oct 2025), the loss is a Jensen–Shannon divergence 6 between attention and layout masks. For GradMem, it is autoregressive reconstruction loss over the context when prefixed by the chunk memory.
Table: Representative Memory Operations
| Model | Insert | Read | Update | Delete |
|---|---|---|---|---|
| TTOM | Yes | Yes | Yes | Yes |
| GradMem | Yes | Yes | Yes | — |
| LycheeMemory | Yes | Yes | — | — |
The interplay of these operations enables test-time adaptation, efficient recall, and bounded memory growth in streaming or long-context scenarios.
4. Computational Complexity and Scaling Properties
The primary advantage of chunkwise test-time memorization is the decoupling of compute and memory requirements from total sequence length:
- Linear/constant scaling: For models such as ETT (Zahirnia et al., 8 Jul 2025), per-chunk operations maintain 7 compute and memory per chunk, yielding total compute 8 and constant per-chunk memory—unlike vanilla Transformers' quadratic scaling.
- Parallelization trade-offs: TNT (Li et al., 10 Nov 2025) shows that chunkwise processing can massively increase throughput by breaking sequential barriers, especially when chunk sizes are hardware-aligned and local memory resets enable full parallelism within shards.
- Chunk size vs. fidelity trade-off: Larger chunks improve throughput but can degrade memory fidelity due to coarser approximation. Resolvable via hierarchical or two-stage training procedures that reconcile pre-training efficiency with fine-grained test-time adaptation.
Empirical scaling ablations show, for instance, that GradMem's capacity increases monotonically with both memory token count and number of optimization steps (e.g., 100% exact match at 16 key-value pairs with 9, compared to baseline RMT's limit at 8 pairs) (Kuratov et al., 14 Mar 2026). LycheeMemory further scales to 1.75M-token contexts with a 2× memory reduction and 6× speedup over MemAgent (Chen et al., 9 Feb 2026).
5. Empirical Performance and Applications
Chunkwise memorization mechanisms have demonstrated robust gains in various tasks:
- Video Generation: TTOM achieves 14% improvement in compositional alignment over test-time optimization alone on the Motion category of T2V-CompBench; similar gains observed for numeracy and spatial relations (Qu et al., 9 Oct 2025).
- Language Modeling:
- MesaNet: Matches softmax transformer perplexity at 2048 tokens, succeeds on long (32K) context benchmarks with constant recurrent state and chunk-parallel CG computation (Oswald et al., 5 Jun 2025).
- ETT: Extends GPT-Large and Phi-2 context length by up to 32× and 16×, achieving 30% improvement in LongBench scores with only selective layer adaptation (Zahirnia et al., 8 Jul 2025).
- GradMem: Outperforms forward-only memory writers (RMT, ARMT) at fixed memory size on both synthetic and real-world NLP tasks (e.g., SQuAD, bAbI) (Kuratov et al., 14 Mar 2026).
- TNT: Achieves up to 17× faster training and improved perplexity (PPL = 23.09) compared to standard Titan/TTT baselines, closing the performance gap with Transformers (Li et al., 10 Nov 2025).
- Long-Context Question Answering and Reasoning: LycheeMemory supports reasoning over ultra-long inputs (up to 1.75M tokens), substantially reducing GPU and wall-clock requirements compared to full-attention or retrieval-augmented baselines (Chen et al., 9 Feb 2026).
6. Limitations and Future Directions
Current chunkwise test-time memorization approaches face several bounded challenges:
- Test-time compute cost: Iterative optimization and per-chunk updates (e.g., GradMem's K-step descent, MesaNet's CG iterations) introduce extra inference-time flops compared to one-pass attention or memoryless models (Oswald et al., 5 Jun 2025, Kuratov et al., 14 Mar 2026).
- Training–inference chunk mismatch: Pre-training on large chunks may require fine-tuning or architectural adjustment to recover resolution and accuracy at test time (as in TNT (Li et al., 10 Nov 2025)).
- Memory capacity and eviction: Fixed-capacity parametric memories necessitate careful management of insert, update, and delete policies to prevent performance degradation as context or task diversity scales (Qu et al., 9 Oct 2025).
- Design choices for memory compression: Choices such as overlap size, chunk granularity, and selective parameter adaptation have a substantial impact on trade-offs between recall fidelity, latency, and hardware efficiency (Zahirnia et al., 8 Jul 2025, Chen et al., 9 Feb 2026).
Potential future directions include the development of custom accelerator support for chunkwise kernels (noted for TNT), dynamic per-chunk adaptation strategies balancing compute and reconstruction error, deeper local memory hierarchies, and hybridization with attention for continued gains on extreme long-context and compositional tasks (Li et al., 10 Nov 2025, Oswald et al., 5 Jun 2025, Chen et al., 9 Feb 2026). Evidence suggests that chunkwise test-time memorization is foundational for both theoretical insights into online memory compression and practical advances in large-scale, streaming, and multimodal inference.