MemFold: Learning Compact Soft Memory for Long-Context Personalization via On-Policy Optimization

Jingxuan Wu; Yuzhe Yang; Yiqiao Huang; Chengzhi Liu; Qingni Wang; Chengxuan Qian; Shutong Wu; Jiawei Zhang

cs.CLSep 29, 2026

MemFold: Learning Compact Soft Memory for Long-Context Personalization via On-Policy Optimization

Authors: Jingxuan Wu, Yuzhe Yang, Yiqiao Huang, Chengzhi Liu, Qingni Wang, Chengxuan Qian, Shutong Wu, Jiawei Zhang,

Organizations: University of North Carolina at Chapel Hill · University of California, Santa Barbara · Harvard University · University of Wisconsin–Madison

Abstract

An assistant that serves the same user over a long horizon has to answer from what that user has revealed: which preferences still hold, which were revised, and which constraints apply now. Retaining that information is not the same as acting on it, and the two are usually optimized as if they were. Keeping the information as text makes the reader's input grow with the retained history, while compressing it into a fixed number of latent vectors bounds the interface but is typically trained to reconstruct text or imitate reference answers, both of which are scored on sequences the reader never produced. We present MemFold, which optimizes a fixed-budget soft memory by the behavior it supports. A query-conditioned textual memory is compressed into K continuous vectors that form the reader's memory interface, and the reader is then trained on its own rollouts under two complementary signals: group-relative rewards for task outcomes, and confidence-gated on-policy distillation in which a frozen textual-memory teacher re-scores the student's sampled tokens under the textual memory. The teacher is never sampled from, so supervision stays on the student's current distribution and adds no autoregressive decoding; at inference it is removed entirely. Across three Qwen backbones, MemFold attains the highest accuracy we measure on PersonaMem-32K and PersonaMem-128K, with margins that widen at the longer history length, and transfers to PrefEval and LongMemEval without target-domain training. Ablations attribute most of the task gain to the reward term and a smaller additional gain to the teacher signal, and memory interventions show that the reader depends on the instance-specific content of its soft memory.

Figures & tables

Figure 2: ❶ On-policy rollout: The student generates responses conditioned on the query and $K$ compressed soft memory tokens. ❷ On-policy distillation: A frozen textual-memory teacher scores the student’s tokens, and a detached teacher–student confidence gate $\bar{g}_{i,t}$ weights the student’s own tokens. ❸ Policy optimization: GRPO task rewards and the gated distillation term jointly update the student, while the teacher and compressor remain frozen.

Panel of Table 1 — Table 1: Main results for in-domain and direct cross-dataset evaluation, reporting accuracy and end-to-end token-equivalent counts. Weighted averages use benchmark sample counts. Best and second-best accuracies are highlighted for each backbone and benchmark.

Method	Initialization		On-Policy Objectives		Memory Interface	PersonaMem-32K		PersonaMem-128K
Method	Writer	Reader	OPD	GRPO	Memory Interface	Mean $\uparrow$	Pass@16 $\uparrow$	Mean $\uparrow$	Pass@16 $\uparrow$
w/o Writer Initialization	✗	✓	✓	✓	Soft	70.8	82.0	80.0	85.7
w/o Reader Initialization	✓	✗	✓	✓	Soft	47.9	76.0	46.2	67.0
w/o OPD	✓	✓	✗	✓	Soft	71.9	84.0	86.9	91.8
w/o GRPO	✓	✓	✓	✗	Soft	58.6	78.0	70.2	82.3
Initialization Only	✓	✓	✗	✗	Soft	60.5	82.0	70.8	85.1

Table 2: Ablation of initialization and on-policy objectives. The text-space control uses textual memory, whereas MemFold uses fixed-budget soft memory. ✓ = satisfies, ✗ = does not satisfy, and N/A = not applicable.

Figure 3: Accuracy across budgets and history lengths (Qwen3-4B).

Figure 4: Token cost relative to

K=256

(Qwen3-4B); labels show absolute counts.

Figure 5: Accuracy across memory conditions and the text teacher (Qwen3-4B).

Panel of Figure 6 — Figure 6: A PrefEval implicit-persona example with Qwen3-4B. The displayed preference summarizes the interaction history and is not given explicitly to the models. GRPO and OPSD recommend wearable trackers, whereas MemFold respects the preference by suggesting non-wearable alternatives. Colored spans mark preference-relevant content. Token counts cover generated answers only; this example illustrates a failure mode, not its frequency or end-to-end efficiency.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: A representative structured textual memory extracted from a long PersonaMem interaction history. For readability, we show a representative subset of the extracted entries.

Backbone	Layer	Dev. Score	Update Top-1	Update AUC	Margin
Qwen2.5-3B	L1	98.98	96.12	97.55	0.207
	L4	100.00	100.00	100.00	0.319
	L12	94.39	79.61	90.37	0.091
	L24	86.13	67.96	79.04	0.042
	L36 (final)	77.49	37.86	51.76	$-0.041$
Qwen3-4B	L1	95.92	93.20	96.35	0.148

Appendix

Table 3: Hidden-state layer selection. Dev. Score is the mean development AUC over value and update discrimination; the remaining metrics are measured on the held-out update task. Bold indicates the selected layer.

Memory representation and compressor
Hyperparameter	Setting
Soft-memory budget $K$	256 soft tokens
Mean-pooling window	32 hidden states
Compressor depth	2 layers
Latent dimension	768
Attention heads	12

Appendix

Table 4: Shared architecture and optimization hyperparameters used by MemFold.

Procedure	Learning Rate	Duration	Updated Modules
Memory-interface initialization
Memory-writer initialization	$1\times 10^{-5}$	3 epochs	Adapter
Compressor reconstruction	$3\times 10^{-4}$	25 updates	Compressor, projector
Representation warmup	$1\times 10^{-4}$	100 updates	Compressor, projector
Auxiliary reasoning adaptation	$1\times 10^{-4}\to$ $5\times 10^{-5}/3\times 10^{-5}$	3 epochs	Temporary LoRA $\omega$ , compressor, projector
Reader initialization	$1\times 10^{-5}$	3 epochs	Adapter, selected compressor modules, projector

Appendix

Table 5: Training configuration for PersonaMem-32K.

Method	Student Input	Initialization	Supervision	Loss
GRPO	Full interaction history	Base model	Task reward	$\mathcal{L}_{\mathrm{GRPO}}$
OPSD	Full interaction history	Base model	On-policy self-distillation	$\mathcal{L}_{\mathrm{OPSD}}$
SDPO	Full interaction history	Base model	Successful-rollout feedback; EMA teacher	$\sum_{t}D_{\mathrm{KL}}(p_{t}\\|q_{t})$
GRPO+OPD	Full interaction history	Base model	Task reward; textual-memory teacher	$\begin{aligned} &\lambda_{\mathrm{GRPO}}\mathcal{L}_{\mathrm{GRPO}}\\[-2.0pt] &{}+\lambda_{\mathrm{OPD}}\mathcal{L}_{\mathrm{OPD}}\end{aligned}$
MemFold	$K=256$ soft-memory vectors	Reader-initialization checkpoint	Task reward; textual-memory teacher	$\begin{aligned} &\lambda_{\mathrm{GRPO}}\mathcal{L}_{\mathrm{GRPO}}\\[-2.0pt] &{}+\lambda_{\mathrm{OPD}}\mathcal{L}_{\mathrm{OPD}}\end{aligned}$

Appendix

Table 6: Configurations compared in Figure 1 .

Method	#Tok.	Accounting Note
Full Text / GRPO / OPSD	$P+A$	$P$ contains the full visible history, query, and instructions.
xRAG	$T_{\mathrm{retriever}}+P+K_{\mathrm{xRAG}}+A$	Counts retriever inputs and projected retrieval positions consumed by the reader.
AutoCompressor	$\sum_{s}T_{\mathrm{AC}}^{(s)}+P+K_{\mathrm{AC}}+A$	Each compression pass includes processed text and carried summary positions.
MemGen	$T_{\mathrm{Weaver}}+K_{\mathrm{Weaver}}+T_{\mathrm{reasoner}}+K_{\mathrm{reasoner}}+A$	Weaver and reasoner computation are counted separately; repeated history reads are retained.
MemFold	$T_{\mathrm{writer\text{-}in}}+T_{\mathrm{memory\text{-}out}}+T_{\mathrm{compressor\text{-}in}}+P+K+A$	Counts memory generation and subsequent compressor encoding as separate operations.

Appendix

Table 7: End-to-end token-equivalent accounting by method.

Persistent Memory for Language Models

Persistent memory for language models studies how models retain, update, and use information across interactions and beyond their context windows.

See all 272 papers

HasMem: Hard-Origin Adaptively Softened Memory for Long-Term LLM Agents
Sep 25, 2026Zihong He, Junxiao Shen, Chen Liang +1Memory Management for LLM Agents Long-Term Conversational Memory
The Hong Kong University of Science and Technology (Guangzhou) · University of Bristol · The Hong Kong University of Science and Technology
Latent Personal Memory: Represent personal memory as dynamic soft prompts
Jun 18, 2026Debrup Das, Avinash Amballa, Yashas Malur Saidutta +3LLM Inference Efficiency Memory-Efficient Fine-Tuning
University of Massachusetts Amherst · Samsung Research America
AdaMem: Learning What to Remember with Adaptive Memory Policies for Personalized Agents
Jun 19, 2026Xingyu Chen, Rui Wang, Zhaopeng Tu +1LLM Agent Memory Persistent Memory for Language Models
Shanghai Jiao Tong University, Tencent · Shanghai Jiao Tong University · Tencent

Sep 25, 2026J/K move · Enter open · S save

HasMem: Hard-Origin Adaptively Softened Memory for Long-Term LLM Agents

Zihong He, Junxiao Shen, Chen Liang, Hai-Ning Liang

The Hong Kong University of Science and Technology (Guangzhou) · University of Bristol · The Hong Kong University of Science and Technology

Memory Management for LLM Agents Long-Term Conversational Memory Agent Memory Management Memory-Augmented Language Model Agents

Open paper AskPDF arXiv

Text-based memory and context compression support reuse of past interactions. Resizing continuous memory changes the input to a frozen LLM, coupling capacity allocation with readout. We propose Hard-Origin Adaptively Softened Memory (HasMem). Frozen hard-prompt embeddings provide a verifiable initial state. A controller adjusts memory widths, a Writer re-encodes resized entries, and Reader and Global provide readout adaptation and cross-turn state. On all

535

questions in a reconstruction probe derived from the Multi-Session Chat (MSC) development split, the main configuration achieves lexical F1 of

95.3

(

+4.4

percentage points) at

93.6\%

of the hard reference's framed memory positions. With approximately matched per-question target body budgets, six configurations at mean per-entry retention around

0.83

0.91

exceed rule-based re-encoding by

8.0

23.6

exact-match (EM) percentage points. With fixed model parameters and rule target width ratio

0.75

, Global's EM gain passes a user-level exact paired test with Bonferroni correction over eight comparisons. On all

500

LongMemEval-S questions, local lexical F1 rises from the hard reference's

3.4

8.9

, and answer negative log-likelihood (NLL) falls from

12.257

5.274

. F1 gains accompany lower EM on both evaluations.

Abstract

Persistent Memory for Language Models

Explore similar work

HasMem: Hard-Origin Adaptively Softened Memory for Long-Term LLM Agents

Latent Personal Memory: Represent personal memory as dynamic soft prompts

AdaMem: Learning What to Remember with Adaptive Memory Policies for Personalized Agents