PaperlayerView exact revision on arXiv

Source-linked AI summary

UMEM: Unified Memory Extraction and Management Framework for Generalizable Memory

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

arXiv:2602.10652v1cs.CL

TL;DR

Existing approaches often decouple memory extraction from management, limiting generalization. UMEM jointly optimizes both using Semantic Neighborhood Modeling and GRPO, outperforming competitive baselines in cross-task generalization and execution efficiency.

  • Problem

    Prior methods predominantly optimize memory management while treating extraction statically, without explicitly optimizing generalization or aligning extracted memories with management policies.

  • Method

    UMEM jointly optimizes memory extraction and management using Semantic Neighborhood Modeling and GRPO with a Marginal Utility Reward.

  • Results

    UMEM consistently outperforms competitive baselines across most benchmarks, with gains in cross-task generalization and execution efficiency.

  • Takeaways & Limitations

    UMEM provides a robust paradigm for continuously refining memory banks during interaction toward lifelong learning in open-ended environments.

  • Takeaways & Limitations

    System performance is strictly bounded by the quality of the retrieved memory.

Abstract

from arXiv · show

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 predominately 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 Large Language Model to simultaneous 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 a 10.67% improvement in multi-turn interactive tasks. Futhermore, UMEM maintains a monotonic growth curve during continuous evolution. Codes and models will be publicly released.

1. Introduction

UMEM addresses the limited self-evolution of frozen-parameter agents by jointly optimizing memory extraction and management for generalizable memories. It uses semantic-neighborhood-based reward optimization and demonstrates improved, stable performance across benchmarks and continued interactions.

  • Motivation: Long-term memory acts as trainable parameters that let agents learn from continuous online interactions despite frozen LLM parameters.The framework conceptualizes self-evolution as executing tasks with retrieved memories and consolidating extracted insights into an external memory bank.
  • Problem: Existing memory optimizers predominantly manage memories while treating extraction as static prompting, causing accumulation of instance-specific noise without explicit generalization.The introduction identifies static extraction as retaining instance-specific details rather than robust, generalizable principles.
  • UMEM Framework: UMEM jointly optimizes memory extraction and management through a frozen Agent Executor, Memory Bank, and learned Mem-Optimizer.The Mem-Optimizer is the core component responsible for transforming executor experience into memories while coordinating extraction and management.
  • Optimization: Semantic Neighborhood Modeling clusters semantically related queries, while a Marginal Utility Reward optimized with GRPO guides end-to-end Mem-Optimizer training.The design simulates cross-task variation to reduce instance-specific noise and target generalizable extracted memories.
  • Results: UMEM significantly outperforms competitive baselines including ReMem and Memp across five benchmarks, while maintaining consistent performance gains as interactions continue.Ablations show that optimizing memory management alone causes significant performance degradation, supporting joint optimization; test-time scaling evolution demonstrates stable self-evolution and widening gaps over baselines.

2. Related Work

Related work has shifted from computationally costly parametric memory mechanisms toward non-parametric external memory banks, while memory evolution has progressed from rule- or prompt-based extraction to optimization-based approaches. Existing methods nevertheless leave memory extraction and management capabilities limited or focus optimization on working memory.

  • From Parametric Memory to Non-Parametric Memory: Memory-augmented language models evolved from architectural mechanisms and scalable lookup frameworks toward external memory banks treated as evolvable agent parameters.Earlier approaches require computationally costly fine-tuning.
  • Self-Evolving Memory without Optimization: Non-parametric memory methods improve experience representation by replacing raw trajectory retrieval with executable programs, reusable entries, and semantic compression.Synapse retrieves raw historical trajectories, causing severe noise and context-window inefficiency; later methods introduce structured abstraction.
  • Self-Evolving Memory without Optimization: Prompting-based and hand-crafted extraction and management policies prevent further improvement of these capabilities.This limitation applies to methods including Memp, ReasoningBank, and SimpleMem.
  • Self-Evolving Memory with Optimization: Optimization-based self-evolving agents use reinforcement learning in two streams, including methods that optimize working or short-term memory.DeepAgent, MemAgent, and Mem-α apply RL to manage working or short-term memory and handle long-context inputs.

3. Task Formulation of Self-Evolving Agents

UMEM formulates self-evolving agents as systems with frozen executors and evolvable external memory banks, organized into inference and memory-evolution passes. Its central focus is jointly optimizing memory extraction and management policies because the memory optimizer is the core bottleneck.

  • Agent Formulation: A self-evolving agent has frozen executor parameters Θ0 and an evolvable, nondifferentiable memory bank B of key–value pairs.Keys correspond to queries, while values store associated memory content.
  • Optimization View: UMEM conceptualizes self-evolution as network optimization with a forward inference pass and a backward memory-evolution pass.This decomposition coordinates memory-augmented execution with subsequent memory updates.
  • Feedforward Pass: At time t, the agent retrieves Top-K relevant entries and uses the frozen executor to generate a trajectory and prediction conditioned on the retrieved context.The retrieved memory is denoted Btopk_t, and the resulting trajectory is τq with prediction ŷt.
  • Feedforward Pass: Because Θ0 is fixed, system performance is strictly bounded by the quality of the retrieved memory and the memory bank.This makes memory quality the determining factor in the frozen-executor setting.
  • Backward Pass: The Mem-Optimizer extracts memory entries from trajectories, samples operations such as ADD or UPDATE, and evolves the bank via Bt+1 ← Apply(Bt, aq).The extracted memory is Δq, the action is aq, and the formulation can extend to successful or failed trajectory pairs.
  • Optimization Objective: UMEM jointly optimizes the Mem-Optimizer’s extraction and management policies, treating it as the core bottleneck in self-evolving agents.The Mem-Optimizer is parameterized by ϕ and denoted πϕ.

4. Method

UMEM jointly extracts and manages memories by modeling semantic neighborhoods and optimizing neighborhood-level marginal utility with GRPO. Its training loop evaluates candidate updates across related queries, enforces valid memory operations, and applies the best update online.

  • Semantic Neighborhood Modeling: Semantic Neighborhood Modeling clusters semantically related queries to approximate cross-task variation and reduce overfitting to instance-specific noise.Queries are projected into a shared semantic space and organized into local neighborhoods.
  • Mem-Optimizer Training: The Mem-Optimizer distills executor trajectories into structured memory actions and samples candidate updates using GRPO.A frozen executor first produces a trajectory and prediction from retrieved Top-K memory entries; the optimizer then proposes memory actions.
  • Marginal Utility Reward: Marginal Utility Reward evaluates each candidate update across the semantic neighborhood rather than only its source query.It combines Success Gain, which rewards correcting failures and penalizes introduced errors, with Efficiency Regularization, gated so brevity is not rewarded at the cost of accuracy.
  • Optimization via GRPO: GRPO trains the memory policy with a joint objective combining marginal utility and a binary XML-format reward.The format reward validates adherence of extracted memories and management operations to the specified XML schema.
  • Online Memory Evolution: After optimization, the highest-reward memory action is immediately applied to the memory bank, enabling online refinement during training.This exposes the agent to an evolving memory rather than a static memory bank.

5. Experiments

Experiments across five benchmarks show that UMEM consistently outperforms representative memory baselines, with gains amplified by stronger executors and larger policy models. Ablations and continual-evolution evaluations support the importance of semantic neighborhoods, joint optimization, and stable memory updates.

  • Benchmark: UMEM is evaluated across five benchmarks spanning single-turn reasoning and multi-turn embodied interaction, including AIME, GPQA-Diamond, HLE, HotpotQA, and ALFWorld.The study tests mathematical, scientific, multidisciplinary, and embodied interaction capabilities, while assessing memory stability and reusability.
  • Main Results: 82.84% Success Rate is achieved by UMEM-Qwen3-4B with GPT-5.1, while UMEM consistently outperforms baselines across most benchmarks.The gains are especially pronounced on complex reasoning tasks such as AIME and GPQA Diamond and in ALFWorld.
  • Main Results: Stronger frozen executors, including GPT-5.1 and Gemini-2.5-Flash, produce more pronounced gains than the Qwen3-8B-Thinking baseline.Higher-quality reasoning trajectories and interaction traces provide better source material for distilling actionable insights.
  • Main Results: Scaling the policy model from 1B to 4B consistently yields additional performance dividends across nearly all tasks.Even UMEM-Llama-3.2-1B substantially improves over the base model and often surpasses larger models.
  • Ablation Study: Removing Semantic Neighborhood Modeling drops AIME performance for GPT-5.1 from 51.67 to 41.67, while N = 3 provides the optimal balance.Both N = 1 and N = 5 perform worse because narrow neighborhoods miss task shifts and broad neighborhoods introduce reward noise.
  • Ablation Study: 4.7 points is the average performance decline across all metrics when memory extraction optimization is disabled.Breaking joint optimization degrades performance across most benchmarks, with extraction optimization contributing more than management optimization in the reported comparison.
  • Continual Evolution: UMEM maintains the highest Success Rate across all epochs in a 10-epoch ALFWorld evaluation with GPT-5.1 as executor.Its evolved memories support higher success with fewer execution steps, indicating improved efficiency rather than longer interaction trajectories.

6. Conclusion

UMEM jointly optimizes memory extraction and management for self-evolving agents using Semantic Neighborhood Modeling and GRPO with a Marginal Utility Reward. It mitigates instance-specific noise, outperforms competitive baselines in cross-task generalization and execution efficiency, and supports continuous memory refinement for lifelong learning.

  • Contribution: UMEM jointly optimizes memory extraction and management rather than treating them as static or decoupled processes.This framework is designed for self-evolving agents.
  • Method: Semantic Neighborhood Modeling and GRPO with a Marginal Utility Reward align extracted memories with the agent’s management policy.The design mitigates accumulation of instance-specific noise.
  • Results: UMEM significantly outperforms highly competitive baselines in cross-task generalization and execution efficiency.The empirical results support improvements across both dimensions.
  • Implications: UMEM continuously refines the memory bank during interaction, offering a paradigm for lifelong learning in open-ended environments.This continuous refinement is presented as a robust approach to lifelong learning.

A. Implementation Details · B. Mem-Optimizer Action Template

UMEM trains the Mem-Optimizer with GRPO using specified sampling, neighborhood, retrieval, regularization, and optimization settings. Its actions use a structured experience format that separates extracted memory content from the memory-evolution operation.

  • A. Implementation Details: 128 training queries are sampled per update, with G=8 rollouts generated for each query over 3 training epochs.These settings define the batch construction and rollout budget for GRPO optimization.
  • A. Implementation Details: Semantic neighborhoods use Top-N=3 neighbors, while memory evolution retrieves memories with Top-K=3.Both neighborhood construction and memory retrieval use three-item limits.
  • A. Implementation Details: KL regularization uses coefficient β=0.001, and policy updates use clipping ratio ϵ=0.2.These settings control regularization and clipping during GRPO training.
  • A. Implementation Details: The learning rate is set to 1 × 10−6.This is the reported optimization step size for Mem-Optimizer training.
  • B. Mem-Optimizer Action Template: Each Mem-Optimizer action is represented as a structured output following a defined template.The template organizes each action into an experience wrapper with separate value and operation fields.
  • B. Mem-Optimizer Action Template: The action template is <experience><value>...</value><operation>...</operation></experience>.This exact format specifies the required structure of each optimizer action.
  • B. Mem-Optimizer Action Template: The <value> field stores extracted memory content, while <operation> specifies the memory-evolution decision, such as addition or replacement.The two fields respectively encode what to remember and how to update the memory bank.

C. Theoretical Analysis · C.1. Cosine Neighborhood as a Proxy for Reuse-Semantic Proximity · D. Prompt Templates

The theoretical analysis justifies cosine neighborhoods as proxies for reuse-semantic proximity because high cosine similarity stabilizes retrieval scores and rankings. The prompt templates operationalize memory reuse across mathematical reasoning, multiple-choice, and question-answering tasks through task-specific inputs, instructions, and output formats.

  • C.1. Cosine Neighborhood as a Proxy for Reuse-Semantic Proximity: The lemma assumes ℓ2-normalized embeddings when analyzing retrieval-score stability under cosine proximity.It applies to any two queries and any candidate key.
  • C.1. Cosine Neighborhood as a Proxy for Reuse-Semantic Proximity: High cosine similarity between queries guarantees nearly identical relevance scores for any candidate memory key.This score stability underpins the proxy relationship between cosine proximity and reuse-semantic proximity.
  • C.1. Cosine Neighborhood as a Proxy for Reuse-Semantic Proximity: Stable retrieval scores produce highly overlapping rankings and similar Top-K memory sets for nearby queries.The cosine neighborhood therefore captures future queries likely to retrieve and reuse the same memory.
  • D. Prompt Templates: The framework provides detailed instruction templates for both the Memory Optimizer and Executor LLM.The Memory Optimizer refines and organizes retrieved past experiences, while Executor prompts enforce domain-specific output formats.
  • D. Prompt Templates: Mathematical reasoning prompts combine past experiences with the question and require step-by-step derivations followed by a boxed final answer.The system instructions direct agents to identify concepts, retrieve relevant methods, and show calculations clearly.
  • D. Prompt Templates: Multiple-choice prompts supply past experiences, the question, and options, then instruct the agent to evaluate candidates and select the best answer.The templates distinguish index-based and letter-based task outputs and wrap the selected answer in a boxed format.
  • D. Prompt Templates: Question-answering prompts combine context, the question, and past experiences, requiring contextual reasoning before a boxed final answer.They direct the agent to extract relevant facts, apply successful reasoning patterns, and synthesize the information.

E. Case Study

The case study shows that retrieved experiences help UMEM interpret implicit task requirements, transfer cleaning procedures across analogous tasks, and complete tasks more efficiently than a memoryless baseline. UMEM’s agent self-corrects during exploration, while the baseline repeatedly performs ineffective pick-and-place actions without cleaning.

  • UMEM Enhanced Agent: UMEM retrieves analogous cleaning experiences and applies the generalized sequence locate object → pick up → go to sinkbasin → clean → place on target.The agent initially explores incorrect locations and picks up the wrong object, then self-corrects after discovering the cloth.
  • Baseline: Without prior experiences, the baseline treats the task as simple pick-and-place and never invokes the clean command, even after viewing the relevant help action.This exposes its inability to connect an available action to the implicit task requirement.
  • Key Insights: 13 steps complete the task with UMEM-extracted memories, whereas the baseline exhausts 30 steps in a repetitive action loop without solving it.The case study attributes this difference to better semantic understanding, meaningful exploration, and transfer from cleaning plates, knives, and pans.

F. Procedure for Evolutionary Memory Management

UMEM evolves the Mem-Optimizer and memory bank jointly through offline semantic-neighborhood construction and online GRPO-based memory evolution. Training evaluates candidate memory operations using structural correctness and neighborhood-level utility before updating both policy and memory.

  • Training Overview: Algorithm 1 jointly trains the Mem-Optimizer πϕ and evolves the memory bank B, producing trained parameters ϕ and an evolved B.The procedure takes query corpus D, frozen executor E, neighborhood size N, and group size G as inputs.
  • Semantic Neighborhood Modeling: Offline semantic neighborhood modeling retrieves the N nearest neighbors NN(q) for each query q from D \ {q}.Neighborhoods are identified before online optimization based on embedding similarity to prevent shortcut learning.
  • Online Memory Evolution: The frozen executor E first generates a memory-augmented trajectory τq, after which πϕ samples G candidate ADD or UPDATE operations.Candidates use the query, trajectory, and retrieved memory as rollout context.
  • Marginal Utility Reward: The marginal utility reward rg averages success-rate and efficiency gains across NN(q), alongside format reward rf for structural correctness.This neighborhood-level reward evaluates whether extracted memories generalize beyond the individual query.
  • Policy and Memory Updates: GRPO updates πϕ with group-relative advantages from the combined rewards rf + rg, enabling stable refinement without a critic network.The memory bank is then updated by committing the best-performing operation from the rollout group.
Loading 2602.10652v1…