Search papers, labs, and topics across Lattice.
The paper introduces GradMem, a method for writing context into a compact memory representation using test-time gradient descent, enabling language models to condition on long contexts without storing large KV-caches. GradMem optimizes a self-supervised context reconstruction loss by performing gradient descent on a small set of prefix memory tokens while freezing model weights. Experiments on associative retrieval and natural language tasks (bAbI, SQuAD) demonstrate that GradMem outperforms forward-only memory writers and effectively scales capacity with additional gradient steps.
Forget KV-caches: GradMem lets you compress long-context knowledge into a tiny memory bank using test-time gradients, rivaling standard methods on question answering.
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.
