GradMem: Context in Geheugen Schrijven met Test-Tijd Gradiëntendaling
GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent
March 14, 2026
Auteurs: Yuri Kuratov, Matvey Kairov, Aydar Bulatov, Ivan Rodkin, Mikhail Burtsev
cs.AI
Samenvatting
Veel toepassingen van grote taalmodellen vereisen conditionering op lange contexten. Transformers ondersteunen dit doorgaans door een grote per-laag KV-cache van eerdere activeringen op te slaan, wat aanzienlijke geheugenoverhead met zich meebrengt. Een wenselijk alternatief is compressief geheugen: lees een context eenmaal, sla deze op in een compacte staat, en beantwoord vele queries vanuit die staat. Wij bestuderen dit in een contextverwijderingsscenario, waarbij het model een antwoord moet genereren zonder toegang tot de oorspronkelijke context tijdens inferentie. Wij introduceren GradMem, dat context in het geheugen schrijft via per-sample optimalisatie tijdens testtijd. Gegeven een context voert GradMem een paar stappen van gradient descent uit op een kleine set prefix-geheugentokens, terwijl de modelgewichten bevroren blijven. GradMem optimaliseert expliciet een zelfgesuperviseerd contextreconstructieverlies op modelniveau, wat resulteert in een verliesgedreven schrijfoperatie met iteratieve foutcorrectie, in tegenstelling tot forward-only methoden. Bij associatieve key-value retrieval presteert GradMem beter dan forward-only geheugenschrijvers met dezelfde geheugengrootte, en extra gradientstappen schalen de capaciteit veel effectiever op dan herhaalde forward writes. Wij tonen verder aan dat GradMem verder reikt dan synthetische benchmarks: met voorgetrainde taalmodellen behaalt het competitieve resultaten op natuurlijke-taaltaken, waaronder bAbI- en SQuAD-varianten, waarbij het uitsluitend vertrouwt op informatie die in het geheugen gecodeerd is.
English
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.