GradMem: Imparare a Scrivere il Contesto in Memoria con Discesa del Gradiente al Momento del Test

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

March 14, 2026

Autori: Yuri Kuratov, Matvey Kairov, Aydar Bulatov, Ivan Rodkin, Mikhail Burtsev

cs.AI

Abstract

Molte applicazioni di grandi modelli linguistici richiedono il condizionamento su contesti lunghi. I Transformer supportano tipicamente questa funzionalità memorizzando una grande cache KV per layer delle attivazioni passate, che comporta un sovraccarico di memoria sostanziale. Un'alternativa desiderabile è la memoria compressiva: leggere un contesto una volta, memorizzarlo in uno stato compatto e rispondere a molte query da quello stato. Studiamo questo approccio in uno scenario di rimozione del contesto, dove il modello deve generare una risposta senza accesso al contesto originale al momento dell'inferenza. Introduciamo GradMem, che scrive il contesto in memoria tramite un'ottimizzazione per campione al momento del test. Dato un contesto, GradMem esegue pochi passi di discesa del gradiente su un piccolo insieme di token di memoria del prefisso, mantenendo congelati i pesi del modello. GradMem ottimizza esplicitamente una perdita auto-supervisionata a livello di modello per la ricostruzione del contesto, risultando in un'operazione di scrittura guidata dalla perdita con correzione iterativa dell'errore, a differenza dei metodi di sola forward pass. Sul recupero associativo chiave-valore, GradMem supera i scrittori di memoria di sola forward pass con la stessa dimensione di memoria, e passi aggiuntivi del gradiente scalano la capacità in modo molto più efficace rispetto a scritture forward ripetute. Mostriamo inoltre che GradMem si trasferisce oltre i benchmark sintetici: con modelli linguistici preaddestrati, ottiene risultati competitivi su compiti di linguaggio naturale includendo varianti di bAbI e SQuAD, basandosi solo sulle informazioni codificate in memoria.

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.