GradMem: Kontext durch Gradientenabstieg zur Testzeit in den Speicher schreiben lernen
GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent
March 14, 2026
Autoren: Yuri Kuratov, Matvey Kairov, Aydar Bulatov, Ivan Rodkin, Mikhail Burtsev
cs.AI
Zusammenfassung
Viele Anwendungen großer Sprachmodelle erfordern die Konditionierung auf lange Kontexte. Transformer-Modelle unterstützen dies typischerweise durch das Speichern eines großen KV-Cache vergangener Aktivierungen pro Schicht, was einen erheblichen Speicheraufwand verursacht. Eine wünschenswerte Alternative ist komprimierter Speicher: Ein Kontext wird einmal gelesen, in einem kompakten Zustand gespeichert, und viele Abfragen werden aus diesem Zustand beantwortet. Wir untersuchen dies in einem Kontextentfernungsszenario, bei dem das Modell zur Inferenzzeit eine Antwort generieren muss, ohne Zugriff auf den ursprünglichen Kontext zu haben. Wir stellen GradMem vor, das Kontext durch Optimierung zur Laufzeit pro Beispiel in den Speicher schreibt. Für einen gegebenen Kontext führt GradMem einige Schritte des Gradientenabstiegs auf einer kleinen Menge von Präfix-Speicher-Tokens durch, während die Modellgewichte eingefroren bleiben. GradMem optimiert explizit einen modellbasierten, selbstüberwachten Kontext-Rekonstruktionsverlust, was zu einem verlustgesteuerten Schreibvorgang mit iterativer Fehlerkorrektur führt, im Gegensatz zu rein vorwärtsgerichteten Methoden. Beim assoziativen Schlüssel-Wert-Abruf übertrifft GradMem rein vorwärtsgerichtete Speicherschreiber mit derselben Speichergröße, und zusätzliche Gradientenschritte skalieren die Kapazität wesentlich effektiver als wiederholte Vorwärtsschreibvorgänge. Wir zeigen weiter, dass GradMem über synthetische Benchmarks hinaus generalisiert: Mit vortrainierten Sprachmodellen erzielt es wettbewerbsfähige Ergebnisse auf natürlichen Sprachaufgaben, einschließlich bAbI- und SQuAD-Varianten, und stützt sich dabei nur auf die im Speicher kodierten Informationen.
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.