PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2603.13875v2cs.CLcs.LG

TL;DR

Long-context applications need compact memory that supports answering queries after the original context is removed. GradMem writes context into memory through a few test-time gradient steps, outperforming forward-only writers and transferring to natural-language tasks.

  • Problem

    The paper asks whether a small fixed-size memory can preserve task-relevant context well enough to predict targets from queries after the original context is removed.

  • Method

    GradMem optimizes writable memory-token embeddings at test time with a model-level self-supervised reconstruction objective while keeping base-model weights fixed.

  • Results

    GradMem outperforms forward-only writers at equal memory size, gains capacity with additional gradient writes, and transfers to pretrained-language-model tasks including bAbI and SQuAD variants.

  • Takeaways & Limitations

    A few gradient-based WRITE steps provide a stronger context-memorization rule than forward-only updates across controlled retrieval and natural-language evaluations.

  • Takeaways & Limitations

    GradMem increases compute and GPU-memory usage because training differentiates through its unrolled WRITE optimization steps.

Abstract

from arXiv · show

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 compressive 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.

1. Introduction

GradMem compresses long contexts into a compact memory state through a few test-time gradient updates on writable memory tokens, allowing queries without the original context. Its explicit reconstruction objective enables iterative error correction, improves capacity over forward-only writing, and transfers to natural-language tasks.

  • Motivation: Long-context applications need reusable compression: read a context once, store what matters, and answer many queries without repeatedly processing the same tokens.The dominant KV-cache approach retains intermediate activations, creating substantial memory overhead.
  • Method: GradMem performs a small number of test-time gradient descent updates on special memory-token embeddings while keeping the base model weights fixed.The memory initialization is meta-learned so useful context representations can be written in few steps.
  • Method: GradMem uses a model-level self-supervised context-reconstruction WRITE objective, providing an explicit signal for what remains unencoded and a compute–compression tradeoff.Additional gradient steps improve the memory state through iterative error correction, without specialized per-layer update rules.
  • Results: GradMem’s capacity scaling is characterized on associative retrieval, and the same task-agnostic reconstruction objective transfers to pretrained language models on bAbI, SQuAD variants, and language modeling.These evaluations operate with the original context removed and rely on information encoded in memory.
  • Results: On associative key–value retrieval, gradient-based writing stores more information than same-size forward-only methods, while additional gradient updates consistently increase capacity.Repeated forward-only writes provide limited or inconsistent gains, and even one gradient WRITE update can outperform one forward-only update.

2. GradMem

GradMem addresses context removal by compressing each context into a fixed-size memory that supports prediction from memory and query alone. It writes this memory through a few per-example gradient updates on frozen-model prefix tokens, using task-agnostic context reconstruction and meta-learned initialization.

  • Context-removal setting: GradMem compresses context C into fixed-size memory M so the model predicts target Y from query Q without direct access to C.The context-removal constraint requires all task-relevant information to pass through M during READ.
  • Memory parameterization: For each example, GradMem updates only the memory state with a few test-time gradient-descent steps while keeping model parameters frozen.Memory consists of m vectors of dimension d, initialized from a shared meta-learned state M0 and updated for K WRITE steps.
  • WRITE objective: GradMem minimizes an autoregressive cross-entropy reconstruction loss over context tokens while prepending the current memory.The loss forces memory to encode context content that cannot be predicted from the preceding context prefix alone.
  • Meta-learning and READ: Training differentiates through the WRITE updates to optimize model parameters and M0 for downstream task performance in the READ phase.The WRITE objective is generic and task-agnostic, while the downstream task loss is computed from the resulting memory, query, and target.
  • Iterative writing: Unlike fixed forward-only writes, GradMem uses loss-driven iterative refinement that can correct write errors and trade additional test-time compute for greater memory capacity.The mechanism obtains a direct signal of how well memory explains the current context and adapts computation to that context.

3. Experiments and Results

Experiments evaluate context removal across associative key–value retrieval, bAbI, Short SQuAD, and WikiText language modeling. GradMem’s gradient-based memory writing outperforms forward-only writing on retrieval and remains competitive on downstream language tasks, with inference-time WRITE steps and memory size affecting performance.

  • Experimental setup: Each context-removal example separates context C, query Q, and target Y, requiring the key–value mapping to be written into memory during WRITE.Associative retrieval uses key–value pairs, and the model answers by retrieving the value associated with the queried key.
  • Downstream language tasks: On downstream language tasks, GradMem remains competitive: larger memory helps Short SQuAD substantially, improves language modeling slightly, and does not improve bAbI in these runs.The evaluation covers bAbI reasoning, Short SQuAD question answering, and WikiText language modeling; full-context GPT-2 remains the upper bound.
  • Associative KV-retrieval: On associative KV-retrieval, GradMem achieves higher accuracy than forward-only RMT with the same single memory state of 8 vectors.Repeated RMT passes provide weak or inconsistent gains because each pass replaces rather than accumulates or explicitly corrects the previous memory.
  • Associative KV-retrieval: Increasing GradMem’s evaluation-time WRITE iterations substantially improves exact match across settings with different numbers of key–value pairs.The gains occur with fixed model parameters and reflect additional inference-time refinement toward the inner-task solution.
  • Associative KV-retrieval: Larger evaluation-time WRITE budgets trade additional inference compute for higher accuracy without requiring model re-optimization.By contrast, larger training-time WRITE budgets are expensive because the meta-learning objective backpropagates through longer WRITE trajectories.

4. Discussion and Conclusions … A. Related Work

GradMem writes a compact context memory through test-time gradient descent, achieving strong retrieval and transfer while introducing compute and memory costs. It differs from related compression, fast-weight, and test-time-training methods through its single context-level memory state and explicit optimization-based WRITE operation.

  • 4. Discussion and Conclusions: GradMem’s WRITE phase increases compute and GPU memory usage because training and inference backpropagate through unrolled optimization steps.The retained computational graph also constrains attention implementations, including common high-performance kernels such as FlashAttention.
  • 4. Discussion and Conclusions: With only 1–2 memory vectors and 1–2 WRITE steps, GradMem reconstructs 32-token spans on pretrained Llama-3.2-1B and Llama-3.2-3B models.The READ model remains frozen during this text-reconstruction evaluation.
  • 4. Discussion and Conclusions: When contexts are reused across queries, one-time WRITE computation can be amortized because READ attends over compact memory and the query instead of the full context.The benefit applies when |C| ≫ |M|.
  • 4. Discussion and Conclusions: GradMem uses a task-agnostic token-level reconstruction objective that yields strong associative-retrieval gains and transfers to natural-language tasks, though it may not be optimal universally.The discussion identifies downstream-task-specific objectives as a possible direction beyond reconstruction.
  • 4. Discussion and Conclusions: GradMem writes context into a small set of memory tokens using a few test-time gradient-descent steps while keeping model weights fixed.This WRITE/READ mechanism is presented as the paper’s central contribution.
  • 4. Discussion and Conclusions: Increasing gradient WRITE steps reliably improves associative retrieval, outperforming forward-only updates with the same architecture and memory size.The comparison comes from controlled KV-retrieval experiments.
  • A. Related Work: Prior work compresses long contexts through architectural memory, recurrent or segment-level memory tokens, compact representations, autoencoding, and associative or fast-weight mechanisms.Examples include Compressive Transformers, RMT-style architectures, ICAE, SelfCP, and associative-memory transformers.
  • A. Related Work: Unlike prior memory and TTT approaches, GradMem uses one input-level state written once per context, optimizes memory tokens with model-level reconstruction, and avoids online per-token updates.TTT methods instead adapt layer parameters or states online using self-supervised reconstruction of layer inputs or activations.

B. Implementation, Training, and Hyperparameter Details

This section describes GradMem’s task-specific initialization, memory configuration, optimization settings, and frozen-head inference procedure. It also notes implementation availability and training-stability considerations for associative retrieval.

  • Implementation: The implementation code is publicly available in the GradMem GitHub repository.The repository URL is https://github.com/yurakuratov/gradmem.
  • Memory configuration: GradMem keeps the memory embedding dimension fixed to dmem = 64 in all ARMT experiments and varies only the number of memory tokens by task.Memory tokens in RMT, ARMT, and GradMem otherwise match the models’ input-embedding dimension.
  • Training initialization: For associative retrieval, models train from scratch with a context-length curriculum, while GradMem begins at 32 key–value pairs.Larger-length training is initialized from the final checkpoint at a smaller length; extra inner-loop WRITE steps can make training less stable.
  • Training initialization: Most bAbI and language-modeling models are fine-tuned from pretrained checkpoints, while RMT on bAbI starts from a GPT-2 (124M) checkpoint already fine-tuned on bAbI.Language modeling also reports a GradMem variant initialized from an RMT checkpoint trained on the same objective.
  • Optimization hyperparameters: GradMem tunes the inner learning rate α ∈[0.01, 10], uses α = 0.4 as a strong NLP default, and searches α for the best KV-retrieval setup at each K.Performance is not strongly dependent on α within a reasonable range, and exhaustive tuning is not performed for every experiment.
  • Heads and inference: For pretrained models, READ uses the original LM head, whereas a separate WRITE head is initialized from it and fine-tuned for reconstruction; inference updates only memory state M.Both heads and the rest of the model remain frozen during inference, and language-modeling experiments omit a separate WRITE head.

C. Accelerating Double Backwards Through Attention

The inner-loop optimization is bottlenecked by backwards-over-backwards attention, motivating optimized implementations that reduce runtime and GPU memory. Among these methods, Fast forward →autograd is fastest on long sequences, Manual HVP is most memory-efficient, and Flash HVP offers the best balance.

  • C. Accelerating Double Backwards Through Attention: For L=1024 tokens, optimized double-backward attention reduces backward time from ∼1000 ms to ∼600 ms and peak GPU memory from ∼60 GB to ∼30 GB.The measurements come from the authors’ setup.
  • C. Accelerating Double Backwards Through Attention: The implementations span eager autograd, SDPA-based forward passes with manual or autograd backward construction, fully analytical Manual HVP, and fused-kernel Flash HVP.Fast forward →autograd recomputes the attention forward during backward to avoid storing forward intermediates, while Flash HVP combines fused forward and backward kernels with analytical double backward.
  • C. Accelerating Double Backwards Through Attention: On shorter sequences, eager attention is most practical, whereas longer contexts benefit from the proposed optimizations.The comparison evaluates the speed of backward methods across sequence lengths.
  • C. Accelerating Double Backwards Through Attention: For longer contexts, Fast forward →autograd is fastest, Manual HVP is by far most memory-efficient, and Flash HVP ranks second in both speed and memory requirements.Flash HVP is characterized as the most balanced approach overall.

D. Computational Analysis: When GradMem is Compute-Efficient

GradMem trades higher one-time WRITE-time computation for cheaper READ-time inference over a compact memory prefix. This trade-off becomes favorable when contexts are long and reused across many queries, with measured latency showing amortized advantages over baseline models.

  • Compute trade-off: GradMem pays additional WRITE-time compute once, then reduces READ-time cost because queries attend only over a short memory prefix.The method performs K gradient updates during WRITE, with each update costing R times a forward pass over the context.
  • Compute break-even: When c > m and the same context is reused beyond the break-even query threshold, GradMem uses less total compute than repeated full-context inference.The amortized READ savings grow linearly with N, while the WRITE cost is paid once per context.
  • Measured latency: GradMem has a higher initial latency from gradient-based WRITE but a smaller per-READ cost than cached-transformer and recurrent baselines.This initial overhead becomes less pronounced as the number of READ operations increases.
  • Measured latency: GradMem is about 1.6x faster than Mamba when the same context is reused across many READ operations.Figure 7 reports A100 GPU results with query length 24, batch size 16, and K = 1.
  • Measured latency: GradMem consistently outperforms Mamba across evaluated context lengths and breaks even with GPT-2 after approximately 64 READ phases for context sizes 256 and 1024.These comparisons use measured WRITE and subsequent READ operations for the same context.
  • Larger-model efficiency: For a single READ after WRITE, LaCT, TTT-MLP, and TTT-Linear are generally faster than GradMem, but GradMem amortizes better as context reuse increases.The alternative methods avoid backpropagation through the full model.

E. Relation between Exact Match and Inner Loss

Increasing WRITE iterations lowers the reconstruction inner loss and correlates with improved Exact Match at larger evaluation iteration counts. The loss decomposition indicates selective refinement: value-token loss decreases while key-token loss remains comparatively stable.

  • E. Relation between Exact Match and Inner Loss: Increasing K reduces the inner loss, indicating more accurate memory states under the reconstruction objective.The analysis links this reduction to improved Exact Match when evaluating with larger Keval.
  • E. Relation between Exact Match and Inner Loss: The reduction in inner loss correlates with improved Exact Match when evaluating with larger Keval.This relationship explains why extrapolating the number of WRITE iterations improves downstream performance.
  • E. Relation between Exact Match and Inner Loss: Loss on key tokens remains comparatively stable across WRITE iterations, while loss on value tokens decreases as Keval increases.This pattern suggests that learned memory selectively refines value information rather than storing the full context verbatim.

F. Memory Size vs. WRITE Steps (K) vs. Number of KV-pairs

The section examines how memory-token count and WRITE optimization steps affect associative retrieval across contexts containing 8 or 16 KV-pairs. More WRITE steps improve performance and partially offset smaller memory, but larger contexts still require more memory.

  • Experimental setup: The experiment sweeps m ∈ {1, 2, 4, 8} memory tokens, K = 1 or K = 2 WRITE steps, and contexts containing 8 or 16 KV-pairs.All runs start from a checkpoint trained with 8 memory tokens and continue training after reducing memory to the target size.
  • WRITE steps and memory efficiency: Increasing WRITE steps consistently improves Exact Match accuracy at a fixed memory budget, especially when memory is small.For 8 KV-pairs, the K = 1 to K = 2 improvement is substantial at m = 1 and m = 2; for 16 KV-pairs, gains remain large in the low-memory regime.
  • Context size and memory: Reducing the number of memory tokens decreases the amount of context that can be reliably stored, while larger contexts require larger memory budgets.The figure compares 8- and 16-KV-pair contexts across the same memory-token and WRITE-step settings.
  • Capacity limits: With K = 2, 8 KV-pairs can be stored almost perfectly using m = 2 memory tokens.This indicates that additional WRITE optimization can substantially compensate for a smaller memory budget in the 8-KV-pair setting.
  • Capacity limits: With K = 2, 16 KV-pairs reach near-perfect performance only when m ≥ 4, showing that additional optimization steps cannot fully remove the bottleneck of insufficient memory.The results support a partial-compensation effect rather than unlimited substitution of WRITE steps for memory tokens.

G. Memory State Size Utilization: GradMem vs. Mamba

GradMem compresses context more effectively than Mamba under limited state size, using 1024 floats across 8 memory tokens versus Mamba’s roughly 41k floats across layers. The comparison fixes the hidden state size at 128 and evaluates associative retrieval with exact-match performance averaged across three runs.

  • Memory State Size: GradMem uses 1024 floats across 8 memory tokens, whereas Mamba uses around 41k floats across all layers.The main experiments reuse most Mamba parameters from the mamba-130m checkpoint.
  • Compression Comparison: With limited state size, GradMem performs compression much better than Mamba.The comparison matches hidden state size at d model = 128 while varying convolution kernel and state size.
  • Evaluation Setup: Associative retrieval performance is reported as exact match (%) across task lengths, averaged across 3 runs with hidden state size 128.

H. KV Retrieval with LaCT and TTT-Linear

GradMem is compared with TTT-Linear and LaCT on associative KV retrieval, using prefix-token memory while preserving the base architecture. Despite a much smaller model-level memory state, additional GradMem WRITE steps remain strongest overall, whereas one step is comparable to TTT-Linear and below LaCT’s best larger-context runs.

  • The evaluation compares GradMem, TTT-Linear, and LaCT on the same associative KV-retrieval data and model scale.
  • GradMem preserves the base architecture without internal layer changes, per-layer writable states, or layer-local objectives.Its input-level memory is supplied as prefix tokens and optimized in a dedicated WRITE phase with a model-level context-reconstruction objective.
  • GradMem uses 1024 floats across 8 memory tokens, versus about 49k for LaCT and 17k for TTT-Linear.The comparison is not matched by memory state size.
  • One GradMem WRITE step is comparable to TTT-Linear but below LaCT’s best runs at larger context sizes.LaCT is a strong baseline with high variance, especially as the number of key–value pairs increases.
  • Additional GradMem WRITE steps remain strongest overall despite using a much smaller model-level memory state.

I. Text Compression with GradMem on Larger Pretrained Models

GradMem extends text compression to pretrained Llama-3.2-1B and 3B models, reconstructing context from only 1–2 learned memory vectors while keeping the READ model frozen. Increasing WRITE steps from 1 to 2 improves reliable storage capacity, corresponding to roughly 16x–32x compression based on coarse estimates.

  • Text reconstruction: GradMem compresses and reconstructs text on pretrained Llama-3.2-1B and Llama-3.2-3B models using token-level reconstruction accuracy.The context is written into memory and reconstructed during READ without access to the original context.
  • Write and read: During WRITE, learned LoRA adapters optimize memory vectors at test time, while the pretrained base model remains frozen during READ.The setup trains configurations varying sequence length N, WRITE steps K, and memory-vector count nmem.
  • Text reconstruction: With only 1–2 memory vectors, GradMem remains effective on 1B- and 3B-parameter pretrained models.Table 9 evaluates target token-accuracy thresholds across sequence lengths N ∈{8, 16, 32, 64}.
  • Write optimization: Increasing WRITE steps from K = 1 to K = 2 consistently improves the amount of text that can be reliably stored.This trend mirrors the observed effect of increasing memory size and WRITE compute in key–value retrieval.
  • Compression: 16x–32x compression is achieved when up to 32–64 input tokens are stored in only 1–2 memory vectors.Because the sequence-length sweep is coarse, these values are approximate capacity estimates.

J. Input-level vs. Per-layer Memory · K. GradMem Algorithm and Minimal Code

Section J shows that input-level memory is more parameter-efficient than size-matched per-layer memory, although a much larger per-layer state is strongest. Section K specifies GradMem’s gradient-based WRITE/READ procedure and its minimal implementation.

  • J. Input-level vs. Per-layer Memory: The ablation changes only the test-time trainable memory parameters while keeping the model-level reconstruction objective fixed.This isolates whether memory is stored at the input level or separately in each layer.
  • J. Input-level vs. Per-layer Memory: GradMem uses 8 input memory tokens totaling 8 × 128 = 1024 parameters, matched against a per-layer KV-cache with 1 × 4 × 128 × 2 = 1024 parameters.Experiments use 4-layer models with hidden size 128 and associative retrieval of 8, 16, and 32 KV-pairs.
  • J. Input-level vs. Per-layer Memory: At matched memory size, input-level memory consistently outperforms per-layer memory, with the gap increasing as associative retrieval becomes harder.At 32 KV-pairs, size-matched per-layer memory is far below input-level memory for both K = 1 and K = 2.
  • J. Input-level vs. Per-layer Memory: An 8× larger per-layer KV-cache becomes the strongest variant, indicating that input-level memory is more parameter-efficient rather than universally more accurate.The comparison contrasts compact matched states with a substantially larger per-layer state.
  • K. GradMem Algorithm and Minimal Code: GradMem’s WRITE procedure optimizes per-example memory with gradients while keeping the model parameters θ frozen.The algorithm encodes context into memory through a meta-learned initialization and multiple WRITE steps.
  • K. GradMem Algorithm and Minimal Code: READ removes the original context, and training backpropagates the outer loss through the WRITE updates.This makes the written memory the sole source of information during inference after context removal.
  • K. GradMem Algorithm and Minimal Code: The minimal PyTorch-like implementation mirrors GradMem while omitting padding, separate WRITE heads, and other engineering details.Helper functions hide tokenization, padding, label shifting, and memory-prefix construction; the full implementation is available in the cited repository.
Loading 2603.13875v2…