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

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

cs.CLMar 14, 2026

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

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

Organizations: AXXX, Cognitive AI Systems Lab, Moscow, Russia · MIRAI, Moscow, Russia · MBZUAI, Abu Dhabi, UAE · London Institute for Mathematical Sciences, London, UK

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

Figures & tables

Figure 1: GradMem learns to write context into memory via per-sample test-time optimization. Given a context, GradMem performs a few test-time gradient updates on memory state to minimize a self-supervised reconstruction loss (WRITE). The memory initialization is meta-learned so that useful context representations can be written with only a few gradient steps. At inference, the model answers queries using only the optimized memory and the query (READ), without access to the original context.

Figure 2: GradMem overview. (a) Each task sample is represented as context $C$ , query $Q$ , and target $Y$ . A context encoder $E_{\theta}$ compresses $C$ into a fixed-size memory $\hat{\mathcal{M}}$ in a WRITE phase, and the model predicts $Y$ from $[\hat{\mathcal{M}};Q]$ in a READ phase, without access to $C$ . (b) Meta-learning view: a shared initialization $\mathcal{M}_{0}$ and model parameters $\theta$ are learned across training examples (outer loop), while at test time each context $C_{i}$ adapts its own memory $\hat{\mathcal{M}}(C_{i})$ via a few gradient steps (dotted trajectory). (c) Test-time gradient descent on memory. Starting from the meta-learned initialization $\mathcal{M}_{0}$ , GradMem updates the per-sample memory state during WRITE with $K$ steps of gradient descent on the context reconstruction loss $\mathcal{L}_{\text{write}}(\mathcal{M};C)$ . At READ, the model predicts the task target using only $[\hat{\mathcal{M}};Q]$ . During training, $\theta$ and $\mathcal{M}_{0}$ are optimized by backpropagating through the WRITE inner loop, so the model learns to use gradient descent as an operation that writes useful information about $C$ into memory; at inference, only the per-sample memory state is updated.

	Number of KV-pairs
Model	4	8	16	32	64	96
Transformer: $\mathcal{M}$ =KV-cache	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	99.8 $\pm$ 0.0	99.8 $\pm$ 0.3	96.5 $\pm$ 2.9	98.8 $\pm$ 0.0
Mamba: $\mathcal{M}$ =per-layer recurrent state	99.9 $\pm$ 0.1	98.9 $\pm$ 0.4	98.7 $\pm$ 0.2	90.2 $\pm$ 10.2	95.2 $\pm$ 0.1	92.2 $\pm$ 0.4
ARMT: $\mathcal{M}$ =per-layer associative matrix	99.0 $\pm$ 0.3	98.5 $\pm$ 0.5	97.4 $\pm$ 0.3	54.9 $\pm$ 2.1	22.6 $\pm$ 3.9	15.2 $\pm$ 0.2
Forward-only write (RMT): $\mathcal{M}$ =8 mem vectors	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	45.5 $\pm$ 0.2	44.3 $\pm$ 0.0	19.3 $\pm$ 3.1	12.9 $\pm$ 0.2
x2 memory updates			69.6 $\pm$ 28.1	18.7 $\pm$ 3.4	–	–

Table 1: KV-retrieval: gradient-based WRITE outperforms forward-only WRITE. Exact match retrieval accuracy (mean

\pm

std over 3 runs) for predicting a 2-token value from a 2-token key as the number of key–value pairs increases. Upper bound: a standard Transformer with full KV-cache retains all past activations. Per-layer memory baselines: Mamba and ARMT maintain state in every layer. Same memory state: RMT (forward-only write) and GradMem both use the same base architecture and the same memory size of

m{=}8

vectors, but a different WRITE rule. Repeating the forward-only WRITE by re-reading the same context segment (x2–x5) gives limited or inconsistent gains, whereas additional gradient-based WRITE steps (x2, x5) consistently improve retrieval at larger numbers of pairs.

Figure 3: Gradient-based memory updates (GradMem) outperform forward-only updates at the same memory size. With a memory state of 8 vectors, GradMem retrieves any of 16 key–value pairs with 95% accuracy, whereas a forward-only update rule stores only 8 pairs with high accuracy. More gradient steps increase the capacity of the same 8-vector memory to 96 pairs at 88% retrieval accuracy. Transformer with KV-cache as non-compressive memory serves as an upper bound.

	bAbI (EM $\uparrow$ )					SQuAD (EM $\uparrow$ )	LM (CE $\downarrow$ )
	QA1	QA2	QA3	QA4	QA5	short	128–255
Input length (tokens)	$\sim$ 40	$\sim$ 100	$\sim$ 300	$\sim$ 20	$\sim$ 20	$\sim$ 40	256
Full context models (upper bound)
GPT-2-124m	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	99.8 $\pm$ 0.1	100.0 $\pm$ 0.0	99.4 $\pm$ 0.1	64.2 $\pm$ 0.3	2.72 $\pm$ 0.00
Limit context to 128 tokens							3.20 $\pm$ 0.00
Pythia-160m	100.0 $\pm$ 0.0	99.7 $\pm$ 0.1	95.5 $\pm$ 2.7	100.0 $\pm$ 0.0	99.0 $\pm$ 0.1	48.9 $\pm$ 0.4	2.84 $\pm$ 0.05

Table 2: GradMem remains competitive on downstream language tasks. Results on bAbI (QA1–5), SQuAD (short) shown as Exact Match (EM) in %, and language modeling (cross-entropy). RMT is the matched memory state baseline; full-context Transformers are upper bounds, while Mamba/ARMT use different and larger memory states. We use 32 memory tokens for SQuAD and LM, and 8 on bAbI. We report mean

\pm

std over 3 runs.

Panel of Figure 4 — Figure 4: More gradient steps at test-time lead to better performance without fine-tuning. (a) Results for $K_{\text{train}}$ values of GradMem setups that could not achieve 99% exact match during fine-tuning. Red dashed lines denote the beginning of extrapolation. (b) Downstream task quality (Exact Match) correlates with WRITE objective in inner loop (reconstruction). Notably, GradMem reconstructs values better than keys , despite the WRITE objective treating them equally. Results on 96-pair KV-retrieval.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

	TTT layers ( Sun et al., 2025 )	GradMem (ours)
Usage pattern	Sequence-modeling layer: updates state online while processing tokens	Explicit two-phase WRITE/READ setting
Inner-loop input	Token $x_{t}$ (or token mini-batches)	Whole context segment $C$ (WRITE once per context)
Test-time parameters	Layer-specific parameters $W_{t}$ (updated from the layer’s inputs/activations), per layer	Prefix memory tokens $\mathcal{M}_{k}\in\mathbb{R}^{m\times d}$ (single memory state), per model
Self-supervised loss	Activation/input reconstruction, e.g. $\ell(W;x_{t})=\\|f(\tilde{x}_{t};W)-x_{t}\\|_{2}^{2}$ (or learned multi-view projections)	Context reconstruction $\mathcal{L}_{\text{write}}(\mathcal{M};C)$
Outer-loop objective	Next-token prediction (LM training)	Downstream task loss with $C$ removed at READ
Outer-loop parameters	Model params $\theta$ + reconstruction task/view params	Model params $\theta$ + memory init $\mathcal{M}_{0}$ (optional: memory projections / control tokens)

Appendix

Table 3: Comparison of GradMem to test-time training (TTT) layers. Both approaches perform learning at inference time via a self-supervised loss, but differ in what is adapted (layer parameters vs. memory tokens), when it is adapted (token-level online vs. context-level WRITE), and what the self-supervised signal reconstructs (layer inputs/activations vs. context tokens).

Property	TTT layers ( Sun et al., 2025 )	Titans ( Behrouz et al., 2025c )	Atlas ( Behrouz et al., 2025a )	LaCT / TTT Done Right ( Zhang et al., 2026 )	GradMem (ours)
Memory level	Layer	Layer	Layer	Layer	Model
Memory update	Per-token/small-mini-batch	Per-token /chunk $+$ momentum/decay	Per-token /chunk with sliding window updates	Per chunk	One WRITE per full context $C$
Inner-loop input	Token $x_{t}$ / views of $x_{t}$	Key–value proj. $(k_{t},v_{t})$ of token $x_{t}$	Local window of keys and values	Chunk of keys and values	Full context $C$
Test-time parameters	Weights $W_{t}$ , per layer	MLP mem. $+$ momentum, per layer	MLP mem. $+$ momentum $+$ gates, per layer	Weights $W_{t}$ , per layer	Single memory prefix $\mathcal{M}\!\in\!\mathbb{R}^{m\times d}$
WRITE objective	View reconst., $L_{2}$	KV recall, $L_{2}$	Window KV recall, $L_{2}$	Chunk KV recall, dot product	Context reconst. $\mathcal{L}_{\mathrm{write}}(\mathcal{M};C)$
Outer-loop objective	Task loss / next-token prediction	Task loss / next-token prediction	Task loss / next-token prediction	Task loss; next-token prediction	Task loss at READ, context removed

Appendix

Table 4: GradMem and TTT methods. GradMem is characterized by explicit model-level memory update operation, performed once per context.

Task	num_mem_tokens	Training technique
Associative retrieval (AR)	8	Curriculum learning by number of KV-pairs (4, 8, 16, …); from randomly init model (4 layers, 4 heads, 128 hid)
Short SQuAD	32	From pretrained model; GradMem trained without learned memory projection
bAbI	8	From pretrained model
Language modeling	32	From pretrained model (all models); GradMem trained without separate WRITE head;

Appendix

Table 5: Task-specific memory hyperparameters and training initialization. In experiments ARMT uses

d_{\text{armt\_mem}}=64

for associative matrix memory. Memory tokens in RMT, ARMT, and GradMem have the same dimension as models input embeddings.

Task	Model	Base LM	Layers	Hidden dim	Heads	# mem tokens ( $m$ )	$d_{\text{mem}}$	$K$	$\alpha$	Initialization, curriculum	Total memory state size
Associative retrieval (AR)	Transformer	Llama	4	128	4	–	–	–	–	Random	–
	Mamba	Mamba	4	128	–	–	16 (state_size)	–	–	Random	40,960
	ARMT	Llama	4	128	4	8	64	–	–	Random, 1-2-4-8-16-32-64-96	198,144
	RMT	Llama	4	128	4	8	–	1–5	–	Random, 4-8-16-32-64-96	1,024
	TTT-Linear	TTT-Linear	4	128	4	–	–	1	1.0	Random	16,896
	LaCT	LaCT	4	128	4	–	–	1	0.001	Random	49,152

Appendix

Table 6: Hyperparameters. “–” indicates the parameter is not applicable. All transformer-based models for associative retrieval (AR) use 4 layers, hidden size 128, 4 attention heads. For NLP tasks, GPT-2 (124M), Pythia (160M), and Mamba (130M) are used as base models.

K

denotes the number of WRITE gradient steps (GradMem) or forward memory updates (RMT).

\alpha

is the inner learning rate for GradMem.

Panel of Figure 5 — Figure 5: Comparison of speed and memory consumption for attention backwards-over-backwards in GradMem. These results were obtained using an A100 GPU, with GPT-2 as the base model for GradMem with 8 memory tokens, a query size of 24, 1 inner SGD step and a batch size of 16.

Figure 6: Break-even number of queries for GradMem compute efficiency. The curve shows the minimum number of queries per context

N

required for GradMem to use less total compute than full-context inference with cached context representations. Estimates use query length

q{=}128

and memory size

m{=}32

; points above the curve favor GradMem.

Panel of Figure 7 — Figure 7: GradMem is competitive in latency when the same context is reused across many READ operations, and becomes about 1.6x faster than Mamba. Results obtained on an A100 GPU. WRITE+READ latency is shown for contexts of length 64 (a), 256 (b), and 1024 (c) tokens. The query length is 24 tokens and the batch size is 16. GradMem uses GPT-2 as the base model with $K=1$ . For GPT-2 and Mamba-130M, the WRITE operation corresponds to building the KV-cache and the recurrent state, respectively. The y-axis uses a logarithmic scale; the absolute latency difference (ms) increases with the number of queries per context.

Figure 8: GradMem can be efficiently scaled to longer contexts and larger models. We benchmark GradMem against (a) Llama-3.2-1B and (b) Llama-3-8B chosen as its base transformer model, as well as (a) TTT-Linear, TTT-MLP and LaCT. After repeated queries to its memory tokens, GradMem achieves lower overall latency compared to transformers decoding from full KV cache and test-time training methods reusing their state. Query length is 24 tokens, batch size is 1, GradMem has 8 memory tokens and

K=1

, measurements obtained on A100 GPU.

Figure 9: On KV-retrieval, GradMem learns to decrease inner loss on value tokens only. This figure was obtained for GradMem fine-tuned on 64-pair KV-retrieval with

K_{\text{train}}=2

Figure 10: More WRITE steps improve memory efficiency, but larger contexts still require larger memory. Exact Match accuracy on associative retrieval as a function of the number of memory tokens

m\in\{1,2,4,8\}

for two numbers of WRITE steps,

K\in\{1,2\}

, and two context sizes: 8 and 16 KV-pairs. All runs are initialized from the checkpoint trained with 8 memory tokens and then further trained after reducing the memory size. Increasing

K

consistently shifts the curve upward, especially with smaller memory size, but the memory size requirement still grows with the number of KV-pairs: for 8 KV-pairs,

K=2

reaches near-perfect performance already with

m=2

, whereas for 16 KV-pairs near-perfect performance requires at least

m=4

Model	Config	Total State Size	N8	N16	N32
Mamba	1L state 4 conv 4	4096	69.2 $\pm$ 6.4	56.8 $\pm$ 6.9	12.4 $\pm$ 15.5
Mamba	2L state 2 conv 2	4096	60.6 $\pm$ 41.2	15.9 $\pm$ 9.1	4.9 $\pm$ 6.2
Mamba	4L state 4 conv 4	16384	96.6 $\pm$ 0.5	82.9 $\pm$ 1.6	19.8 $\pm$ 19.3
Mamba	4L state 16 conv 4	40960	98.8 $\pm$ 1.5	98.7 $\pm$ 0.8	90.4 $\pm$ 4.7
GradMem	4L k=1	1024	99.7 $\pm$ 0.0	96.3 $\pm$ 0.9	86.9 $\pm$ 0.5
GradMem	4L k=2	1024	100.0 $\pm$ 0.0	99.6 $\pm$ 0.1	98.3 $\pm$ 0.1

Appendix

Table 7: Associative retrieval performance, exact match (%) (mean

\pm

std) across task lengths. All models have the same hidden state size of 128, the results are averaged across 3 runs.

	Number of KV-pairs
Model	8	16	32	64	96
TTT-Linear	97.2 $\pm$ 0.9	96.9 $\pm$ 0.7	82.1 $\pm$ 3.4	38.6 $\pm$ 14.4	11.9 $\pm$ 10.8
LaCT	95.7 $\pm$ 6.8	99.9 $\pm$ 0.1	99.7 $\pm$ 0.4	62.5 $\pm$ 25.2	46.9 $\pm$ 24.3
LaCT (best of 3)	99.9	99.9	99.9	91.6	74.9
GradMem ( $K=1$ )	99.7 $\pm$ 0.0	96.3 $\pm$ 0.9	86.9 $\pm$ 0.5	58.6 $\pm$ 0.7	32.6 $\pm$ 0.1
GradMem ( $K=5$ )	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	99.9 $\pm$ 0.1	99.1 $\pm$ 0.3	88.4 $\pm$ 2.3

Appendix

Table 8: KV-retrieval comparison with TTT-Linear and LaCT. Exact match retrieval accuracy (mean

\pm

std over 3 runs unless otherwise noted) on the same associative KV-retrieval setup as Table 1 . LaCT has high variance, so we also report the best of 3 runs.

Model	$K$	$n_{\text{mem}}$	$N@95$	$N@99$
Llama-3.2-1B	1	1	8	8
	1	2	16	8
	2	1	32	16
	2	2	32	32
Llama-3.2-3B	1	1	16	8
	1	2	16	16

Appendix

Table 9: GradMem is applicable to larger pretrained LLMs: with only 1–2 memory vectors and 1–2 WRITE steps, it can compress and reconstruct short text spans while keeping the READ model frozen. Maximum sequence length

N

that reaches target token-accuracy thresholds on text reconstruction with GradMem. Higher is better. We evaluate Llama-3.2-1B and Llama-3.2-3B with

K\in\{1,2\}

WRITE steps and

n_{\text{mem}}\in\{1,2\}

memory vectors. During WRITE, the model uses learned LoRA adapters, while during READ the pretrained base model is frozen and used without LoRA. Sequence lengths are evaluated on a coarse grid

N\in\{8,16,32,64\}

, so the reported thresholds should be interpreted as approximate capacity estimates.

Figure 11: Input-level memory is more parameter-efficient than per-layer memory, while a much larger per-layer state can outperform it. Exact Match on associative KV retrieval for 8/16/32 KV-pairs using 4-layer models with hidden size 128. We compare GradMem’s input-level memory with 8 memory tokens (

8\times 128=1024

parameters) against per-layer memory implemented as test-time trainable KV-cache in two regimes: (i) matched total memory size, corresponding to KV-cache for one token across all 4 layers (

1\times 4\times 128\times 2=1024

), and (ii) matched KV-cache length for 8 tokens (

8\times 4\times 128\times 2=8192

), which is

8\times

larger. All variants use the same model-level reconstruction objective; only the parameterization of test-time trainable memory changes. With matched memory size, input-level memory is consistently stronger, especially at longer contexts; with the much larger per-layer state, per-layer memory becomes strongest.

Memory-Augmented Language Models

Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.

See all 291 papers

PReM: Learning What to Preserve and When to Refresh for Context Compression
Jul 15, 2026Bohan Yu, Lei Shen, Chenxi Zhou +5LLM Inference Acceleration Long-Context Language Model Inference
Alibaba Group
$δ$ -mem: Efficient Online Memory for Large Language Models
May 12, 2026Jingdi Lei, Di Zhang, Junxian Li +7Memory-Augmented Language Models Memorization in Language Models
Nanyang Technological University · Mind Lab · Fudan University +3
Context Distillation as Latent Memory Management
May 27, 2026Ziyang Zheng, Zeju Li, Xiangyu Wen +5Memory-Augmented Language Models Language Model Distillation
The Chinese University of Hong Kong · Huawei Noah’s Ark Lab

Jul 15, 2026J/K move · Enter open · S save

PReM: Learning What to Preserve and When to Refresh for Context Compression

Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han+4

Alibaba Group

LLM Inference Acceleration Long-Context Language Model Inference KV-Cache Compression Context Compression Memory-Augmented Language Models LLM Context Compression+1

Open paper AskPDF arXiv

Efficient long-context inference is not only about reducing memory cost, but also about keeping useful contextual evidence accessible as generation proceeds. However, existing compression-oriented approaches, such as key-value (KV) cache compression and context compression, often either make an early decision about which contextual information to keep or rely on an external compressor. Such designs make it difficult to adapt the compressed context to the evidence needed by later reasoning steps. This paper introduces PReM (Preserve and Refresh Memory), a context-compression framework that maintains the long context as the model's internal layer-wise KV memory and learns what to preserve and when to refresh it. Specifically, PReM uses a dedicated memory layer to make memory-selection decisions, and a special memory token <m> to trigger refreshes during generation. To train this behavior, PReM introduces Phase-Separated Refresh Training, aligning memory selection with memory-conditioned generation while preserving continuity across refreshes. Experiments with 32K-token contexts show that PReM outperforms strong baselines under both 16x and 32x compression, while maintaining a favorable balance between answer quality and inference efficiency.

Abstract

Memory-Augmented Language Models

Explore similar work

PReM: Learning What to Preserve and When to Refresh for Context Compression

δδδ-mem: Efficient Online Memory for Large Language Models

Context Distillation as Latent Memory Management

$δ$ -mem: Efficient Online Memory for Large Language Models