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

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.

Machine learning models often need to use long documents, code, or conversation histories, but keeping all of that text available for every question is expensive. Compressing the text into a smaller memory can help, but it is hard to do without losing important details. We study whether a model can read a context once and store the useful information in a small memory that can later replace the original text. We introduce GradMem, a method that teaches models to write into this memory using a few small learning steps at the moment the context is read. The model itself stays unchanged; only the memory is updated for the current context. Users can choose how many writing steps to use, trading more computation for a better memory. In controlled memory tests, GradMem stored more information than comparable methods that write memory in a single forward pass. We also show that the same idea works with pretrained language models on tasks such as question answering and language modeling. GradMem points toward AI systems that can read a document once, store its key details in a small learned memory, and answer multiple later questions without rereading the full text; when more detail is needed, users can spend extra writing steps to improve the memory.