💬 NLP

UMEM: Unified Memory Extraction and Management Framework for Generalizable Memory

UMEM is a self-evolving agent framework that improves memory generalizability by jointly optimizing memory extraction and management through Semantic Neighborhood Modeling and neighborhood-level marginal utility rewards via GRPO.

Original authors: Yongshi Ye, Hui Jiang, Feihu Jiang, Tian Lan, Yichao Du, Biao Fu, Xiaodong Shi, Qianghuai Jia, Longyue Wang, Weihua Luo

Published 2026-02-12

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Yongshi Ye, Hui Jiang, Feihu Jiang, Tian Lan, Yichao Du, Biao Fu, Xiaodong Shi, Qianghuai Jia, Longyue Wang, Weihua Luo

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

Imagine you are teaching a robot how to navigate a busy kitchen.

At first, the robot is a bit of a "rote memorizer." If you teach it, "To clean a plate, pick up the blue plate and put it in the sink," it learns that specific instruction perfectly. But the next time you give it a red plate, it freezes. It doesn't realize that "plate" is a general concept; it thinks it only knows how to handle the blue one. This is what the researchers call the "Rote Memorization Trap"—the robot is collecting useless, hyper-specific clutter instead of actual wisdom.

The paper introduces UMEM, a framework designed to turn that robot from a mindless recorder into a wise student.

Here is how UMEM works, explained through three simple metaphors:

1. The "Wisdom vs. Gossip" Filter (Memory Extraction)

Most AI agents act like people who remember every single detail of a conversation, even the boring parts (e.g., "At 2:03 PM, John wore a red tie and said 'hello'"). This is "noise."

UMEM acts like a Master Philosopher. Instead of recording the red tie and the exact time, it distills the experience into a universal principle: "When greeting someone, acknowledge them politely." It ignores the "gossip" (the specific details) and extracts the "wisdom" (the general rule). This ensures that when the robot encounters a new situation, it has a rule it can actually use.

2. The "Neighborhood Test" (Semantic Neighborhood Modeling)

How do you know if a rule is actually useful or just a lucky guess? UMEM uses a clever trick called the Neighborhood Test.

Imagine a student studying for a math test. If they only memorize the answer to one specific problem, they might pass that one question but fail the whole exam. UMEM doesn't let the AI "study" just one problem. Instead, it creates a "neighborhood" of similar problems.

Before the AI is allowed to save a new memory, the system asks: "Does this new rule help you solve all the similar problems in this neighborhood, or just this one?" If the rule only works for one specific instance, the AI is told, "No, that's just a shortcut. Find a rule that works for the whole group!" This forces the AI to develop general intelligence rather than just memorizing a cheat sheet.

3. The "Smart Librarian" (Unified Management)

In older systems, the "Learning" part and the "Organizing" part were two different people who didn't talk to each other. The Learner would bring in a pile of notes, and the Librarian would just shove them into a drawer.

UMEM makes them a single, unified brain. The "Learner" (Extraction) and the "Librarian" (Management) work in perfect sync. The AI doesn't just decide what to learn; it simultaneously decides how to file it away so it can be found easily later. It’s like a chef who, while cooking, is also simultaneously organizing their spice rack so that the next time they make a different dish, the right ingredients are already within reach.

The Result: A Self-Evolving Pro

Because of this, the researchers found that UMEM-powered agents don't just get better; they get steadily smarter the more they work.

While other robots eventually get "confused" by their own messy memories (like a cluttered desk making it harder to work), UMEM agents maintain a "monotonic growth curve." This means their "desk" stays organized, their "wisdom" stays relevant, and they keep getting faster and more efficient at solving complex, real-world tasks.

Technical Summary: UMEM (Unified Memory Extraction and Management Framework)

UMEM is a self-evolving agent framework designed to transform interaction experiences into generalizable, long-term memories. It addresses a critical bottleneck in current self-evolving agents: the inability to distinguish between generalizable principles and instance-specific noise.

1. The Problem: The "Rote Memorization" Trap

Existing self-evolving agents typically treat memory extraction (distilling insights) and memory management (updating the memory bank) as decoupled processes. Most current methods focus on optimizing management (selection/retrieval) while treating extraction as a static process (using fixed prompts). This leads to two major issues:

Accumulation of Instance-Specific Noise: Agents tend to memorize specific details (e.g., "the answer to $x+y=10$ is $x=6$ ") rather than general rules (e.g., "solve systems of equations by substitution"). This "rote memorization" pollutes the memory bank.
Management Misalignment: Because extraction and management are not optimized together, the extracted memories often do not align with the policies used to retrieve or update them, rendering the memory system ineffective.

2. Methodology: Joint Optimization via UMEM

UMEM proposes a unified framework where a learned Memory Optimizer (Mem-Optimizer) is jointly trained to perform both extraction and management. The framework consists of three components: a frozen Agent Executor, a Memory Bank (non-parametric parameters), and the Mem-Optimizer.

Key Technical Innovations:

Semantic Neighborhood Modeling (SNM): To prevent overfitting to a single query, UMEM constructs a "neighborhood" of semantically similar queries using a pre-trained encoder (e.g., BGE-M3). During training, the model is evaluated not just on the current query, but on its ability to improve performance across the entire neighborhood. This forces the optimizer to distill high-level, generalizable principles.
Marginal Utility Reward via GRPO: The framework uses Group Relative Policy Optimization (GRPO) to train the Mem-Optimizer. The reward function is a Marginal Utility Reward, which consists of:
1. Success Gain ( $G_{succ}$ ): Rewards the memory if it corrects failures in the neighborhood.
2. Efficiency Regularization ( $R_{eff}$ ): Rewards concise reasoning (token reduction) only if correctness is maintained, preventing the agent from learning "shortcuts" that are short but wrong.
Online Memory Evolution: Unlike static training, UMEM updates the memory bank dynamically during the training process with the best-performing memory actions. This forces the agent to learn how to operate within a continuously evolving system.

3. Key Contributions

Unified Framework: The first approach to jointly optimize the entire memory lifecycle (extraction + management) in a single end-to-end pipeline.
Generalization Mechanism: Introduced Semantic Neighborhood Modeling to shift the focus from instance-specific memorization to cross-task generalization.
Robust Self-Evolution: Demonstrated that an agent can maintain a monotonic growth curve in performance, avoiding the "error accumulation" common in long-term interaction.

4. Experimental Results

The authors evaluated UMEM across five benchmarks (including AIME, GPQA, and ALFWorld) using various executors (Qwen3, GPT-5.1, Gemini-2.5-Flash).

Superior Performance: UMEM significantly outperformed competitive baselines like ReMem and Memp. In multi-turn interactive tasks (ALFWorld), it achieved up to a 10.67% improvement in success rates.
Cross-Model Portability: The memories distilled by UMEM (even when trained on smaller models) proved highly effective when applied to much stronger, unseen executors like GPT-5.1.
Efficiency Gains: In embodied tasks, UMEM not only increased success rates but also reduced the average number of steps required, proving that it learns efficient procedural knowledge rather than just "trial and error."
Stability: Ablation studies confirmed that joint optimization is essential; optimizing management without extraction leads to severe performance degradation.

5. Significance

UMEM represents a significant step toward lifelong learning for AI agents. By moving away from "rote memorization" and toward "principled distillation," it provides a blueprint for agents that can operate in open-ended, dynamic environments. It demonstrates that for an agent to truly evolve, it must not only learn what happened but also why it happened in a way that applies to future, unseen scenarios.