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

1 Introduction

• Preference-consistent behavior is distinct from remembering evidence: models may recover preference changes yet fail to constrain recommendations or stop acting on invalidated information [1]–[5].Implicit objections and condensed histories that lose preference signals are also identified as challenges [2] [3] [5].Source: §1 · p. 2
• Textual memory grows with retained history, while latent memory bounds the reader’s interface with a fixed number of continuous vectors [6–17].A fixed interface alone does not ensure that compressed memory supports the behavior enabled by text.Source: §1 · p. 2
• Reconstruction and reference-answer imitation score sequences the reader did not produce, while scalar task rewards do not identify which response decisions underused memory [18] [19] [20] [21] [22] [23] [24].A sequence-level reward can charge an entire response for a localized preference violation.Source: §1 · p. 2
• MemFold evaluates compression on the same sampled response by comparing its token likelihood under soft memory with that under a frozen textual-memory reader.The token-level differences weight the reader’s own tokens, complementing task reward with information about which tokens textual memory supports more strongly.Source: §1 · pp. 2–3
• MemFold optimizes a fixed-budget soft-memory reader on its own rollouts with confidence-gated on-policy distillation and reports the highest accuracy across four benchmarks and three backbones [20–24].The reported margins widen at longer histories, and transfer to unseen benchmarks requires no target-domain training.Source: §1 · p. 3

2 Related Work

• Long-horizon memory research studies textual stores and recall or multi-session reasoning, while personalization research tests whether retained information changes responses to particular users [6] [8] [9] [10] [11] [28] [29] [32] [33].Benchmarks also examine evolving preferences [1], implicit preferences [2] [3], and forgetting invalidated information [4].Source: §2 · p. 3
• Compact-memory methods shorten or encode context into continuous vectors, but commonly use reconstruction or imitation objectives scored on sequences the reader did not generate [16] [17] [18] [19] [35].Reinforcement learning instead optimizes memory decisions on model samples [20] [21] [22] [23] [24] [42] [43] [44] [45] [46].Source: §2 · p. 3
• MemFold holds the interface fixed and optimizes how it is read, using a frozen copy of the reader to rescore its own tokens under textual memory [49] [50].Unlike approaches using a stronger teacher and signed gap, it uses the same initialized reader with a different memory and a non-negative gate.Source: §2 · pp. 3–4

3 Methodology

• The on-policy loop uses student responses conditioned on query and compressed memory, teacher scoring and confidence gating of those tokens, then joint GRPO and distillation updates.The teacher and compressor remain frozen during the update.Source: Figure 2 · p. 4, §3 · p. 4

3.1 Problem Formulation

• A query-conditioned textual memory is compressed into Z, a fixed-size K×d continuous memory, and the reader generates from the query and Z.The compressor uses the first four Transformer blocks of the frozen backbone; K is the memory budget and d the reader’s input embedding dimension.Source: §3.1 · p. 5

3.2 Fixed-Budget Soft Memory

• A Perceiver-style compressor supplies an interface whose size is independent of history length and that the reader can consume.Initialization reconstructs textual memory and trains answer generation, but these objectives do not directly optimize the reader’s use of memory on its own generations.Source: §3.2 · p. 5

3.3 On-Policy Memory Optimization

• On-policy optimization addresses gaps left by initialization: reference-answer training misses errors in generated responses, and scalar task rewards do not locate token-level underuse of memory.The current reader samples responses, which receive task rewards for group-relative policy optimization.Source: §3.3 · p. 6
• The frozen textual-memory teacher scores student-sampled tokens, and a detached confidence gate strengthens tokens in proportion to the teacher’s relative support.The bounded gate never promotes tokens the student did not sample; task-outcome feedback comes from GRPO.Source: §3.3 · pp. 6–7

4.1 Experimental Setup

• PersonaMem-32K and PersonaMem-128K test dynamic preference tracking; PrefEval tests cross-dataset personalization transfer, while LongMemEval tests broader long-term memory reasoning without task-specific training.The three backbones are Qwen2.5-3B-Instruct, Qwen2.5-7B-Instruct, and Qwen3-4B.Source: §4.1 · p. 7
• Baselines cover GRPO and OPSD objectives, latent-memory methods AutoCompressor, MemGen, and xRAG, and Full Text as the uncompressed-context comparison.The implementations of xRAG, MemGen, and AutoCompressor are adapted implementations rather than exact reproductions of official checkpoints.Source: §4.1 · p. 7
• The study reports accuracy and average end-to-end token-equivalent count per test instance, using greedy decoding by default and source-domain validation for checkpoint selection.PrefEval and LongMemEval predictions are scored through the Gemini-3.8-flash API; LongMemEval also uses structured yes/no outputs.Source: §4.1 · p. 8

4.2 Results

• MemFold has highest accuracy on every benchmark and backbone, tying the strongest baseline in one PrefEval cell; its margin is larger at 128K than 32K across all three backbones.Average end-to-end cost stays close to full-context inference and is lower in aggregate for all three backbones, except on PrefEval, where memory-construction costs are not offset by reader-side savings.Source: Table 1 · p. 8, §4.2 · pp. 8–9

5 Ablation Study

• Reader initialization produces the largest initialization gain, while writer initialization adds a further gain by improving the textual memory used to construct the soft representation.Table 2 ablates initialization and on-policy objectives and contrasts textual-memory control with fixed-budget soft memory.Source: §5 · p. 9, Table 2 · p. 9
• GRPO accounts for most of the task improvement, and OPD adds a complementary gain through token-level guidance from the textual-memory teacher when combined with GRPO.Source: §5 · p. 9
• The text-space control performs comparably but does not outperform the full soft-memory model, suggesting fixed-budget compression is not the primary bottleneck under this training configuration.Source: §5 · p. 9

Out-of-domain transfer

• MemFold maintains more robust performance than the compared methods across PrefEval and LongMemEval, though transfer is not uniform and it ties the strongest baselines in one PrefEval setting.The comparison evaluates frozen PersonaMem-32K checkpoints on PrefEval and checkpoints trained and selected only on LoCoMo on LongMemEval.Source: §6 · p. 10

Training efficiency

• MemFold improves accuracy more rapidly with optimizer updates and reaches comparable accuracy with fewer student rollouts in the Qwen2.5-3B-Instruct PersonaMem-32K experiment.Reward and distillation reuse the same responses, while the teacher scores sampled tokens without autoregressive generation; the comparison reflects complete configurations and excludes initialization and teacher forward passes from total training cost.Source: §6 · p. 10

Budget sensitivity

• Both accuracy curves peak at K = 256 and remain nearly flat from 128 upward, while the smallest budget is clearly under-provisioned.The decline at K = 512 is interpreted as a property of the fixed-budget initialization recipe, not evidence that more vectors inherently carry less.Source: §6 · p. 10, Figure 3 · p. 10

Cost of the budget

• End-to-end token cost varies by less than one percent across budgets on the longer benchmark, because most variation comes from generated tokens.Every configuration still reads the full interaction history once to build textual memory, so enlarging the budget costs little and shrinking it saves little.Source: Figure 4 · p. 11, §6 · p. 11

Memory use

• Shuffled or null memory causes substantial accuracy drops on both datasets, indicating that the reader depends on instance-specific memory content.The intervention does not establish which memory parts are used or whether the reader resolves the correct time-dependent preference.Source: Figure 5 · p. 11, §6 · p. 11

Localized personalization failures

• In one PrefEval example, GRPO and OPSD recommend wearable trackers despite the user’s preference, whereas MemFold suggests preference-consistent alternatives.The example illustrates a failure pattern, not its frequency.Source: §6 · pp. 11–12, Figure 6 · p. 12

7 Conclusion

• MemFold judges compact personalized memory by the generations it supports, and the conclusion reports strong benchmark accuracy, transfer, and dependence on instance-specific memory.Across three backbones, it attains the highest measured accuracy on PersonaMem-32K and PersonaMem-128K, transfers to PrefEval and LongMemEval without target-domain training, and depends on instance-specific memory under shuffled- and null-memory interventions.Source: §7 · p. 12

A.1.1 Textual Memory Construction

• The external extractor builds query-relevant textual memory from history and question only, retaining grounded evidence, temporal relations, and derived facts without access to the reference answer.The adapter is trained on context–query–memory examples and subsequently writes the memory without the external extractor; Figure 7 shows a representative template subset.Source: §A.1.1 · p. 18, Figure 7 · p. 18

A.1.2 Soft-Memory Construction and Initialization

• The soft-memory interface uses a two-layer Perceiver-style compressor with K learned queries and a projector, and its default budget is K = 256.Its input is produced by mean-pooling every 32 consecutive hidden states from chunks of at most 2,048 tokens; separate compressors are trained for each dataset–backbone configuration.Source: §A.1.2 · p. 19
• Initialization proceeds through compressor reconstruction, representation warmup, auxiliary reasoning adaptation, and reader initialization, each with a distinct objective.Reconstruction uses no answer or reasoning supervision; warmup combines cross-context separation, same-context view alignment, and prototype decorrelation.Source: §A.1.2 · pp. 19–21
• Auxiliary reasoning adaptation updates the compressor and projector with evidence-grounded reasoning targets, then reader initialization trains answer generation and uses the resulting adapter as the textual-memory teacher.The temporary rank-8 LoRA used for reasoning adaptation is discarded; the reader-initialization objective also includes a soft-output anchoring loss.Source: §A.1.2 · pp. 21–22

A.1.3 On-Policy Training Details

• On-policy training uses fixed cached memories: the teacher reads textual memory while the student receives its compressed representation from the same memory.The training memories are not refreshed after policy updates, while evaluation memories are generated by the final adapter.Source: §A.1.3 · p. 22
• The reader is optimized on student-generated responses using group-relative task rewards and teacher-scored token log-probabilities, with the OPD mask covering response tokens through the first EOS.PersonaMem rewards exact parsed-option matches; groups with effectively identical rewards receive zero GRPO advantage, while the teacher-score implementation detaches both logged quantities.Source: §A.1.3 · pp. 22–23

A.2 Training Configuration

• MemFold uses a fixed soft-memory budget of 256 tokens and extracts textual-memory representations from the fourth Transformer block unless otherwise specified.Table 4 summarizes shared architectural and algorithmic hyperparameters.Source: §A.2 · p. 23, Table 4 · p. 23
• The procedures use AdamW with weight decay 0.01 and global gradient clipping at 1.0, with procedure-specific learning-rate schedules and auxiliary-loss weights.Memory-writer initialization uses cosine decay and 5% warmup; later procedures use constant learning rates without warmup, while representation warmup weights its three terms at 0.1 each.Source: §A.2 · p. 23, Table 5 · p. 24
• Writer initialization trains a LoRA adapter, reader initialization continues it while updating the final resampler layer, output normalization, and projector, and on-policy optimization updates only the adapter.The backbone, remaining compressor parameters, and textual-memory teacher stay frozen during on-policy optimization; the temporary rank-8 LoRA for auxiliary reasoning adaptation is discarded.Source: §A.2 · pp. 23–24
• PersonaMem rollouts use eight responses per query with λOPD = 0.02 and λGRPO = 0.3, while LoCoMo uses four responses and coefficient 1.0 for both objectives.The configurations also differ in temperature, top-p, and maximum response length; neither uses additional reference-policy or KL regularization.Source: §A.2 · p. 24
• Experiments used NVIDIA H200 GPUs, and the PersonaMem-32K/Qwen2.5-3B efficiency run completed 369 updates in about 12.9 minutes on four GPUs.The reported wall-clock time excludes evaluation and preceding initialization procedures.Source: §A.2 · p. 24

A.3 Datasets & Benchmarks

• PersonaMem evaluates evolving user-preference tracking with 50 test questions for 32K contexts and 233 for 128K contexts; PrefEval uses 1,000 implicit-persona examples, and LongMemEval-S uses 500 questions.PrefEval predictions are scored with the Gemini API; LongMemEval answers use gemini-3.8-flash with LOW thinking and structured yes/no outputs.Source: §A.3 · pp. 24–25

A.4 Baseline Implementation Details

• Baselines share the backbone, source-domain splits, visible-history boundary, and within-setting evaluation protocol, with selection based only on source-domain validation and parameters frozen for OOD evaluation.xRAG, MemGen, and AutoCompressor are adapted implementations rather than exact reproductions of official checkpoints.Source: §A.4 · p. 25
• The xRAG adaptation retrieves individual visible-history messages with a frozen GTE-large encoder and trains a two-layer MLP projector while keeping the retriever and Qwen backbone frozen.Projector training first reconstructs source-domain context messages and then fine-tunes on source-domain QA.Source: §A.4 · p. 25
• MemGen uses the Weaver core adapted to Qwen, while AutoCompressor recurrently carries learned summary vectors across history segments.MemGen uses eight prompt latents and eight inference latents, and its reported PersonaMem and PrefEval checkpoints disable Trigger; AutoCompressor produces 32 summary vectors per 1,536-token segment.Source: §A.4 · p. 25
• OPSD is implemented as a full-history QA baseline using generalized Jensen–Shannon divergence with point-wise clipping and no task reward.Its frozen teacher may access the reference answer during source-domain training, but reference answers are unavailable during validation, OOD evaluation, and test inference.Source: §A.4 · pp. 25–26
• Vanilla GRPO is a full-history baseline without teacher-distribution supervision, using group-normalized task rewards from eight sampled completions per source-domain question.It uses strict multiple-choice exact match on PersonaMem and normalized token F1 for open-ended source-domain QA.Source: §A.4 · p. 26

B.1 Efficiency Comparison Setup

• Figure 1 compares five complete method configurations, not a controlled loss-only ablation, using Qwen2.5-3B-Instruct on PersonaMem-32K.Each method is trained for 369 optimizer updates; checkpoints are evaluated on the same 50 test questions with deterministic answer-option permutations, with one training run per method.Source: §B.1 · pp. 26–27
• At 369 updates, GRPO, GRPO+OPD, SDPO, and MemFold each have 5,904 student rollouts, while OPSD has 5,860.Repeated sampling counts as a new rollout; teacher forward passes do not, and groups with identical rewards still count toward rollouts.Source: §B.1 · p. 27

B.2 End-to-End Token-Equivalent Accounting

• Token-equivalent processing counts discrete tokens processed or generated at each inference stage, treating each soft, summary, or latent position as one token equivalent.It excludes training-only annotation, teacher computation, and evaluation-judge tokens; it measures logical processing rather than FLOPs or latency.Source: §B.2 · p. 27, Table 7 · p. 28

C Theoretical Analysis of LOPD

• The theoretical analysis defines the sampled-token confidence gap as teacher minus student log probability and sets the gate to a sigmoid of that gap.It holds sampled trajectories and masks fixed when distinguishing batch gradients from conditional expectations.Source: C Theoretical Analysis of LOPD · p. 27

C.1 Gradient of the Sampled Surrogate

• With samples, masks, gate, and teacher scores fixed in the backward pass, the surrogate update is gate-weighted negative log-likelihood on student-generated tokens.There is no gradient through the teacher or gate, and this is a fixed-sample surrogate gradient rather than a total derivative through the rollout distribution.Source: §C.1 · p. 28
• Tokens the teacher scores higher receive larger weights; equal log probabilities yield weight 1/2, while tokens the student scores higher receive weights approaching zero.These weights encode relative confidence on sampled tokens, not token correctness or the identity of information lost in compression.Source: §C.1 · p. 28

C.2 Relationship to Reverse KL

• For a fixed prefix, OPD uses bounded, nonnegative per-sample weights instead of reverse KL’s unbounded signed log-ratio, and the expected update is stationary when student and teacher distributions match.The reverse-KL relationship assumes a finite vocabulary and strictly positive probabilities.Source: §C.2 · p. 29, §C.3 · p. 29

C.3 Expected Update: Stationarity and Bounded Influence

• In expectation, OPD replaces the reverse-KL log-ratio with a bounded, sign-preserving coefficient, behaving like scaled reverse KL for small gaps and capping each token’s influence for large gaps.The teacher serves as a reference reader, not an accuracy oracle; OPD re-ranks student candidates rather than importing teacher choices the student did not sample.Source: §C.3 · pp. 30–31

C.4 Remarks on the Implemented Estimator

• The theoretical guarantees are local and conditional: they do not ensure training reaches student-teacher agreement, monotone attenuation, or that the OPD coefficient needs no tuning.Under tempered or nucleus sampling, the estimator’s baseline can act as self-imitation toward the truncated sampling distribution; PersonaMem uses temperature 1.0 and top-p 0.98, while LoCoMo uses 0.8 and 0.95.Source: §C.4 · p. 31

• The memory writer is initialized from memories extracted by a stronger external model, so initial textual-memory quality depends on that model.Source: D Limitations · p. 31
• Learning the memory writer without stronger-model supervision is left to future work.Source: D Limitations · p. 31
• Experiments cover Qwen backbones only up to 7B parameters; applicability to larger models and other model families remains unverified.Source: D Limitations · p. 31