GradMem: Aprendendo a Escrever Contexto na Memória com Descida de Gradiente no Momento do Teste

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

Resumo

Muitas aplicações de modelos de linguagem de grande porte exigem o condicionamento em contextos longos. Os Transformers normalmente suportam isso armazenando um grande cache KV por camada de ativações passadas, o que incorre em uma sobrecarga substancial de memória. Uma alternativa desejável é a memória compressiva: ler um contexto uma vez, armazená-lo em um estado compacto e responder a muitas consultas a partir desse estado. Estudamos isso em um cenário de remoção de contexto, onde o modelo deve gerar uma resposta sem acesso ao contexto original no momento da inferência. Apresentamos o GradMem, que grava o contexto na memória por meio de uma otimização por amostra em tempo de teste. Dado um contexto, o GradMem executa algumas etapas de descida de gradiente em um pequeno conjunto de tokens de memória de prefixo, mantendo os pesos do modelo congelados. O GradMem otimiza explicitamente uma perda de reconstrução de contexto auto supervisionada a nível de modelo, resultando em uma operação de gravação orientada por perda com correção iterativa de erros, diferentemente dos métodos de apenas propagação direta (forward-only). Na recuperação associativa de chave-valor, o GradMem supera escritores de memória de apenas propagação direta com o mesmo tamanho de memória, e etapas adicionais de gradiente escalam a capacidade de forma muito mais eficaz do que gravações diretas repetidas. Mostramos ainda que o GradMem se transfere para além de benchmarks sintéticos: com modelos de linguagem pré-treinados, ele obtém resultados competitivos em tarefas de linguagem natural, incluindo variantes do bAbI e do SQuAD, dependendo apenas das informações codificadas na memória.

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.