GradMem: Aprender a Escribir Contexto en la Memoria con Descenso de Gradiente en Tiempo de Prueba

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

March 14, 2026

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

cs.AI

Resumen

Muchas aplicaciones de grandes modelos de lenguaje requieren condicionarse sobre contextos largos. Los Transformers suelen soportar esto almacenando una gran caché KV por capa de activaciones pasadas, lo que incurre en una sobrecarga de memoria sustancial. Una alternativa deseable es la memoria compresiva: leer un contexto una vez, almacenarlo en un estado compacto y responder muchas consultas a partir de ese estado. Estudiamos esto en un escenario de eliminación de contexto, donde el modelo debe generar una respuesta sin acceso al contexto original en el momento de la inferencia. Introducimos GradMem, que escribe el contexto en la memoria mediante una optimización por muestra en tiempo de prueba. Dado un contexto, GradMem realiza unos pocos pasos de descenso de gradiente sobre un pequeño conjunto de tokens de memoria de prefijo, manteniendo los pesos del modelo congelados. GradMem optimiza explícitamente una pérdida de reconstrucción de contexto auto-supervisada a nivel de modelo, resultando en una operación de escritura impulsada por la pérdida con corrección iterativa de errores, a diferencia de los métodos de solo forward. En la recuperación asociativa clave-valor, GradMem supera a los escritores de memoria de solo forward con el mismo tamaño de memoria, y los pasos de gradiente adicionales escalan la capacidad mucho más efectivamente que las escrituras forward repetidas. Además, mostramos que GradMem se transfiere más allá de los benchmarks sintéticos: con modelos de lenguaje preentrenados, obtiene resultados competitivos en tareas de lenguaje natural que incluyen variantes de bAbI y SQuAD, dependiendo únicamente de la información codificada en la 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.