# GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent

arXiv: https://arxiv.org/abs/2603.13875 · HTML: https://mlsanity.com/paper/2603.13875/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 49 references that are in the map.
- First posted on arXiv 14 Mar 2026.

## Builds on (most cited first, 30 of 41)

- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [The Power of Scale for Parameter-Efficient Prompt Tuning](https://mlsanity.com/paper/2104.08691/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [Distributed Representations of Sentences and Documents](https://mlsanity.com/paper/1405.4053/) (2014)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [RepoCoder: Repository-Level Code Completion Through Iterative Retrieval and Generation](https://mlsanity.com/paper/2303.12570/) (2023)
- [Adapting Language Models to Compress Contexts](https://mlsanity.com/paper/2305.14788/) (2023)
- [Learning to (Learn at Test Time): RNNs with Expressive Hidden States](https://mlsanity.com/paper/2407.04620/) (2024)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [Skip-Thought Vectors](https://mlsanity.com/paper/1506.06726/) (2015)
- [Towards AI-Complete Question Answering: A Set of Prerequisite Toy Tasks](https://mlsanity.com/paper/1502.05698/) (2015)
- [Recurrent Memory Transformer](https://mlsanity.com/paper/2207.06881/) (2022)
- [In-context Autoencoder for Context Compression in a Large Language Model](https://mlsanity.com/paper/2307.06945/) (2023)
- [Test-Time Training Done Right](https://mlsanity.com/paper/2505.23884/) (2025)
- [Neural Variational Inference for Text Processing](https://mlsanity.com/paper/1511.06038/) (2015)
- [Cartridges: Lightweight and general-purpose long context representations via self-study](https://mlsanity.com/paper/2506.06266/) (2025)

## Built on by (most cited first, 9 of 9)

- [Test-Time Learning with an Evolving Library](https://mlsanity.com/paper/2605.14477/) (2026)
- [Evidence-Informed LLM Beliefs for Continual Scientific Discovery](https://mlsanity.com/paper/2606.29182/) (2026)
- [Proteus: Incremental Memory Activation for Long-Context Sequence Modeling](https://mlsanity.com/paper/2608.16844/) (2026)
- [RPMem: Learning Long-Term Recurrent Parametric Memory Across Sessions for LLM Agents](https://mlsanity.com/paper/2609.23466/) (2026)
- [FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention](https://mlsanity.com/paper/2609.24089/) (2026)
- [ProTTT: Learning to Learn Semantic User Memory with Test-Time Training](https://mlsanity.com/paper/2609.32564/) (2026)
- [Towards In-Parameter Memory Augmentation for Large Language Models](https://mlsanity.com/paper/2610.08630/) (2026)
- [Encoding Context into LLM Weights: What Makes It Context?](https://mlsanity.com/paper/2610.17686/) (2026)
- [Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules](https://mlsanity.com/paper/2610.19210/) (2026)

## Abstract (arXiv)

Many large language model applications require conditioning on long contexts. Transformers typically support this by storing a large per-layer KV-cache of past activations, which incurs substantial memory overhead. A desirable alternative is compressive memory: read a context once, store it in a compact state, and answer many queries from that state. We study this in a context removal setting, where the model must generate an answer without access to the original context at inference time. We introduce GradMem, which writes context into memory via per-sample test-time optimization. Given a context, GradMem performs a few steps of gradient descent on a small set of prefix memory tokens while keeping model weights frozen. GradMem explicitly optimizes a model-level self-supervised context reconstruction loss, resulting in a loss-driven write operation with iterative error correction, unlike forward-only methods. On associative key--value retrieval, GradMem outperforms forward-only memory writers with the same memory size, and additional gradient steps scale capacity much more effectively than repeated forward writes. We further show that GradMem transfers beyond synthetic benchmarks: with pretrained language models, it attains competitive results on natural language tasks including bAbI and SQuAD variants, relying only on information encoded in memory.