💬 NLP

UniMem: Complementary Episodic-to-Parametric Memory for Boundary-Agnostic Task Streams

UniMem is a self-routing framework that addresses the stability-plasticity dilemma in boundary-agnostic task streams by dynamically coordinating an episodic buffer for novel tasks and an expandable parametric memory for recurring patterns, thereby achieving superior performance without requiring explicit task labels or fixed parameter budgets.

Original authors: Siyu Xia, Chenheng Zhang, Yanting Wu, Haoxuan Li, Jiajun Chai, Xiaohan Wang, Guojun Yin, Wei Lin, Zhouchen Lin, Haifeng Zhang, Jun Wang

Published 2026-07-29

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Siyu Xia, Chenheng Zhang, Yanting Wu, Haoxuan Li, Jiajun Chai, Xiaohan Wang, Guojun Yin, Wei Lin, Zhouchen Lin, Haifeng Zhang, Jun Wang

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 super-smart robot how to do chores. You want it to learn from every single task it encounters, from folding laundry to fixing a leaky faucet. But here's the catch: the robot has a brain that is already full of general knowledge, and you can't easily rewrite its entire brain every time it learns something new. This is the world of Large Language Models (LLMs), the AI brains behind many of the chatbots and assistants we use today.

To make these robots smarter, scientists usually try two main tricks. The first is like giving the robot a notebook (External Memory). When it faces a new problem, it flips through its notebook to find a similar example it saw before. This is great for learning new things quickly, but it's slow because the robot has to stop and search every time, and it doesn't really "learn" the skill—it just copies the answer. The second trick is rewiring the brain (Parametric Memory). This is like practicing a skill until your muscles remember it. Once learned, the robot is fast and efficient, but it's hard to teach it new things without messing up what it already knows, and it needs to know exactly what it's learning before it starts.

The big question scientists are asking is: How do we build an AI that can handle a never-ending stream of new, mixed-up tasks without getting confused or slowing down? It needs to be flexible enough to grab a notebook for weird, one-off jobs, but smart enough to practice and "muscle-memory" the common ones. This paper, UniMem, proposes a solution inspired by how human brains actually work.

The Brain's Best Trick: A Two-System Team

Think of your own brain. When you have a wild, one-time dream about flying a pink elephant, you might remember it for a day or two, but you don't try to build a permanent muscle memory for it. That's your episodic memory—a temporary storage for specific, fresh experiences. But when you practice riding a bike for the first week, your brain slowly moves that skill into your procedural memory. Now, you can ride without thinking about it, and you don't need to look up instructions every time.

The authors of this paper noticed that AI agents face a similar problem. They are being asked to handle "boundary-agnostic" task streams. Imagine a robot assistant that doesn't know when one job ends and another begins. It just gets a constant flow of requests: "Write a poem," "Calculate a tip," "Debug this code," "Write a poem again," "Order pizza," "Write a poem again."

If the robot uses only a notebook (retrieval), it gets slow and clumsy when it has to look up the same "write a poem" instructions over and over. If it tries to rewire its brain for every single request, it gets confused and forgets how to order pizza because it's trying to learn a new poem style at the same time.

Enter UniMem: The Smart Traffic Cop

The researchers propose UniMem, a system that acts like a super-smart traffic cop for the robot's memory. Instead of forcing the robot to choose between a notebook or a rewired brain, UniMem lets it use both, switching between them automatically.

Here is how it works, using a playful analogy:

Imagine the robot has a Magic Ticket Booth (the Routing Tokens). Every time a new request comes in, the robot checks the request against its current "known tasks."

The "I've Seen This Before" Ticket: If the request looks familiar (like "Write a poem"), the Magic Ticket Booth hands the robot a special Key (a Parametric Memory Block). This key unlocks a dedicated, high-speed muscle memory for that specific task. The robot executes the task instantly and efficiently, just like you riding a bike.
The "What Is This?" Ticket: If the request is weird, rare, or the robot isn't sure (like "Write a poem about a toaster in the style of a pirate"), the Magic Ticket Booth sends the request to a Waiting Room (the Episodic Buffer). Here, the robot doesn't try to force a permanent memory. Instead, it pulls up a "cheat sheet" (retrieval) from its notebook to help answer the question.

The Magic of "Growing" Memory

The real magic happens when the robot starts seeing the "weird" requests over and over again.

In the Waiting Room, the robot keeps a list of all the "toaster pirate" requests it has seen. Once it sees enough of them to form a pattern (say, 50 or more), the system says, "Hey, this isn't a one-off anymore! This is a real skill."

At that moment, UniMem performs a consolidation. It takes that pattern, builds a new "Key" (a new Parametric Memory Block) specifically for "toaster pirate poems," and moves it from the Waiting Room to the permanent muscle memory. The robot now has a new, permanent skill without ever having to stop and relearn everything else.

If a request is truly rare and never comes back, it stays in the Waiting Room, and the robot just uses its notebook to handle it. This prevents the robot from wasting brain space on things it will never do again.

What the Experiments Showed

The researchers tested this system on a long stream of tasks, simulating a robot that had to handle everything from simple math to complex reasoning, with some tasks appearing often and others appearing only a few times.

They found that UniMem consistently outperformed other methods.

Compared to robots that just used a notebook (Retrieval), UniMem was faster and more accurate because it didn't have to search every time for common tasks.
Compared to robots that tried to rewire their whole brain for everything (Parametric Memory), UniMem didn't get confused or forget old skills. It handled the mix of common and rare tasks much better.

In one specific test with a model called LLaMA-3.2-3B, UniMem improved the accuracy of getting the exact right answer by about 4.0 points on average compared to other methods. When they tested it on a stream of 100 different tasks, UniMem maintained high performance, while other methods started to struggle or forget.

The Verdict

The paper suggests that by copying the human brain's strategy of having a temporary "scratchpad" for new things and a permanent "skill library" for common things, we can build AI agents that are both flexible and efficient. UniMem doesn't just store data; it learns when to store it permanently and when to keep it temporary.

It's a step toward creating AI that can grow and adapt in the real world, where tasks are messy, mixed up, and constantly changing, without needing a human to tell it exactly what to learn next. The system suggests that the future of AI memory isn't about choosing one big brain or one big notebook, but about having a smart system that knows how to use both.

Technical Summary: UniMem

Problem Statement

The paper addresses the stability-plasticity dilemma faced by Large Language Model (LLM) agents deployed in real-world, boundary-agnostic, and evolving task streams. Existing memory paradigms fall into two categories, each with significant limitations in this setting:

External Memory (e.g., RAG): Offers high plasticity by allowing immediate incorporation of new evidence without gradient updates. However, it fails to internalize recurring execution patterns, leading to suboptimal performance on repeated tasks and incurring inference-time retrieval overhead.
Parametric Memory (e.g., PEFT, LoRA): Provides stable and efficient execution by internalizing task knowledge into trainable parameters. However, it typically relies on explicit task boundaries and fixed parameter budgets, making it difficult to adapt to heterogeneous streams where task composition and scale change dynamically.

Relying solely on either approach results in a trade-off where agents either lack generalization on recurring patterns or cannot efficiently accommodate new, sparse tasks. The authors frame this as a challenge of managing heterogeneous, boundary-agnostic, and evolving task streams without predefined task labels.

Methodology

Inspired by Complementary Learning Systems (CLS) theory—which posits that the human brain resolves the stability-plasticity trade-off through complementary episodic storage (hippocampus) and gradual consolidation into long-term knowledge (neocortex)—the authors propose UniMem.

UniMem is a self-routing framework that decouples task identification from task execution, enabling autonomous memory expansion. The core components are:

1. Self-Routing Architecture

UniMem utilizes learnable routing tokens as memory controllers to select the appropriate memory pathway for each incoming query.

Routing Tokens ( $E$ ): A matrix containing tokens for known tasks ( $e_k$ ) and a Novelty Sentinel ( $e_{NOVEL}$ ).
Routing Mechanism: Given an input query $q$ , the model computes a routing distribution over the tokens. A query is routed to a known task token $e_k$ only if its probability exceeds both the novelty sentinel's probability and a confidence threshold $\tau_{route}$ . Otherwise, it is routed to the novelty pathway.
Decoupled Execution:
- Known Tasks: When a known token is selected, a corresponding Parametric Memory Block ( $P_k$ ) is activated. In this work, $P_k$ is instantiated as Procedural KV Memory (learnable, layer-wise key-value pairs) injected via cross-attention into the frozen LLM backbone.
- Novel Tasks: Queries routed to $e_{NOVEL}$ are stored in an Episodic Buffer and handled via Retrieval-Augmented Generation (RAG) for execution, avoiding premature parameter allocation.

2. Episodic-to-Parametric Lifecycle

The system operates in two phases to manage memory expansion:

Phase I: Initialization: The system learns routing tokens for an initial set of known tasks and calibrates the novelty sentinel using a subset of pseudo-unseen tasks to distinguish between known and novel queries.
Phase II: Autonomous Discovery and Consolidation:
- Buffering: Novel queries accumulate in the Episodic Buffer.
- Discovery: When the buffer reaches capacity, UniMem triggers unsupervised task discovery. It uses Normalized Compression Distance (NCD) to measure procedural similarity and HDBSCAN to form boundary-agnostic candidate clusters.
- Filtering: Clusters undergo embedding-based outlier filtering and cohesion checks.
- Consolidation: High-quality clusters are promoted to new parametric memory units. A fresh routing token is allocated and paired with a new Procedural KV Memory block, which is then optimized via Supervised Fine-Tuning (SFT) on the clustered samples. Sparse or low-quality clusters remain in the episodic buffer.

Key Contributions

Problem Formulation: The authors formalize the challenge of memory management for LLM agents in heterogeneous, boundary-agnostic streaming environments, highlighting the limitations of existing single-paradigm approaches.
UniMem Framework: They propose a unified memory system that autonomously routes rare tasks to episodic memory (RAG) and consolidates frequent patterns into expandable parametric memory, inspired by CLS theory.
Empirical Validation: Extensive evaluations across diverse long-horizon streaming settings demonstrate that UniMem consistently outperforms baselines while maintaining robust task execution and memory stability.

Experimental Results

The authors evaluated UniMem on two benchmarks: Super-Natural Instructions (SNI) (scaling from 10 to 100 tasks) and SuperGLUE Mixed Stream (boundary-agnostic reasoning tasks).

Performance: UniMem achieved an average gain of 4.0 Exact Match (EM) points across three backbone models (LLaMA-3.2-3B, LLaMA-3.1-8B, Qwen3-8B) on SNI, outperforming baselines including Zero-Shot, RAG, Replay LoRA, and the token-centric TOKMEM.
Comparison with TOKMEM: UniMem significantly outperformed TOKMEM (e.g., 47.56% EM vs. 40.54% on LLaMA-3.2-3B with 100 tasks), validating the benefit of decoupling lightweight routing tokens from high-capacity procedural memory blocks.
Comparison with Replay LoRA: UniMem improved average EM by +6.20 and +2.44 points over Replay LoRA on LLaMA-3.2-3B and LLaMA-3.1-8B, respectively, demonstrating better handling of heterogeneous task patterns without catastrophic forgetting.
Routing Accuracy: The routing layer maintained high accuracy (e.g., >85% for LLaMA-3.1-8B at 100 tasks), effectively identifying appropriate parametric memory units.
Ablation Studies:
- Per-task vs. Shared KV: Task-specific KV memory significantly outperformed shared KV memory, especially as task diversity increased.
- Gating Mechanism: Disabling the learned gates for injecting KV memory caused substantial performance drops, confirming their necessity for regulating heterogeneous signals.
- Episodic Buffer: Removing the episodic buffer for sparse tasks led to performance degradation, highlighting its role in providing non-parametric coverage for long-tail queries.

Significance and Claims

The paper claims that UniMem offers a scalable direction for building self-evolving agents with expandable parametric memory. By separating task identification from execution and utilizing a novelty-guided lifecycle, UniMem allows agents to:

Handle heterogeneous task streams without explicit task labels.
Maintain execution fidelity while autonomously expanding memory.
Avoid premature parameter allocation for sparse or uncertain patterns.

The authors position UniMem as a solution that bridges the gap between the flexibility of external retrieval and the efficiency of parametric internalization, enabling LLM agents to operate effectively in open-world, streaming environments. They acknowledge limitations, such as the conservative consolidation strategy and the specific instantiation of parametric memory, suggesting future work could explore more adaptive criteria and different memory modules.