GradMem: Обучение записи контекста в память с градиентным спуском во время тестирования

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

March 14, 2026

Авторы: Yuri Kuratov, Matvey Kairov, Aydar Bulatov, Ivan Rodkin, Mikhail Burtsev

cs.AI

Аннотация

Многие приложения больших языковых моделей требуют кондиционирования на длинных контекстах. Трансформеры обычно поддерживают это за счет хранения большого KV-кэша активаций прошлых слоев, что приводит к значительным затратам памяти. Желательной альтернативой является компрессионная память: однократное чтение контекста, сохранение его в компактном состоянии и ответ на множество запросов из этого состояния. Мы изучаем это в условиях удаления контекста, когда модель должна генерировать ответ без доступа к исходному контексту во время вывода. Мы представляем GradMem, который записывает контекст в память посредством оптимизации на этапе тестирования для каждого примера. Для заданного контекста GradMem выполняет несколько шагов градиентного спуска на небольшом наборе префиксных токенов памяти, сохраняя веса модели замороженными. GradMem явно оптимизирует самоконтролируемую функцию потерь реконструкции контекста на уровне модели, что приводит к операции записи, управляемой потерями, с итеративной коррекцией ошибок, в отличие от методов, основанных только на прямом проходе. На задаче ассоциативного извлечения "ключ-значение" GradMem превосходит методы записи памяти только с прямым проходом при том же размере памяти, а дополнительные шаги градиента масштабируют емкость гораздо эффективнее, чем повторные записи прямым проходом. Мы также показываем, что GradMem обобщается за пределы синтетических бенчмарков: с предобученными языковыми моделями он достигает конкурентоспособных результатов на задачах обработки естественного языка, включая варианты bAbI и SQuAD, полагаясь исключительно на информацию, закодированную в памяти.

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.