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

• URL: http://arxiv.org/abs/2603.13875v1
• Date: Sat, 14 Mar 2026 10:17:33 GMT
• Title: GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent
• Abstract summary: We introduce GradMem, which writes context into memory via per-sample test-time optimization.<n> GradMem explicitly optimize a model-level self-supervised context reconstruction loss.<n>It attains competitive results on natural language tasks including bAbI and SQuAD variants.
• Score: 10.16391634040356
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Memory Grafting: Scaling Language Model Pre-training via Offline Conditional Memory [65.39827296429527]
  Scaling conditional memory offers a promising way to increase language-model capacity.<n>Existing methods such as Engram learn large memory tables from scratch during pre-training.<n>We propose Memory Grafting, a conditional memory scaling method that utilizes frozen hidden states from a grafting model as conditional n-gram memory.
  arXiv  Detail & Related papers  (2026-05-20T09:35:04Z)
• Is One Score Enough? Rethinking the Evaluation of Sequentially Evolving LLM Memory [50.857546269660276]
  We introduce SeqMem-Eval, a diagnostic evaluation framework for sequentially evolving large language models (LLMs) memory.<n>Rather than focusing only on final performance, SeqMem-Eval evaluates how memory states evolve, generalize, consolidate experience, and retain useful information during sequential inference.
  arXiv  Detail & Related papers  (2026-05-14T20:15:22Z)
• NextMem: Towards Latent Factual Memory for LLM-based Agents [58.35585202907478]
  NextMem is a latent factual memory framework that utilizes an autoregressive autoencoder to efficiently construct latent memory.<n>Extensive experiments demonstrate that NextMem achieves superior performance.
  arXiv  Detail & Related papers  (2026-02-26T14:35:27Z)
• Evaluating Long-Term Memory for Long-Context Question Answering [100.1267054069757]
  We present a systematic evaluation of memory-augmented methods using LoCoMo, a benchmark of synthetic long-context dialogues annotated for question-answering tasks.<n>Our findings show that memory-augmented approaches reduce token usage by over 90% while maintaining competitive accuracy.
  arXiv  Detail & Related papers  (2025-10-27T18:03:50Z)
• HMT: Hierarchical Memory Transformer for Efficient Long Context Language Processing [33.720656946186885]
  Hierarchical Memory Transformer (HMT) is a novel framework that facilitates a model's long-context processing ability.<n>HMT consistently improves the long-context processing ability of existing models.
  arXiv  Detail & Related papers  (2024-05-09T19:32:49Z)
• Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs [61.40047491337793]
  We present Hierarchical cOntext MERging (HOMER), a new training-free scheme designed to overcome the limitations of large language models. HomeR uses a divide-and-conquer algorithm, dividing long inputs into manageable chunks. A token reduction technique precedes each merging, ensuring memory usage efficiency.
  arXiv  Detail & Related papers  (2024-04-16T06:34:08Z)
• Compressed Context Memory For Online Language Model Interaction [39.72054168889216]
  This paper presents a context key/value compression method for Transformer language models in online scenarios. As the context lengthens, the attention process demands increasing memory and computations, which in turn reduces the throughput of the language model. We propose a compressed context memory system that continually compresses the accumulating attention key/value pairs into a compact memory space.
  arXiv  Detail & Related papers  (2023-12-06T10:50:43Z)
• Augmenting Language Models with Long-Term Memory [142.04940250657637]
  Existing large language models (LLMs) can only afford fix-sized inputs due to the input length limit. We propose a framework, Language Models Augmented with Long-Term Memory (LongMem), which enables LLMs to memorize long history.
  arXiv  Detail & Related papers  (2023-06-12T15:13:39Z)
• A Model or 603 Exemplars: Towards Memory-Efficient Class-Incremental Learning [56.450090618578]
  Class-Incremental Learning (CIL) aims to train a model with limited memory size to meet this requirement. We show that when counting the model size into the total budget and comparing methods with aligned memory size, saving models do not consistently work. We propose a simple yet effective baseline, denoted as MEMO for Memory-efficient Expandable MOdel.
  arXiv  Detail & Related papers  (2022-05-26T08:24:01Z)
• Training Language Models with Memory Augmentation [28.4608705738799]
  We present a novel training approach designed for training language models with memory augmentation. Our approach uses a training objective that directly takes in-batch examples as accessible memory. We demonstrate significant gains over previous memory-augmented approaches.
  arXiv  Detail & Related papers  (2022-05-25T11:37:29Z)
• LaMemo: Language Modeling with Look-Ahead Memory [50.6248714811912]
  We propose Look-Ahead Memory (LaMemo) that enhances the recurrence memory by incrementally attending to the right-side tokens. LaMemo embraces bi-directional attention and segment recurrence with an additional overhead only linearly proportional to the memory length. Experiments on widely used language modeling benchmarks demonstrate its superiority over the baselines equipped with different types of memory.
  arXiv  Detail & Related papers  (2022-04-15T06:11:25Z)
• Memory-Based Semantic Parsing [79.48882899104997]
  We present a memory-based model for context-dependent semantic parsing. We learn a context memory controller that manages the memory by maintaining the cumulative meaning of sequential user utterances.
  arXiv  Detail & Related papers  (2021-09-07T16:15:13Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.