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

TL;DR

GradMem introduces a test-time gradient descent method to write context into memory for large language models, enabling efficient context reuse and improved retrieval without increasing memory overhead.

Contribution

It proposes a novel gradient-based memory writing technique that enhances context retention and retrieval in language models during inference.

Findings

GradMem outperforms forward-only memory methods at the same memory size.

Additional gradient steps significantly increase memory capacity.

Achieves competitive results on natural language tasks like bAbI and SQuAD.

Abstract

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 ompressive 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.