MemCoE: A Two-Stage, Schema-Inspired Memory System for LLM Agents
TL;DR: A USTC + City University of Hong Kong team borrows cognitive psychology's Memory Schema Theory — the prefrontal cortex organizes memory schemas while the hippocampus encodes specific content — and maps this division of labor onto LLM Agent memory. Stage 1 induces a global "memory guideline" via contrastive feedback and textual gradients; Stage 2 trains a multi-turn RL policy under that guideline's constraint. MemCoE beats all baselines on PersonaMem, PrefEval, and PersonaBench, and the induced guideline transfers across LLMs (optimized on Qwen, used directly on GPT-5).
Key points
- Problem: Existing memory systems either use static hand-crafted templates (Mem0, A-MEM, LightMem) that never learn, or rely on RL with sparse, delayed outcome-level rewards (MemAgent, MEM-α), making training unstable due to a huge action space and weak signal.
- Inspiration: Memory Schema Theory (Alba & Hasher, 1983) — decoupling "how to organize" (stable, abstract, transferable schema) from "what to record" (flexible, contextual content).
- Stage 1 — MGI: Treats the guideline S as an optimizable natural-language parameter. Correct trajectories vs. negative ones are compared by an LLM to produce textual gradients; batch-level aggregation identifies common failure modes; natural-language optimization updates the guideline.
- Stage 2 — GMPO: With S* fixed, the memory-evolution + answering policy is trained via GRPO using a weighted reward:
R(τ) = (1-λ) * R_S(τ; S*) + λ * R_ans(τ), where R_S is a dense per-step guideline-compliance score and R_ans is binary answer correctness. - Results: Best on all 8 evaluation settings. LongContext degrades sharply under noisy histories; RL baselines remain competitive but trail, validating process-level rewards.
- Ablations: Removing either stage hurts — MGI removal is most damaging for PrefEval (-8.1), GMPO removal for long PersonaMem histories; removing both collapses performance (-8.6 to -9.6).
- Efficiency: Faster than Mem0/A-Mem (extraction/update/forgetting internalized into the model) and more accurate than MemAgent/MEM-α.
- Transferability: Guidelines induced on Qwen2.5-7B outperform RAG and A-Mem when directly applied to gpt-4o-mini (52.56 vs 48.47), gemini-2.5-flash (64.62 vs 62.37), and GPT-5 (66.67 vs 64.42) — evidence of model-agnostic meta-knowledge.
- Benchmarks: PersonaMem (32K/128K contexts), PrefEval (explicit/implicit preferences), PersonaBench (noisy corpora, noise 0.3/0.5/0.7)
- Backbone: Qwen2.5-7B-Instruct; retrieval: all-MiniLM-L6-v2, Top-10; 4K-token chunks; trained on 300 PersonaMem samples; 4× A6000 GPUs
- Paper: https://arxiv.org/abs/2605.00702
- Code: https://github.com/Applied-Machine-Learning-Lab/ACL2026_MemCoE
- Baselines: Mem0 (https://github.com/mem0ai/mem0), MemAgent, MEM-α
- Theory: Memory Schema Theory (Alba & Hasher, 1983); GRPO (Shao et al., 2024); TextGrad (Yuksekgonul et al., 2025)
Formalization
Memory evolution: M_{t+1} = T(M_t, h_t; S, φ), answering: y_t = A(x, M_t). Guideline induction objective: S* = argmax_S E_(H,x)[R(τ+, {τ_j^-}; S)] with simple binary correctness reward.
Experiment setup
Limitations (acknowledged by the authors)
1. Dependence on an LLM scorer for the guideline-compliance reward 2. Error accumulation over long chunked histories (unintended forgetting, over-generalization) 3. Single-objective optimization (no stability/plasticity trade-off control) 4. Domain dependence of guideline induction
Why it matters
MemCoE marks a paradigm shift from hand-designed memory update rules to learnable memory strategies: "how to organize" is induced from data, "what to store" is learned by RL within the guideline's constraints — more stable, interpretable, and transferable than fully end-to-end learning, analogous to deep learning replacing hand-crafted features.