💬 NLP

CMI-Mem: Toward Generalizable Long-Term Memory Management via CMI-Augmented Reinforcement Learning

CMI-Mem is a lightweight reinforcement learning memory manager that enhances generalizability and training efficiency by combining extrinsic task-performance rewards with an intrinsic Conditional Mutual Information signal to evaluate memory relevance without relying on fixed query distributions.

Original authors: Yubo Wang, Qiuyu Zhao, Zenghui Sun, Shichao Dong, Jinsong Lan, Xiaoyong Zhu, Haoyang Li, Bo Zheng, Lei Chen

Published 2026-08-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yubo Wang, Qiuyu Zhao, Zenghui Sun, Shichao Dong, Jinsong Lan, Xiaoyong Zhu, Haoyang Li, Bo Zheng, Lei Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, large language models are like brilliant scholars who can reason through complex problems and write with fluency, yet they suffer from a peculiar form of amnesia. Once a conversation ends, these models forget everything that was just said, operating as if every interaction is the first time they have ever spoken to a human. To fix this, researchers have built "memory managers," specialized systems designed to decide what information is worth keeping and what should be discarded. The goal is to create an agent that remembers a user's favorite coffee order from months ago or recalls a specific detail mentioned in a previous session to solve a new problem. However, teaching these managers to be good at their job has been difficult. For years, the standard method has been to train them by asking a series of questions after a conversation and seeing if the model can answer correctly based on what it stored. This approach works, but it has a blind spot: it only values memories that happen to answer the specific questions asked during training, potentially missing the subtle, useful details that don't fit a neat quiz format.

A team of researchers from Alibaba Group and the Hong Kong University of Science and Technology has proposed a new way to train these memory managers, one that looks at the information itself rather than just the final test score. They call their system CMI-Mem. Instead of relying solely on whether a model gets the right answer to a pre-written question, this new method teaches the manager to evaluate every single piece of information it considers storing based on how much new, useful value it adds to the existing collection. Imagine a librarian who, instead of waiting for a patron to ask for a specific book, constantly assesses whether a new volume offers a unique perspective that isn't already covered by the books on the shelf. If the new book repeats what is already there, it is rejected; if it fills a gap or connects ideas in a fresh way, it is kept. This dual approach allows the system to build a memory that is not just good at passing a specific test, but is robust and adaptable to many different kinds of future conversations.

The researchers found that the old method, which relies heavily on "question-and-answer" training, creates a narrow focus. When a memory manager is trained only to answer specific questions, it learns to store facts that directly answer those questions but often ignores the broader context or the relationships between ideas. This makes the system brittle; if the user asks a question in a slightly different way or needs the information for a task the system wasn't explicitly trained on, the memory often fails. The new system, CMI-Mem, solves this by adding a second, internal reward signal. This signal measures the "information gain" of a new memory fragment. It asks a simple, fundamental question: given everything the system already knows, how much does this new piece of conversation teach us that we didn't already know? If the answer is "nothing new," the system learns to skip it. If the answer is "a lot," the system learns to store it. This happens in real-time, for every single action the manager takes, providing a constant stream of feedback that is independent of any specific question.

To test this idea, the team trained their memory manager using a combination of these two signals: the traditional ability to answer questions correctly and the new, internal measure of information value. They tested the resulting system on three different benchmarks designed to challenge long-term memory, including tasks that required recalling facts from very long conversations, reasoning across multiple separate sessions, and selectively forgetting irrelevant details while keeping important ones. The results were clear. The new system outperformed previous methods that relied only on question-based training. It was better at remembering details across long periods and more successful at handling tasks that required understanding the flow of a conversation rather than just retrieving a specific fact. Perhaps most importantly, the system trained more efficiently. Because it received constant feedback on the quality of each memory operation, it learned faster and required fewer attempts to reach a high level of performance.

The study also revealed that the two types of rewards work best when used together. The internal information signal helps the system build a rich, diverse memory that isn't limited to a specific set of questions, while the question-and-answer signal ensures that the memory remains useful for actual tasks. When the researchers tried to use only the internal signal, the system struggled to know what was actually important for the user's goals. When they used only the question signal, the system became too rigid and failed to generalize to new situations. By combining them, the researchers created a memory manager that is both flexible and reliable. This approach suggests a path forward for building AI agents that can truly remember and learn from their interactions, moving beyond simple fact retrieval to a more human-like ability to curate and value their own experiences. The work demonstrates that by grounding memory management in the fundamental principles of information, rather than just the outcome of a test, we can create systems that are better prepared for the unpredictable nature of real-world conversation.

Technical Summary: CMI-Mem

Problem Statement

Large Language Models (LLMs) inherently lack persistent memory of past interactions, limiting their ability to maintain long-term dependencies and achieve robust personalization. While Memory Managers have emerged as critical infrastructure to decide what to store, update, or retrieve, existing approaches face significant limitations:

Query-Conditioned Supervision: Most current methods rely on Reinforcement Learning (RL) trained with synthetic Question-Answer (QA) pairs evaluated by an "LLM-as-a-Judge." This paradigm values memory only through the lens of a fixed downstream reader and a sampled query distribution. Consequently, information outside the specific training questions receives little to no learning signal, even if it is valuable for future summarization, reasoning, or personalization.
Distribution Mismatch: Synthetic QA pairs often emphasize explicit fact-seeking retrieval, failing to capture the broader, associative, or trigger-based nature of real-world memory use.
Lack of Intrinsic Criteria: Existing methods rely on heuristic rules or black-box judgments that lack a unified theoretical foundation for evaluating the marginal information value of a memory operation. This can lead to memory bloating or catastrophic forgetting, as managers struggle to distinguish novel, relevant content from redundancy.

Methodology: CMI-Mem

The authors propose CMI-Mem, a lightweight RL framework that trains a memory manager using a hybrid reward system. This system combines an extrinsic task-grounded signal with an intrinsic information-theoretic signal.

1. Structured Multi-dimensional Memory Architecture

The memory state $M_t$ is partitioned into four cognitively motivated slots, each managed by a type-specific agent sharing a single policy $\pi_\theta$ :

Core: Stable user attributes.
Episodic: Timestamped events.
Semantic: Factual knowledge.
Procedural: Routines and workflows.
Agents emit actions (e.g., ADD, DELETE, MERGE, REPLACE) and candidate memory fragments. Crucially, these fragments serve as retrieval indices to their source dialogue turns, allowing downstream readers to ground answers in raw context rather than just compressed summaries.

2. Hybrid Reward Mechanism

The training objective blends two complementary signals:

Extrinsic QA Reward ( $r_{QA}$ ): Measures end-task correctness. A downstream LLM (Qwen3.7-Max) answers sampled questions based on the updated memory state. This anchors the memory construction to downstream utility.
Intrinsic CMI Reward ( $r_{CMI}$ ): Based on Conditional Mutual Information, this evaluates the information contributed by a new conversational input relative to the current memory state without conditioning on a sampled QA query.
- Definition: $r_{CMI} = I(C_t; m^{new}_t | M^{local}_{t-1})$ , where $C_t$ is the current dialogue context, $m^{new}_t$ is the candidate fragment, and $M^{local}_{t-1}$ is a local conditioning set of existing memories.
- Estimation: Since direct CMI calculation is intractable for natural language, the authors approximate it using residual projection in embedding space. They compute the partial correlation between the context and the new fragment after linearly removing the effect of existing memories.
- Function: High CMI indicates content that is relevant to the new dialogue but novel relative to stored memory; low CMI signals redundancy. This provides dense, per-operation feedback.

3. Training Framework

Optimization: The policy is optimized using Group Relative Policy Optimization (GRPO) at the session level to ensure training stability.
Curriculum Learning: To stabilize early training without supervised fine-tuning (SFT) warm-starts, the authors employ a stratified curriculum that progressively increases sample difficulty based on dialogue depth, fragment count, and QA question count.
Reward Composition: The final reward is a weighted sum: $r_t = \alpha \tilde{r}_{CMI} + (1-\alpha) r_{QA}$ .

Key Contributions

Intrinsic-Extrinsic Complementarity: The paper introduces a hybrid reward that combines downstream task grounding (QA) with an intrinsic signal (CMI). This reduces reliance on specific query distributions while ensuring memory remains useful for target tasks.
Information-Theoretic Modeling: By utilizing Conditional Mutual Information, the method replaces ad-hoc valuation rules with a principled metric for relevance and novelty relative to the evolving memory state.
Per-Operation Optimization: Unlike QA rewards which are often sparse (session-level), the CMI signal is available at every memory operation. This provides denser feedback, improving training stability and rollout efficiency.
Retrieval-Indexed Architecture: The system treats memory fragments as indices to source dialogue turns, enabling provenance-aware retrieval that preserves critical label information often lost in pure summarization.

Experimental Results

The authors evaluated CMI-Mem on three benchmarks: LongMemEval, LoCoMo, and MemoryAgentBench.

Performance: CMI-Mem-4B outperformed RL-based baselines (Memory-R1, MemBuilder-RL) across all three benchmarks. Notably, on the out-of-distribution MemoryAgentBench, CMI-Mem-4B surpassed the MemBuilder-RL baseline by +11.1 in overall score, demonstrating improved transferability.
Scaling: The 8B parameter version (CMI-Mem-8B) outperformed the 8B baseline on MemoryAgentBench and the strongest prompting baselines.
Training Dynamics: Ablation studies showed that training with CMI+QA resulted in more stable policy gradient losses compared to QA-only training. The CMI signal smoothed reward transitions and prevented the progressive negative drift observed in QA-only variants.
Ablation Insights:
- CMI alone lacked an outcome anchor, resulting in lower accuracy.
- QA alone on a flat memory store was effective but suboptimal for complex action spaces.
- Stacking QA on retrieval-indexed memory without CMI caused performance regression, highlighting the necessity of the intrinsic signal for the richer action space.
- The full combination (CMI + QA + Retrieval-Index) achieved the highest accuracy.

Significance and Claims

The paper claims that CMI-Mem addresses the fundamental limitation of "question-driven" memory valuation. By decoupling memory quality assessment from specific downstream queries, the method enables generalizable memory management that transfers effectively to unseen tasks and query types.

The authors emphasize that the CMI signal does not replace task supervision but acts as a necessary complement, providing a direct criterion for the marginal information value of memory operations. This approach allows for more efficient training and inference while supporting cross-scenario transfer, moving beyond the constraints of heuristic rules or black-box LLM judgments.

Limitations Acknowledged:
The authors note three main limitations:

Multimodal Constraints: The system's performance is bounded by the perceptual capabilities of the underlying foundation model.
Model Scale: RL training was validated primarily on 4B and 8B models; scalability to 14B+ models remains unexplored.
Evaluation Metrics: Current benchmarks still rely predominantly on QA testing. The full potential of the method is difficult to demonstrate without direct metrics for memory quality independent of QA.