UMEM: Unified Memory Extraction and Management Framework for Generalizable Memory
Abstract
Self-evolving memory serves as the trainable parameters for Large Language Models (LLMs)-based agents, where extraction (distilling insights from experience) and management (updating the memory bank) must be tightly coordinated. Existing methods predominantly optimize memory management while treating memory extraction as a static process, resulting in poor generalization, where agents accumulate instance-specific noise rather than robust memories. To address this, we propose Unified Memory Extraction and Management (UMEM), a self-evolving agent framework that jointly optimizes a LLM to simultaneously extract and manage memories. To mitigate overfitting to specific instances, we introduce Semantic Neighborhood Modeling and optimize the model with a neighborhood-level marginal utility reward via GRPO. This approach ensures memory generalizability by evaluating memory utility across clusters of semantically related queries. Extensive experiments across five benchmarks demonstrate that UMEM significantly outperforms highly competitive baselines, achieving up to 11.49 points improvement in multi-turn interactive tasks. Furthermore, UMEM maintains a stable improvement trend during continuous evolution.
Lay Summary
AI agents can use saved memories from past tasks, much like notes, to help with new tasks. But this system has two parts: a task agent that uses the notes, and a memory editor that writes or updates them. If the editor saves details that only fit one example, the task agent may follow the wrong lesson later. We propose UMEM to train the memory editor to write notes that help beyond the single task that produced them. UMEM checks whether a new note also helps on related tasks, not just the original one. This encourages general lessons and filters out noisy shortcuts. With better notes, agents can answer more questions and complete more step-by-step tasks without retraining the whole language model.