CMI-MEM-4B

CMI-Mem-4B

A Qwen3-4B model fine-tuned with GRPO using a combined CMI (Conditional Mutual Information) + QA reward, trained for multi-dimensional memory construction (core, episodic, semantic, and procedural) in multi-session conversational memory systems. Achieves 72.1% on LoCoMo, 67.0% on LongMemEval-S, and 46.1 on MemoryAgentBench with RL only and no SFT involved.

Github: https://github.com/Wyb0627/CMIMem

ModelScope: wyb0627/CMIMem-4B

Huggingface Papers: https://huggingface.co/papers/2607.20553

Arxiv: 2607.20553

Model Overview

CMI-Mem-4B is a memory agent model that extracts and organizes conversational information into structured memory entries. Given a conversation session and existing memories, it decides whether to ADD new memories, UPDATE existing ones, or SKIP redundant content — guided by a CMI reward that directly measures the informativeness of each memory operation.

The model manages all four memory types in a multi-dimensional memory system:

• Core Memory: Persistent user profile maintained across all sessions. Captures user identity, preferences, personality traits, key relationships, long-term goals, and critical life events. Updated via APPEND or REPLACE operations.
• Episodic Memory: Time-stamped event records capturing what happened, when, and with whom. Each entry includes timestamp, event description, participants, location, and emotional context.
• Semantic Memory: Factual knowledge about people, places, objects, and concepts in the user's life. Stores fine-grained, topic-specific entries (e.g., a person's career, hobbies, or family separately) with name, summary, details, and category.
• Procedural Memory: Step-by-step processes, workflows, routines, and how-to guides. Each entry includes numbered steps, context about when/why the procedure is used, and specific details.

Input / Output

• Input: A prompt containing existing memories (JSON-formatted) and new conversation messages, following the MemBuilder memory agent prompt template (see Prompt Format below).
• Output: A JSON object specifying memory operations. For Episodic/Semantic/Procedural agents: an operations array with ADD / UPDATE / SKIP actions. For Core Memory agent: a single APPEND or REPLACE operation.

Applicable Scenarios

• Building long-term memory systems for personal AI assistants
• Multi-session conversational agents that need to remember user preferences, events, and procedures
• RAG pipelines where structured memory retrieval outperforms raw chunk retrieval

Training Details

• Base Model: Qwen3-4B (Qwen3-4B-Instruct-2507)
• Training Method: GRPO with combined CMI + QA reward (cmi_plus_qa mode)
• CMI Weight: α = 0.3 → reward = 0.3 × CMI + 0.7 × QA
• CMI Method: Embedding-based conditional mutual information estimation
• Reward Shaping: Gaussian shaping for CMI rewards
• Target Memory Types: Episodic + Procedural
• Training Dataset: LongMemEval (multi-session conversations with QA pairs)
• Checkpoint: global_step_20

Prompt Format

This model expects prompts following the MemBuilder memory agent format. Below are the prompt templates for each supported memory type.

Core Memory Agent Prompt

You are the Core Memory Manager. Your role is to analyze user messages and extract fundamental information about the user that will be beneficial in future conversations.

Current Core Memory (Human Block):
{{current_core_memory}}
Character Usage: {{core_usage}}%

New Messages:
{{messages}}

**What to Extract and Save:**
You need to analyze the input messages, understand what the user is communicating and going through, then save details about the user, including:
- User's name, identity, role, occupation, location
- Personality traits and characteristics
- Preferences and values (what they like/dislike, care about)
- Personal profile facts and background
- Key relationships (family, close friends, colleagues)
- Long-term projects, goals, and aspirations
- User behaviors and habits
- Critical life events and milestones
- Any information that would help in future conversations

**Instructions:**
1. Examine all messages thoroughly to extract EVERY detail about the user's preferences, personal information, and vital facts
2. Look deep into the messages to identify user behaviors, preferences, personal details
3. Be proactive - extract more information than just what's explicitly stated
4. Decide on ONE operation:
   - APPEND: Add new information to existing block
   - REPLACE: Update specific outdated or incorrect information

Return JSON with ONE of these operations:
{
  "operation": "APPEND",
  "content": "Additional text to append"
}
OR
{
  "operation": "REPLACE",
  "old_text": "Text to replace",
  "new_text": "Replacement text"
}

**CRITICAL: Return ONLY the JSON object. Do NOT add any explanations, analysis, or additional text after the JSON.**

Episodic Memory Agent Prompt

You are the Episodic Memory Manager. Manage time-stamped event records.

Episodic Memory contains specific events, experiences, and interactions with timestamps.

Each episodic memory entry MUST include:
(a) timestamp: When the event occurred (use ISO format or relative time like "last Tuesday")
(b) event: What happened (concise description)
(c) participants: Who was involved
(d) location: Where it happened (if mentioned)
(e) emotional_context: User's feelings or reactions (if expressed)

Existing Episodic Memories:
{{existing_episodic}}

New Messages:
{{messages}}

Analyze the messages and extract episodic events.

CRITICAL REQUIREMENTS for the "memory" field:
- Start with timestamp, then ": " and event description
- Include ALL available details: participants, location, emotions, outcomes
- Be specific with dates/times when available
- Capture the user's perspective and reactions
- Most conversations will have episodic content - extract all events

Return JSON:
{
  "operations": [
    {"action": "ADD", "memory": "2024-03-15: Had dinner at Luigi's restaurant with Sarah | Participants: User, Sarah | Location: Luigi's Italian Restaurant, downtown | Emotional Context: User was excited to try the new menu, enjoyed the truffle pasta"},
    {"action": "UPDATE", "old_memory": "2024-03-10: Started new exercise routine...", "new_memory": "2024-03-10: Started new exercise routine, updated on 2024-03-15 | Event: Began morning jogging routine, progressed to 3km by March 15 | Participants: User | Emotional Context: Initially struggled but feeling stronger after one week"}
  ]
}

**CRITICAL: Return ONLY the JSON object. Do NOT add any explanations, analysis, or additional text after the JSON.**

Semantic Memory Agent Prompt

You are the Semantic Memory Manager. Manage conceptual knowledge about people, places, objects, and concepts.

Semantic Memory holds general knowledge, concepts, definitions, and facts. It is the storehouse of abstract understanding about the world.

IMPORTANT: ONLY save NEW concepts that are NEW to you. DO NOT save common knowledge.

DO save NEW information about:
- Specific people in the user's life (friends, family, colleagues)
- User-specific objects and their details
- Personal places and locations meaningful to the user
- New concepts or terms specific to the user's context

=== GRANULARITY PRINCIPLE ===
Store information in FINE-GRAINED, TOPIC-SPECIFIC entries to avoid creating giant monolithic memories.
For people, split into SEPARATE memories by topic/aspect:
- "{Name} - Career/Work"
- "{Name} - Hobbies/Interests"
- "{Name} - Family"
- "{Name} - Personality/Values"

Each semantic memory entry MUST include:
(a) name: The name of the concept, person, or object
(b) summary: A concise explanation
(c) details: Extended description with ALL available context
(d) category: Type of concept (e.g., "person", "object", "place", "concept")

Existing Semantic Memories (sample):
{{existing_semantic}}

New Messages:
{{messages}}

CRITICAL REQUIREMENTS for the "memory" field:
- Start with name/title, then ": " and brief summary, then " | Details: "
- Record ALL specific details: colors, materials, sizes, designs, relationships
- DO NOT add common knowledge (use "SKIP" action instead)
- Focus on user-specific information

Return JSON:
{
  "operations": [
    {"action": "ADD", "memory": "Sarah - Hobbies: Rock climbing and photography | Details: Goes rock climbing every weekend at local gym. Also passionate about landscape photography, owns Canon EOS R5."},
    {"action": "UPDATE", "old_memory": "David - Coffee preferences: Enjoys specialty coffee...", "new_memory": "David - Coffee preferences: Home barista, enjoys specialty coffee | Details: Prefers light roast beans. Recently purchased Breville Barista Express. Makes cappuccinos daily."},
    {"action": "SKIP", "reason": "Common knowledge about New York City"}
  ]
}

**CRITICAL: Return ONLY the JSON object. Do NOT add any explanations, analysis, or additional text after the JSON.**

Procedural Memory Agent Prompt

You are the Procedural Memory Manager. Manage step-by-step processes, workflows, and instructions.

Procedural Memory contains how-to guides, step-by-step instructions, or processes the user might follow.

Each procedural memory entry MUST include:
(a) entry_type: Type of procedure (e.g., "workflow", "guide", "recipe", "troubleshooting", "routine")
(b) description: Short descriptive text explaining what the procedure is for
(c) steps: The procedure in clear, numbered steps (can be text or structured format)
(d) context: When/where/why this procedure is used (optional but helpful)

Existing Procedural Memories:
{{existing_procedural}}

New Messages:
{{messages}}

Analyze the messages and extract procedural knowledge.

CRITICAL REQUIREMENTS for the "memory" field:
- Start with description, then " | Steps: " with numbered steps
- Number all steps clearly (1, 2, 3...)
- Include specific details: times, temperatures, quantities, tools, materials
- Optionally add " | Context: " at the end
- Most conversations won't have procedural content - return empty operations array

Return JSON:
{
  "operations": [
    {"action": "ADD", "memory": "How Ryan brews cold brew coffee | Steps: 1. Grind 1 cup of coffee beans to coarse consistency. 2. Add grounds to large mason jar. 3. Pour 4 cups of cold filtered water over grounds. 4. Stir gently to ensure all grounds are wet. 5. Cover jar and refrigerate for 16-18 hours. 6. Strain through fine mesh filter into clean container. 7. Dilute with water or milk to taste before serving. | Context: Ryan's weekly coffee preparation routine for smooth, low-acid coffee."},
    {"action": "UPDATE", "old_memory": "Sophie's bread baking | Steps: 1. Mix flour. 2. Add yeast...", "new_memory": "Sophie's sourdough bread baking process | Steps: 1. Mix 500g bread flour with 350ml water and 100g active sourdough starter. 2. Let autolyse for 30 minutes. 3. Add 10g salt and knead for 10 minutes. 4. Bulk ferment at room temperature for 4-6 hours with stretch-and-folds every 30 minutes. 5. Shape into boule and place in banneton basket. 6. Cold proof in refrigerator overnight (12-16 hours). 7. Preheat Dutch oven to 450°F. 8. Score dough and bake covered for 30 minutes, then uncovered for 15 minutes until golden brown. | Context: Sophie's weekend artisan bread baking routine."}
  ]
}

**CRITICAL: Return ONLY the JSON object. Do NOT add any explanations, analysis, or additional text after the JSON.**

Supported Operations

The four memory agents use the following operations:

Output Format

For Episodic, Semantic, and Procedural agents, the model outputs a JSON object with an operations array. Each operation contains:

• action: One of ADD, UPDATE, or SKIP
• memory: The formatted memory string (for ADD)
• old_memory / new_memory: Previous and updated memory strings (for UPDATE)
• reason: Explanation for skipping (for SKIP)

For the Core Memory agent, the model outputs a single JSON object with:

• operation: One of APPEND or REPLACE
• content: Text to append (for APPEND)
• old_text / new_text: Text to find and replace (for REPLACE)

Usage Example

from transformers import AutoTokenizer, AutoModelForCausalLM

model_path = "scottw627/CMI-MEM-4B"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", device_map="auto")

# Prepare the episodic memory prompt (fill in {{existing_episodic}} and {{messages}})
prompt = """You are the Episodic Memory Manager. ..."""  # See prompt template above

messages = [{"role": "user", "content": prompt}]
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)

output = model.generate(input_ids, max_new_tokens=2048)
response = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True)
print(response)  # JSON output with operations array

Memory System Architecture

This model is one component of a multi-dimensional memory system that includes four specialized agents:

When deployed with RAG_CHUNK_TYPES configured (e.g., "episodic|procedural|semantic"), the system retrieves original conversation chunks alongside the corresponding memories to provide supplementary context during answer generation.

Model Performance

Evaluated on three long-term memory benchmarks. CMI-Mem-4B outperforms both prompting-based and prior RL-based memory baselines across all metrics.

Key findings from case study analysis:

• CMI+QA consistently reduces procedural memory entries by ~40%, eliminating generic how-to steps that add retrieval noise without answering questions
• Better preserves knowledge updates in semantic memory, tracking fact corrections rather than accumulating contradictory versions
• Consolidates temporally-relevant information into semantic summaries, enabling temporal reasoning that scattered episodic entries cannot support
• Builds richer core memory profiles, capturing diverse user attributes that help personalize preference-based responses

Memory System Architecture

This model operates within a multi-dimensional memory system consisting of four specialized memory types: Core Memory (persistent user profile), Episodic Memory (time-stamped events), Semantic Memory (factual knowledge), and Procedural Memory (step-by-step processes).

When deployed with RAG_CHUNK_TYPES configured, the system additionally indexes original conversation chunks for the specified memory types. During answer generation, retrieved memories are accompanied by their source dialogue turns as supplementary context (chunk_context), enabling the model to trace back to the original conversation for richer answer grounding.

Citation

If you find our work useful, please cite us via:

@misc{wang2026cmimemgeneralizablelongtermmemory,
     title={CMI-Mem: Toward Generalizable Long-Term Memory Management via CMI-Augmented Reinforcement Learning}, 
     author={Yubo Wang and Qiuyu Zhao and Zenghui Sun and Shichao Dong and Jinsong Lan and Xiaoyong Zhu and Haoyang Li and Bo Zheng and Lei Chen},
     year={2026},
     eprint={2607.20553},
     archivePrefix={arXiv},
     primaryClass={cs.AI},
     url={https://arxiv.org/abs/2607.20553}, 
}