Learning How and What to Memorize: Cognition-Inspired Two-Stage Optimization for Evolving Memory
| Authors | Derong Xu et al. |
| Year | 2026 |
| Field | NLP |
| arXiv | 2605.00702 |
| Download | |
| Categories | cs.CL |
Abstract
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.
Engineering Breakdown
Plain English
This paper addresses a critical limitation in LLM agents: they struggle to maintain and update user memory across long conversations because context windows are finite and existing memory systems use rigid, hand-crafted rules. The authors propose MemCoE, a two-stage framework inspired by cognitive neuroscience that learns both how to organize memory and what information to update, replacing static rules with learned optimization. The key innovation is using contrastive learning in the first stage (Memory Guideline Induction) to establish better memory organization principles, then optimizing what to update in the second stage, which improves stability compared to sparse reward RL approaches that previous work relied on.
Core Technical Contribution
The core contribution is a cognition-inspired architecture that separates memory management into two learnable components: (1) a Memory Guideline Induction stage that uses contrastive learning to discover optimal memory organization principles globally, and (2) a dynamic update mechanism that learns what information to store or modify. This differs fundamentally from prior memory systems that either use static hand-crafted rules (brittle and inflexible) or pure RL with sparse rewards (unstable and slow to converge). By drawing from memory schema theory and neuroscience insights about how the prefrontal cortex and hippocampus divide memory responsibilities, the authors created a more structured and trainable approach than monolithic end-to-end RL.
How It Works
MemCoE operates in two distinct stages. In stage one, Memory Guideline Induction uses contrastive learning to optimize a global guideline—essentially a learned schema that defines how memory should be organized conceptually (e.g., user preferences by category, temporal patterns, interaction patterns). This stage works by comparing positive memory organizations (those that led to good interactions) against negative ones using a contrastive loss, creating a shared understanding of optimal memory structure. In stage two, given this global guideline, the system learns a dynamic updater that decides which information to add, modify, or discard from the agent's memory at each interaction step. The updater receives the user input, current memory state, and the learned guideline, then produces targeted memory operations rather than full memory rewrites. This two-stage decomposition reduces the search space and provides much denser training signal (from the contrastive stage) compared to waiting for sparse long-horizon rewards.
Production Impact
In production, this approach would transform how conversational AI systems handle user personalization at scale. Instead of losing important user preferences when context limits force older messages out, agents could maintain a compressed, learned memory of user patterns that grows richer over time. The two-stage design reduces training instability—a major issue when deploying RL-based systems in production where reward signals are often sparse and delayed. However, there are real trade-offs: you need to implement both the contrastive training pipeline and the memory update mechanism, requiring additional model capacity and inference computation for memory operations at each turn. Integration complexity increases because you need a separate memory database and a system to encode/decode memory into the model's working context, which adds latency (likely 50-200ms per interaction for memory operations).
Limitations and When Not to Use This
The paper does not address how to scale this to truly long-term memory (months or years of interactions) where the learned guideline itself may need to evolve. The contrastive learning stage requires paired examples of good vs. bad memory organizations, which may be expensive to generate at scale or require manual annotation in some domains. The approach assumes that a single global guideline generalizes well across users, which likely breaks down when user preferences are highly idiosyncratic or in niche domains. The paper also appears incomplete in the abstract (cuts off mid-sentence at 'contrastive feed'), so we cannot assess whether it addresses computational efficiency for real-time inference or provides comparisons against strong RL baselines like PPO with reward shaping.
Research Context
This work builds on a growing body of research combining cognitive science principles with LLM architectures (similar to how attention mechanisms drew inspiration from human attention). It addresses a known limitation of memory-augmented LLMs and RL-based agents: instability from sparse rewards in long-horizon tasks. The two-stage approach echoes recent trends in separating learning into more interpretable components (similar to modular RL and curriculum learning). This likely contributes to the broader research direction of making LLM agents more efficient and stable through better inductive biases rather than pure scale.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::