Organizations: AXXX, Cognitive AI Systems Lab, Moscow, Russia · MIRAI, Moscow, Russia · MBZUAI, Abu Dhabi, UAE · London Institute for Mathematical Sciences, London, UK
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θ compresses C into a fixed-size memory M^ in a WRITE phase, and the model predicts Y from [M^;Q] in a READ phase, without access to C . (b) Meta-learning view: a shared initialization M0 and model parameters θ are learned across training examples (outer loop), while at test time each context Ci adapts its own memory M^(Ci) via a few gradient steps (dotted trajectory). (c) Test-time gradient descent on memory. Starting from the meta-learned initialization M0 , GradMem updates the per-sample memory state during WRITE with K steps of gradient descent on the context reconstruction loss Lwrite(M;C) . At READ, the model predicts the task target using only [M^;Q] . During training, θ and M0 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: M =KV-cache
100.0 ± 0.0
100.0 ± 0.0
99.8 ± 0.0
99.8 ± 0.3
96.5 ± 2.9
98.8 ± 0.0
Mamba: M =per-layer recurrent state
99.9 ± 0.1
98.9 ± 0.4
98.7 ± 0.2
90.2 ± 10.2
95.2 ± 0.1
92.2 ± 0.4
ARMT: M =per-layer associative matrix
99.0 ± 0.3
98.5 ± 0.5
97.4 ± 0.3
54.9 ± 2.1
22.6 ± 3.9
15.2 ± 0.2
Forward-only write (RMT): M =8 mem vectors
100.0 ± 0.0
100.0 ± 0.0
45.5 ± 0.2
44.3 ± 0.0
19.3 ± 3.1
12.9 ± 0.2
x2 memory updates
69.6 ± 28.1
18.7 ± 3.4
–
–
Table 1: KV-retrieval: gradient-based WRITE outperforms forward-only WRITE. Exact match retrieval accuracy (mean ± 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 ↑ )
SQuAD (EM ↑ )
LM (CE ↓ )
QA1
QA2
QA3
QA4
QA5
short
128–255
Input length (tokens)
∼ 40
∼ 100
∼ 300
∼ 20
∼ 20
∼ 40
256
Full context models (upper bound)
GPT-2-124m
100.0 ± 0.0
100.0 ± 0.0
99.8 ± 0.1
100.0 ± 0.0
99.4 ± 0.1
64.2 ± 0.3
2.72 ± 0.00
Limit context to 128 tokens
3.20 ± 0.00
Pythia-160m
100.0 ± 0.0
99.7 ± 0.1
95.5 ± 2.7
100.0 ± 0.0
99.0 ± 0.1
48.9 ± 0.4
2.84 ± 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 ± std over 3 runs.
Figure 4: More gradient steps at test-time lead to better performance without fine-tuning. (a) Results for Ktrain 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 xt (or token mini-batches)
Whole context segment C (WRITE once per context)
Test-time parameters
Layer-specific parameters Wt (updated from the layer’s inputs/activations), per layer
Prefix memory tokens Mk∈Rm×d (single memory state), per model
Self-supervised loss
Activation/input reconstruction, e.g. ℓ(W;xt)=∥f(x~t;W)−xt∥22 (or learned multi-view projections)
Context reconstruction Lwrite(M;C)
Outer-loop objective
Next-token prediction (LM training)
Downstream task loss with C removed at READ
Outer-loop parameters
Model params θ + reconstruction task/view params
Model params θ + memory init M0 (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 xt / views of xt
Key–value proj. (kt,vt) of token xt
Local window of keys and values
Chunk of keys and values
Full context C
Test-time parameters
Weights Wt , per layer
MLP mem. + momentum, per layer
MLP mem. + momentum + gates, per layer
Weights Wt , per layer
Single memory prefix M∈Rm×d
WRITE objective
View reconst., L2
KV recall, L2
Window KV recall, L2
Chunk KV recall, dot product
Context reconst. Lwrite(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 darmt_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 )
dmem
K
α
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). α is the inner learning rate for GradMem.
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.
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 Ktrain=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∈{1,2,4,8} for two numbers of WRITE steps, K∈{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 ± 6.4
56.8 ± 6.9
12.4 ± 15.5
Mamba
2L state 2 conv 2
4096
60.6 ± 41.2
15.9 ± 9.1
4.9 ± 6.2
Mamba
4L state 4 conv 4
16384
96.6 ± 0.5
82.9 ± 1.6
19.8 ± 19.3
Mamba
4L state 16 conv 4
40960
98.8 ± 1.5
98.7 ± 0.8
90.4 ± 4.7
GradMem
4L k=1
1024
99.7 ± 0.0
96.3 ± 0.9
86.9 ± 0.5
GradMem
4L k=2
1024
100.0 ± 0.0
99.6 ± 0.1
98.3 ± 0.1
Appendix
Table 7: Associative retrieval performance, exact match (%) (mean ± 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 ± 0.9
96.9 ± 0.7
82.1 ± 3.4
38.6 ± 14.4
11.9 ± 10.8
LaCT
95.7 ± 6.8
99.9 ± 0.1
99.7 ± 0.4
62.5 ± 25.2
46.9 ± 24.3
LaCT (best of 3)
99.9
99.9
99.9
91.6
74.9
GradMem ( K=1 )
99.7 ± 0.0
96.3 ± 0.9
86.9 ± 0.5
58.6 ± 0.7
32.6 ± 0.1
GradMem ( K=5 )
100.0 ± 0.0
100.0 ± 0.0
99.9 ± 0.1
99.1 ± 0.3
88.4 ± 2.3
Appendix
Table 8: KV-retrieval comparison with TTT-Linear and LaCT. Exact match retrieval accuracy (mean ± 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
nmem
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∈{1,2} WRITE steps and nmem∈{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∈{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×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×4×128×2=1024 ), and (ii) matched KV-cache length for 8 tokens ( 8×4×128×2=8192 ), which is 8× 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 combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.
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.
Large language models increasingly need to accumulate and reuse historical information in long-term assistants and agent systems. Simply expanding the context window is costly and often fails to ensure effective context utilization. We propose δ-mem, a lightweight memory mechanism that augments a frozen full-attention backbone with a compact online state of associative memory. δ-mem compresses past information into a fixed-size state matrix updated by delta-rule learning, and uses its readout to generate low-rank corrections to the backbone's attention computation during generation. With only an 8×8 online memory state, δ-mem improves the average score to 1.10× that of the frozen backbone and 1.15× that of the strongest non-δ-mem memory baseline. It achieves larger gains on memory-heavy benchmarks, reaching 1.31× on MemoryAgentBench and 1.20× on LoCoMo, while largely preserving general capabilities. These results show that effective memory can be realized through a compact online state directly coupled with attention computation, without full fine-tuning, backbone replacement, or explicit context extension.
Jingdi Lei, Di Zhang, Junxian Li +7
Nanyang Technological University · Mind Lab · Fudan University +3
Context distillation compresses contextual information into model parameters, yet existing methods often ignore how multiple distilled latent memories should be stored, retrieved, and safely activated in non-oracle settings. We formulate context distillation as a latent memory management problem. We distill each context into an independent LoRA adapter, forming a modular memory bank that enables explicit memory selection. Given a query, our framework retrieves candidate memories, routes the query to the most suitable adapter, and uses a Self-Gating mechanism to decide whether latent memory should be activated. To improve efficiency, we further introduce cache sharing to reduce management overhead during inference. Experiments show that our method substantially outperforms baselines with retrieval, while Self-Gating improves robustness by deactivate unnecessary latent memories.
Ziyang Zheng, Zeju Li, Xiangyu Wen +5
The Chinese University of Hong Kong · Huawei Noah’s Ark Lab