GradMem : Apprendre à écrire le contexte en mémoire par descente de gradient au moment du test

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

Résumé

De nombreuses applications de grands modèles de langage nécessitent de conditionner sur de longs contextes. Les Transformers prennent généralement cela en charge en stockant un grand cache clé-valeur (KV-cache) par couche pour les activations passées, ce qui entraîne une surcharge mémoire substantielle. Une alternative souhaitable est la mémoire compressive : lire un contexte une fois, le stocker dans un état compact et répondre à de nombreuses requêtes à partir de cet état. Nous étudions cela dans un cadre de suppression de contexte, où le modèle doit générer une réponse sans avoir accès au contexte original au moment de l'inférence. Nous présentons GradMem, qui écrit le contexte en mémoire via une optimisation par échantillon au moment du test. Étant donné un contexte, GradMem effectue quelques étapes de descente de gradient sur un petit ensemble de jetons de mémoire de préfixe tout en gardant les poids du modèle gelés. GradMem optimise explicitement une perte d'auto-supervision au niveau du modèle pour la reconstruction du contexte, ce qui donne une opération d'écriture pilotée par la perte avec correction d'erreur itérative, contrairement aux méthodes uniquement avant. Sur la tâche de récupération associative clé-valeur, GradMem surpasse les écrivains de mémoire uniquement avant avec la même taille de mémoire, et des étapes de gradient supplémentaires augmentent la capacité bien plus efficacement que des écritures avant répétées. Nous montrons en outre que GradMem se généralise au-delà des benchmarks synthétiques : avec des modèles de langage pré-entraînés, il obtient des résultats compétitifs sur des tâches de langage naturel, y compris des variantes de bAbI et de SQuAD, en ne reposant que sur les informations encodées en mémoire.

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.