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, Xin Eric Wang
cs.CL
2026-09-29
MemFold compresses chat history into 256 vectors and trains directly on the reader's own rollouts, lifting 128K personalization accuracy from ~24% to 94%.
A long-running assistant needs to track which of a user's past statements still hold, which were revised, and which no longer apply. Storing that history as text keeps it readable but makes the reader's input grow with every session. Compressing it into a fixed number of continuous vectors bounds the interface size, but prior compressors are trained to reconstruct the original text or imitate reference answers. Both objectives score the model on sequences it never actually generated, so nothing in training checks whether the compressed memory still changes what the reader says. The paper's example makes the gap concrete: a response can be fluent and on-topic while one recommendation inside it violates a preference the user stated earlier, and a single scalar reward over the whole response can't locate that one bad sentence.
MemFold has a writer, a compressor, and a reader, all sharing one LoRA adapter on a frozen backbone.
Because the gate is detached and the teacher never samples, the signal never asks the student to copy tokens it didn't produce itself; it only reweights the student's existing choices. At inference, the teacher and the extra distillation computation disappear entirely; the model answers from (x, Z) alone.
Across Qwen2.5-3B-Instruct, Qwen2.5-7B-Instruct, and Qwen3-4B, MemFold is the top-accuracy method on every benchmark/backbone combination tested, and the gap over baselines widens as history length grows from 32K to 128K tokens.
| Backbone | Setting | Full Text | Best baseline | MemFold |
| Qwen2.5-7B-Instruct | PersonaMem-128K | 24.0 | MemGen 78.5 | 94.4 |
| Qwen2.5-3B-Instruct | PersonaMem-128K | 4.3 | xRAG 65.2 | 89.4 |
| Qwen3-4B | PersonaMem-128K | 21.9 | MemGen 58.4 | 88.4 |
Full-text conditioning collapses at 128K, down to 4.3% for the 3B model, a sign that simply handing the model the whole history stops working at that length. On PrefEval and LongMemEval, benchmarks MemFold was never trained on, it still tops or ties the best baseline: 19.9% vs MemGen's 13.3% on PrefEval with Qwen3-4B, and 36.8% vs MemGen's 11.0% on LongMemEval with Qwen2.5-7B. Token cost stays close to the full-text baseline rather than dropping sharply, because the writer still has to read the entire history once to produce M. The gain is behavioral, not primarily a token-count saving.
Ablations show reader initialization matters most: removing it drops mean accuracy from 87.9 to 46.2 on PersonaMem-128K. Removing GRPO hurts more than removing OPD (87.9 → 70.2 vs. 87.9 → 86.9), so GRPO drives most of the task gain and OPD adds a smaller complementary one. A text-space control that runs the same training recipe on uncompressed textual memory performs comparably to MemFold but doesn't beat it, suggesting the fixed-budget compression itself isn't the main bottleneck in this setup. Swapping the matched memory for a shuffled or null one causes large accuracy drops, confirming the reader actually uses instance-specific content rather than guessing from task priors.
This gives teams building long-memory assistants a concrete alternative to reconstruction-style compressor training: judge a compressed memory by whether the trained reader gets the downstream answer right, not by whether the compressor can regenerate the source text. A fixed K-vector interface keeps the reader's cost flat as conversations grow, which matters most exactly where full-text conditioning is shown here to fail, around 128K tokens. The teacher-rescoring trick adds a token-level signal without extra autoregressive decoding, so it's cheaper to bolt onto an RL pipeline than standard distillation.
PrefEval histories are short enough that memory-construction overhead outweighs any token savings, and that's also where MemFold's margin over baselines is narrowest, a tie rather than a win on one backbone. The frozen teacher's own task accuracy is well below the final student's, and the paper leans on a bounded per-token gate to stop that inaccuracy from dominating updates, without directly measuring how often the teacher is wrong on the tokens it upweights. The memory-intervention test shows the reader depends on instance content, but the paper is explicit that it doesn't show the reader picks the correct, currently valid version of a preference when the history contains multiple revisions of it. The sweep over memory budget K is flat from 128 vectors up and dips at K=512, which the paper attributes to not retuning initialization per budget, an unresolved confound. Evaluation scoring relies on a Gemini model as judge, adding a judge-dependent source of error to every reported number.