Memory-Augmented Test-Time Training: Integrating External Memories for Longer Contexts

TL;DR: What if you could combine the strengths of TTT-E2E with retrieval-augmented or external memory modules? This hybrid could allow the model to not only compress context into its weights but also leverage persistent, scalable memory.

Research Question: How does augmenting TTT-E2E with a learnable external memory module impact long-term context retention, reasoning, and scaling properties?

Hypothesis: A TTT-E2E model equipped with external memory will outperform pure weight-based TTT on tasks requiring retrieval or reasoning over information far outside the model’s current window, and will provide better scaling with context length.

Experiment Plan: - Implement TTT-E2E with an attached differentiable memory module (e.g., a memory matrix or retrieval component).

• During test time, allow both weight updates (TTT) and memory reads/writes, using meta-learning to coordinate the two.
• Evaluate on synthetic and real datasets requiring long-term recall (e.g., long-form QA, reasoning over book-length documents).
• Analyze performance, scaling behavior, and memory usage compared to standard TTT-E2E and retrieval-augmented transformers.

References:

• Tandon, A., Dalal, K., Li, X., Koceja, D., Rod, M., et al. (2025). End-to-End Test-Time Training for Long Context.
• Hardt, M., & Sun, Y. (2023). Test-Time Training on Nearest Neighbors for Large Language Models. International Conference on Learning Representations.
• Chen, Z., Romanou, A., Weiss, G., & Bosselut, A. (2025). PERK: Long-Context Reasoning as Parameter-Efficient Test-Time Learning. arXiv.org.