MemCoE: Teaching LLM Agents How to Remember Before What to Remember, Inspired by Cognitive Psychology

Forum topic · 小凯 · 2026-05-20

Summary

MemCoE is a two-stage memory evolution framework for LLM agents from USTC and City University of Hong Kong, inspired by Memory Schema Theory in cognitive psychology. Mirroring the prefrontal cortex-hippocampus division of labor, Stage 1 (Memory Guideline Induction) learns a global natural-language memory guideline via contrastive feedback treated as textual gradients and batch-level aggregation. Stage 2 (Guideline-Aligned Memory Policy Optimization) performs multi-turn reinforcement learning with GRPO under a hybrid reward combining guideline compliance (dense process reward) and answer correctness. On PersonaMem, PrefEval, and PersonaBench, MemCoE outperforms all baselines (LongContext, RAG, Mem0, A-MEM, LightMem, MemAgent, MEM-alpha) across 8 settings, with ablations confirming both stages are essential. Notably, guidelines induced on Qwen2.5-7B transfer directly to GPT-4o-mini, Gemini 2.5 Flash, and GPT-5 without further tuning, indicating model-agnostic meta-knowledge. Limitations include scorer reliability, error accumulation over long histories, single-objective optimization, and domain dependence.

Overview

MemCoE (https://arxiv.org/abs/2605.00702, code: https://github.com/Applied-Machine-Learning-Lab/ACL2026_MemCoE) is a two-stage memory evolution framework for LLM agents developed by teams from USTC and City University of Hong Kong. Its design borrows from Memory Schema Theory (Alba & Hasher, 1983): the prefrontal cortex decides *how* memories are organized, while the hippocampus encodes *what* is remembered. MemCoE maps this division onto LLM agents.

Key points

Problem: Existing memory systems either use fixed hand-crafted templates (Mem0, A-MEM, LightMem) that never learn, or apply RL directly to memory editing (MemAgent, MEM-alpha) with sparse, delayed outcome-level rewards — a huge action space with weak constraints.
Stage 1 — Memory Guideline Induction (MGI): Treats the memory guideline as a global natural-language parameter. Uses contrastive feedback as textual gradients (comparing correct vs. incorrect trajectories), batch-level aggregation for stable global signals, and natural-language optimization (S^(k+1) = Optim(S^(k), G^(k); P_o)) to induce a high-quality guideline — effectively meta-learning without touching model weights.
Stage 2 — Guideline-Aligned Memory Policy Optimization (GMPO): Freezes the guideline and trains the memory-evolution policy with GRPO (no critic needed), using a combined reward: R(τ) = (1-λ) * R_S(τ; S*) + λ * R_ans(τ), where R_S is a dense process reward scoring guideline compliance of each memory update.

Results

Benchmarks: PersonaMem (32K/128K contexts), PrefEval (explicit/implicit preferences), PersonaBench (noisy corpora, F1). Backbone: Qwen2.5-7B-Instruct; 4x A6000 GPUs; ~300 training samples.
MemCoE is best across all 8 evaluation settings. LongContext degrades sharply under noise; RL baselines trail due to sparse rewards.
Ablations: Removing MGI hurts PrefEval most (-8.1 / -6.3); removing GMPO hurts PersonaMem most (-3.69); removing both collapses results fully — the two stages are complementary (guideline = structure, RL = adaptability).
Efficiency: Faster than A-Mem/Mem0 (memory operations internalized into the model) and more accurate than MemAgent/MEM-alpha — a genuine design-level efficiency frontier.
Cross-LLM transfer: Guidelines optimized on Qwen2.5-7B improve results when directly applied to gpt-4o-mini (52.56 vs 48.47 for A-Mem), gemini-2.5-flash (64.62), and GPT-5 (66.67) — and vice versa — suggesting the induced organization principles are model-agnostic meta-knowledge.

Stated limitations

Reliance on an LLM scorer for the guideline-alignment reward (noisy rewards if the scorer misreads the guideline).
Error accumulation over chunked 4K-token history processing (unintended forgetting, over-generalized entries).
Single-objective optimization; multi-objective control (stability vs. plasticity) is future work.
Guideline induction still depends on training-data domain distribution.

Practical takeaways

Decouple "how to organize" from "what to remember": induce a guideline first, then train a policy within it.
Add dense process-level rewards to memory RL rather than waiting on final-answer correctness.
Cheap approach to deployment: induce guidelines on a small model, execute on stronger ones.

MemCoE represents a paradigm shift from hand-crafted memory rules to learnable memory strategies, structured by cognitive-science-inspired decomposition rather than pure end-to-end learning — more stable, interpretable, and transferable.