Self-Evolving LLM Memory Extraction Across Heterogeneous Tasks
| Authors | Yuqing Yang et al. |
| Year | 2026 |
| HF Upvotes | 6 |
| arXiv | 2604.11610 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
As LLM-based assistants become persistent and personalized, they must extract and retain useful information from past conversations as memory. However, the types of information worth remembering vary considerably across tasks. We formalize the heterogeneous memory extraction task and introduce BEHEMOTH, a benchmark that repurposes 18 existing datasets spanning personalization, problem-solving, and agentic tasks, using a downstream utility-driven metric for systematic evaluation. Our empirical analysis confirms that no single static extraction prompt dominates across all task categories, and that existing self-evolving prompt optimization frameworks, originally designed for homogeneous distributions, degrade when training tasks are heterogeneous. To address this, we propose CluE, a cluster-based self-evolving strategy that groups training examples into clusters by extraction scenarios, analyzes each cluster independently, and synthesizes cross-cluster insights to update the extraction prompt. Experiments on BEHEMOTH show that CluE generalizes effectively across heterogeneous tasks (+9.04% relative gain), consistently outperforming prior self-evolving frameworks.
Engineering Breakdown
Plain English
This paper tackles the problem of extracting and storing useful information from LLM conversations as persistent memory, where different types of conversations (personalization, problem-solving, agentic tasks) require remembering different kinds of information. The authors created BEHEMOTH, a benchmark combining 18 existing datasets with a utility-driven metric to evaluate what an LLM should remember across these heterogeneous task categories. They discovered that no single extraction prompt works well across all task types, and that existing self-evolving prompt optimization methods actually get worse when trained on diverse task distributions. To fix this, they propose CluE, a clustering-based approach that adapts memory extraction strategies per task cluster rather than using a one-size-fits-all prompt.
Core Technical Contribution
The core insight is that memory extraction is fundamentally heterogeneous—different tasks require remembering fundamentally different signals—and existing homogeneous prompt optimization frameworks fail under this distribution shift. The authors formalize this as a multi-task learning problem where prompt strategies must be task-aware rather than universal. CluE's main innovation is using task clustering to group similar memory extraction patterns together, then optimizing specialized extraction prompts for each cluster rather than searching for a global optimum. This is a departure from prior work that assumed a single high-quality extraction strategy could work everywhere, introducing task-conditional memory extraction as a necessary design pattern for persistent LLM assistants.
How It Works
CluE operates in phases: first, it clusters downstream tasks based on their inherent characteristics and the type of information actually useful for solving them (this clustering likely happens offline based on task representations or metadata). For each task cluster, the system maintains a separate extraction prompt that is continuously optimized using self-evolution techniques—meaning the prompts iteratively improve based on utility feedback (did remembering this information actually help the model perform better on downstream tasks?). When a new conversation arrives, the system first categorizes which cluster it belongs to, then applies the corresponding extraction prompt to decide what to remember. The key mechanism is the utility-driven evaluation: instead of measuring extraction quality in isolation, it directly measures whether extracted memory actually improves downstream task performance, ensuring only genuinely useful information is retained.
Production Impact
For teams building persistent LLM assistants (like customer support bots or personalized agents), this work directly addresses a critical problem: memory bloat and poor generalization. Rather than storing everything or using naive summarization, you can now adopt cluster-aware memory extraction that adapts to your specific mix of task types. In a production system, you'd profile your conversation types upfront (which conversations are personalization-heavy, which are technical problem-solving, etc.), train CluE's clustering offline, then deploy task-aware extraction prompts that run in-line during conversations. The trade-off is added complexity: you need to maintain multiple extraction prompts and perform task classification on each conversation, increasing latency slightly and requiring more careful prompt management. However, the memory quality improvement likely reduces hallucinations and context window waste downstream, since the model won't be drowning in irrelevant historical details.
Limitations and When Not to Use This
The paper assumes task clustering is stable and pre-definable, which may not hold for novel or emerging task types not seen during training. It also requires strong utility-driven evaluation signals (i.e., you must be able to measure whether extracted information actually helped), which is feasible for closed-loop tasks but harder for open-ended assistant scenarios where impact is subjective. The approach may not generalize well across very different domains (e.g., training on customer support data but deploying on medical consultations) without expensive re-clustering. Additionally, the paper doesn't address the computational cost of task classification at inference time or how to handle conversations that span multiple task clusters—real conversations are often messy combinations of personalization, problem-solving, and agentic planning.
Research Context
This work builds on the growing literature of in-context learning and memory-augmented LLMs, extending ideas from prompt optimization (like prompt gradient descent and in-context learning) to the heterogeneous setting. It relates to multi-task learning and domain adaptation in deep learning, but applies those insights to the discrete prompt space rather than weight adaptation. The BEHEMOTH benchmark consolidates prior work on personalization (user preference extraction), problem-solving (domain knowledge retention), and agent memory (action-outcome tracking) into a unified evaluation framework. The research opens a new direction: task-aware in-context learning, where the structure of future interactions shapes what gets encoded and stored from past interactions.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::