Source-linked AI summary
Learning How and What to Memorize: Cognition-Inspired Two-Stage Optimization for Evolving Memory
Derong Xu, Shuochen Liu, Pengfei Luo, Pengyue Jia, Yingyi Zhang, Yi Wen, Yimin Deng, Wenlin Zhang, Enhong Chen, Xiangyu Zhao, Tong Xu
TL;DR
Limited context windows and static or sparsely rewarded memory updates make it difficult for LLM agents to track evolving preferences over long interactions. MemCoE learns a global memory guideline from contrastive feedback, then uses guideline-aligned rewards to train a multi-turn memory-update policy. Across three personalization benchmarks, it consistently outperforms strong baselines, with reported robustness to longer histories and noisier evidence.
Problem
Limited context windows hinder tracking evolving preferences, while static update rules and sparse outcome rewards leave memory-update policies difficult to adapt and optimize over long horizons.
Source: §1 · pp. 1–2Method
MemCoE first induces a global memory guideline from contrastive feedback, then uses guideline-aligned rewards and multi-turn reinforcement learning to train a memory-evolution policy.
Source: §1 · p. 2, §4.1 · p. 4, §4.2 · p. 5Results
Takeaways & Limitations
The findings support coupling an explicit evolution guideline with policy optimization for structured, consistent long-horizon personalization memory.
Source: §6 · p. 9The second-stage process rewards depend on an LLM-based scorer, making performance sensitive to scorer reliability.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
Large language model (LLM) agents require long-term user memory for consistent personalization, but limited context windows hinder tracking evolving preferences over long interactions. Existing memory systems mainly rely on static, hand-crafted update rules; although reinforcement learning (RL)-based agents learn memory updates, sparse outcome rewards provide weak supervision, resulting in unstable long-horizon optimization. Drawing on memory schema theory and the functional division between prefrontal regions and hippocampus regions, we introduce MemCoE, a cognition-inspired two-stage optimization framework that learns how memory should be organized and what information to update. In the first stage, we propose Memory Guideline Induction to optimize a global guideline via contrastive feedback interpreted as textual gradients; in the second stage, Guideline-Aligned Memory Policy Optimization uses the induced guideline to define structured process rewards and performs multi-turn RL to learn a guideline-following memory evolution policy. We evaluate on three personalization memory benchmarks, covering explicit/implicit preference and different sizes and noise, and observe consistent improvements over strong baselines with favorable robustness, transferability, and efficiency.
1 Introduction
Limited context windows and static update rules make it difficult for LLM agents to track evolving user preferences, while sparse rewards leave RL-based memory updates weakly constrained. MemCoE separates memory organization from content updates and consistently outperforms strong baselines across three personalization benchmarks.
Motivation
- Limited context windows impede use of full dialogue histories, and storing or retrieving snippets alone struggles to capture dynamic user patterns.External memory is motivated as a way to preserve important information over time for consistent personalization.Source: §1 · p. 1
- Static memory pipelines are difficult to adapt, while RL updates trained with sparse, delayed outcome rewards face large action spaces and unstable long-horizon optimization.These challenges motivate mechanisms for both memory organization and updating.Source: §1 · pp. 1–2
Approach
- Inspired by memory schema theory, MemCoE decouples how memory is organized from what content is stored and updated.The paper draws a functional analogy between prefrontal schema selection and hippocampal encoding of episodic details.Source: Figure 1 · p. 1, §1 · p. 2
- Memory Guideline Induction learns a global natural-language guideline from contrastive feedback, and Guideline-Aligned Memory Policy Optimization uses it to reward and train multi-turn memory updates.The guideline is intended to constrain the update policy's action space.Source: §1 · p. 2
Results
- MemCoE consistently outperforms strong baselines across PersonaMem, PrefEval, and PersonaBench, with reported robustness, transferability, and efficiency.The benchmarks span explicit and implicit preferences and increasingly noisy evidence sources.Source: §1 · p. 2
2 Related Work
Prior work builds memory banks using varied storage and retrieval structures, while personalization and experience-memory methods address specialized needs. The paper positions MemCoE as a general memory-evolution mechanism that complements these approaches and adds process-level guidance to RL-based updating.
Memory for LLM Agents
- Memory-bank research uses segmentation, summarization, compression, updating, and structured indices, while personalization and experience-memory methods target user profiles or past trajectories.The paper says its method is orthogonal to memory construction and storage and compatible with existing backends.Source: §2 · p. 3
RL for Memory
- RL-based memory methods optimize sequential memory operations, but typically rely on sparse final-answer rewards rather than process-level guidance for updates.MemCoE addresses this gap with guideline-aligned rewards.Source: §2 · p. 3
Prompt Optimization
- Unlike general prompt optimization, MGI refines a global memory-evolution guideline over long histories and connects it to Stage-2 RL through guideline-aligned process rewards.The paper describes contrastive diagnosis and batch aggregation as stabilizing guideline updates.Source: §2 · p. 3
3 Preliminary
The preliminary formulation treats user memory as a textual bank that evolves with dialogue and supports personalized responses. Its central challenge is designing coherent updates that preserve stable, accurate, temporally consistent user representations over long interactions.
Memory Formulation
- The system maintains a textual user memory bank that evolves as new behaviors and preferences emerge, alongside dialogue history and a learnable update prompt.The evolution operator incorporates new information, refines existing entries, and removes outdated or inconsistent content.Source: §3 · p. 3
- The formulation treats user memory as a continuously adapting latent structure aligned with the user's evolving profile.Source: §3 · p. 4
Problem
- The key challenge is designing memory evolution that stays coherent with dialogue history and supports stable, accurate, temporally consistent user representations.Responses are conditioned on both the task or query and the current memory state.Source: §3 · p. 4
4 Methodology
MemCoE learns a memory-evolution guideline and then trains a policy to update memory in accordance with it. The first stage uses contrastive textual feedback to refine the guideline; the second applies guideline-aligned rewards and multi-turn policy optimization.
4 Methodology
- MemCoE treats the memory-update instruction as an optimizable natural-language parameter, first learning how to perform memory operations and then what to store.The framework replaces hand-crafted update rules with a learned guideline and a policy that evolves memory in accordance with it.Source: §4 · p. 4, §4.1 · p. 4, Figure 2 · p. 5
4.1 Memory Guideline Induction
- MGI compares correct and suboptimal memory-augmented trajectories to produce textual feedback, then aggregates instance-level critiques into a guideline-level update.The resulting guideline is refined through natural-language editing and is intended to encode effective principles for downstream memory evolution.Source: §4.1 · p. 4
4.2 Guideline-Aligned Memory Policy Optimization
- GMPO fixes the induced guideline, combines a dense guideline-adherence reward with an answer-correctness reward, and uses multi-conversation GRPO to optimize the memory policy.The policy learns to follow the guideline while selectively storing information useful for downstream interaction quality.Source: §4.2 · p. 5
5 Experiments
Experiments compare MemCoE with retrieval-based and RL-based memory systems across three personalization benchmarks and eight evaluation settings. The reported results cover overall performance, ablations, efficiency, transfer, retrieval, token budgets, and guideline quality.
5.2 Overall Evaluation
- MemCoE achieves the best overall score across three benchmarks and consistently outperforms baselines across in-domain and out-of-domain settings.The comparison covers PersonaMem at 32K/128K, PrefEval Explicit/Implicit, and PersonaBench under different noise levels.Source: Table 1 · p. 6, §5.2 · pp. 6–7
5.1 Experimental Settings
- The evaluation spans three personalization benchmarks, compares retrieval-based and RL-based memory baselines, and reports accuracy on PersonaMem and PrefEval and F1 on PersonaBench.PersonaMem covers preference evolution over long histories; PrefEval distinguishes explicit and implicit preferences; PersonaBench uses heterogeneous, noisy user corpora.Source: §5.1 · p. 6
5.3 Ablation Study
- Removing MGI most hurt PrefEval preference retention, removing GMPO most hurt PersonaMem long-horizon tracking, and removing both collapsed performance.Removing contrastive feedback or guideline rewards caused smaller but consistent drops.Source: Table 2 · p. 6, §5.3 · p. 7
5.4 Efficiency Analysis
- On PersonaMem, MemCoE combines the best performance with relatively fast memory construction and evolution, rather than trading efficiency entirely for accuracy.The comparison covers 20 dialogue histories at 32K; circle size indicates runtime standard deviation.Source: Figure 3 · p. 7, §5.4 · p. 7
5.5 Cross-LLM Transferability of Guidelines
- Across four LLMs, both MGI variants outperform RAG and A-Mem, and guidelines optimized with gpt-4o-mini achieve the best numbers on three backbones.The table evaluates guidelines optimized with one LLM on different LLMs, without RL.Source: Table 3 · p. 7, §5.5 · pp. 7–8
5.6 Comparison on Different Retrieval
- Both MemCoE retrieval modes remain strong across Top-K values, while vanilla RAG degrades as K increases and can fall below Empty Memory.Ours (RAG) peaks around K=20 and can surpass the full-history variant.Source: Figure 4 · p. 8, §5.6 · p. 8
5.7 Impact of Per-Round Token Budget on Memory Evolution
- A moderate per-round token budget balances update frequency against the complexity of processing each evolution step.Budgets of 1K–2K can require many rounds and compound errors; 8K–32K can make each step more complex.Source: Figure 5 · p. 8, §5.7 · p. 8
5.8 Effect of Guideline Quality
- MGI’s induced guideline performs best on both PersonaMem settings, improving over the manual prompt by 10.4% at 32K and 11.3% at 128K.The reported gains are stable across three random seeds; an LLM rewrite of the manual prompt yields a moderate gain.Source: §5.8 · p. 8, Figure 6 · p. 9
6 Conclusion
MemCoE separates the guideline for organizing memory from the policy that decides what to retain, update, or forget. Experiments support improvements in performance, robustness, efficiency, and transferability for evolving user memory.
6 Conclusion
- MemCoE separates learning how memory should be organized from deciding what information to store.It first induces a transferable, schema-consistent guideline, then optimizes a guideline-aligned policy across multi-session interactions.Source: §6 · p. 9
- Across three personalization memory benchmarks, MemCoE consistently outperforms retrieval-based memory-bank and RL-based memory-agent baselines.The conclusion also reports robustness under longer histories and noisier evidence.Source: §6 · p. 9
- The reported results support coupling an explicit evolution guideline with policy optimization to improve efficiency, robustness, and transferability in conversational-agent memory.Source: §6 · p. 9
Limitations
The method has limits related to scorer reliability, evolution-round settings, and its single-objective policy design. The authors identify balancing competing objectives under a fixed guideline as an open challenge.
Limitations
- Second-stage process rewards rely on an LLM-based scorer, making performance sensitive to scorer reliability.Source: Limitations · p. 9
- The method requires careful tuning of per-round token budgets and evolution rounds; splitting long histories into many rounds can compound errors and cause unintended forgetting or over-generalized entries.Source: Limitations · p. 9
- The current design uses a single-objective policy under a fixed guideline, leaving trade-offs such as stability versus plasticity non-trivial to manage.The authors note that additional control mechanisms may be required.Source: Limitations · p. 9
Ethical considerations
The paper identifies misuse of persistent memory as its primary ethical risk and describes GRPO as optimizing multi-turn memory evolution over grouped trajectories.
Ethical considerations
- Persistent memory could enable over-collection or unwanted profiling without safeguards, so deployments should minimize data and provide user controls.The paper also recommends avoiding sensitive identifiers and enforcing retention policies and access control.Source: Ethical considerations · p. 9
A GRPO for Memory Evolution
- GRPO normalizes rewards within trajectory groups, then assigns each group-relative advantage to token actions across memory updates and answers.In the multi-conversation setting, each trajectory is decomposed into conversations.Source: A GRPO for Memory Evolution · p. 13
C.1 PersonaMem Dataset
The experiments use PersonaMem to test personalization across growing interaction histories, PrefEval to test preference retrieval amid distractors, and PersonaBench to test retrieval from heterogeneous personal corpora.
C.1 PersonaMem Dataset
- PersonaMem evaluates long-term personalization across 20 simulated personas with evolving traits and preferences in 15 domains.Histories combine 10, 20, or 60 sessions, yielding approximate context lengths of 32k, 128k, and 1M tokens.Source: §C.1 · pp. 13–14
C.2 PrefEval Dataset
- PrefEval tests preference inference, retrieval, following, and proactiveness using 1,000 explicit and 1,000 implicit instances with 50 intervening distractor turns.The experiments use its multiple-choice classification setting.Source: §C.2 · p. 14
C.3 PersonaBench Dataset
- PersonaBench evaluates personalized retrieval and question answering from conversations, assistant dialogues, and e-commerce histories for six users.Its short, underspecified queries require grounding in evidence distributed across those heterogeneous corpora.Source: §C.3 · p. 15
D Comparison in Different Categories
Across category breakdowns, MemCoE leads overall on PersonaMem, PrefEval, and PersonaBench, with weaker results in some categories and degradation as PersonaBench memory noise rises.
D Comparison in Different Categories
- 32.27 without noise and 29.89, 25.99, and 25.09 at noise levels 0.3, 0.5, and 0.7: MemCoE leads overall on PersonaBench at every reported level.All methods degrade as noise rises; MemCoE remains particularly strong on Pref. (Easy), while its advantage narrows on Pref. (Hard) and Social.Source: D Comparison in Different Categories · pp. 15–16
- 57.06 vs. 53.58 for MemAgent: MemCoE has the best overall PersonaMem score at 32K, with gains concentrated in memory-dependent personalization categories.At 128K, MemCoE scores 47.24 and leads on Latest prefs (54.44), Prefs evolve (66.47), Aligned recs (51.61), and New Scenarios (38.35).Source: D Comparison in Different Categories · p. 15, Table 8 · p. 16
- 81.30 explicit and 69.90 implicit: MemCoE achieves the best overall accuracy in both PrefEval settings.Its exception is Pet under Explicit Memory, where MemAgent scores 74.42 against MemCoE’s 67.44.Source: D Comparison in Different Categories · p. 15, Table 9 · p. 16
E.1 Comparison with Post-Training Baselines
MemCoE outperforms QA-focused post-training on PersonaMem, retains explicit preferences better over multiple rounds, and most remaining errors involve failures in both memory evolution and answer generation.
E.1 Comparison with Post-Training Baselines
- 57.06: MemCoE scores highest overall on PersonaMem, ahead of the best reported PPO/GRPO score of 50.31 among post-training baselines.The post-training sequence improves from Frozen (34.36) through SFT (46.83) to PPO/GRPO (49.49/50.31), but gains are uneven across personalization skills.Source: §E.1 · p. 16
E.2 Evaluation of Preference Retention
- About 74% vs. roughly 51%: at round 10, MemCoE retains more of the inserted preference than MemAgent.The retention test uses PrefEval (Explicit), a fixed 4K-token dialogue context per round, and Gemini-2.5-Pro as judge.Source: §E.2 · p. 17, Figure 7 · p. 18
E.3 Error Analysis
- 15.6% of cases are FE+FG, the largest failure category, where preference capture and final response generation both fail.SE+SG accounts for 77.1%; the three failure categories together make up 22.9%.Source: Figure 8 · p. 18, §E.3 · p. 18
E.4 Effect of Training Steps on RL-Based Baselines
RL-based memory baselines plateau while MemCoE continues to outperform their best checkpoints; further analyses show scaling behavior, an intermediate optimum for guideline induction, and an advantage over general prompt optimization. Case studies and prompt designs illustrate how the method handles memory evolution and answer generation.
E.5 Scaling Analysis
- From 4K to 128K dialogue tokens, the memory bank grows sublinearly (~500 to ~2,000), while per-round evolution time rises approximately linearly.Each evolution round processes 4K tokens; memory growth gradually flattens as dialogue history lengthens.Source: Figure 9 · p. 19, §E.5 · pp. 19–20
E.4 Effect of Training Steps on RL-Based Baselines
- MemAgent peaks at 53.58 (steps 100/140) and MEM-α at 53.37 (step 120), then plateau; MemCoE beats both best checkpoints throughout.The study used 300 sampled PersonaMem training examples, batch size 4, and evaluation on PersonaMem-32K while varying RL updates up to 200.Source: §E.4 · pp. 18–19, Figure 11 · p. 19
E.6 Optimization Step Analysis
- MGI’s best relative gains over the 10-step setting are +6.1% (PersonaMem 32K), +10.3% (PersonaMem 128K), +3.5% (PrefEval Explicit), and +3.9% (PrefEval Implicit).Performance peaks at an intermediate step budget; too few steps leave updates underdeveloped, while larger budgets can weaken generalization through over-specialization or contradictory edits.Source: Figure 10 · p. 19, §E.6 · p. 20
E.7 Comparison with TextGrad
- On PersonaMem (32K), TextGrad improves over the manual prompt but remains substantially below MGI.The authors attribute the gap to MGI’s trajectory-grounded contrastive feedback and batch aggregation.Source: §E.7 · p. 20, Table 12 · p. 20
F Case Study
- In four cases, MemCoE selects both PersonaMem ground-truth options and returns PersonaBench age 39 and work location University, while baselines miss these answers.Baselines choose different options in the two multiple-choice cases and often return unknown or missing information in the factual QA cases.Source: F Case Study · p. 20, Figure 12 · p. 21, Figure 13 · p. 21
G.1 Meta Prompt for Guideline Optimization
- Guideline optimization diagnoses correct and incorrect updates, aggregates consistent actionable feedback, then revises template_evolve while preserving required placeholders.The three-stage pipeline is illustrated by the contrastive diagnosis, feedback aggregation, and template revision prompts.Source: §G.1 · pp. 20, 22, Figure 16 · p. 23, Figure 17 · p. 23
G.2 Prompt for Memory Evolution and Final Answer Generation
- Memory evolution progressively adds recency/usefulness, evidence bounds, conflict resolution, and selective exclusion of privacy-sensitive or unsupported content.These constraints address unresolved conflicts and over-collection of one-off details, which can make profiles inconsistent or noisy.Source: §G.2 · p. 22, Figure 19 · p. 24
- Shared answer prompts constrain multiple-choice outputs to an option letter and PersonaBench outputs to relevant entity names, standardizing response granularity for fair evaluation.The task-specific prompts are shown for PersonaMem/PrefEval and PersonaBench.Source: §G.2 · pp. 22–23, Figure 20 · p. 25, Figure 21 · p. 25