Task-Focused Memorization for Multimodal Agents

TaskMem: Learning What to Memorize via Reinforcement Learning for Multimodal Agents

Most AI agents attempt to remember everything. This inevitably leads to a combinatorial explosion of irrelevant data. In multimodal agents—systems that perceive and act within physical or virtual environments—this results in cluttered, noisy long-term memory. Such noise obscures the signal needed for actual task execution.

Currently, memory construction in these agents is largely a passive process. Existing frameworks typically rely on heuristic summarization or fixed templates to decide what to store. They treat memory as a secondary transcription task rather than a strategic decision. This approach fails to address the "AI Frame Problem" (the difficulty of identifying contextually relevant information without being overwhelmed by possibilities).

The authors of this paper argue that memory should be an active, goal-driven policy. They introduce TaskMem. This framework moves away from static summarization. Instead, it uses a learnable policy that dynamically selects what to retain based on the agent's current environment and objectives.

The Problem

The status quo for multimodal long-term memory is essentially a glorified transcription service. Current methods aim for high fidelity (accuracy regarding the video). However, they lack intentionality. As shown in, an agent tasked with housework needs a different memory profile than an agent tasked with user interaction.

Because current models do not explicitly optimize for utility, they suffer from two main failures. First, they consume excessive context by storing redundant or trivial details. Second, they fail to capture the specific nuances required for specialized tasks. Currently, memory generation is disconnected from the actual task execution loop. The agent "remembers" in a vacuum. It only realizes the memory was useless when a downstream task fails.

How It Works

TaskMem treats memory generation as a learnable policy $\pi_\theta$ that decides what to store at each timestep $t$. The authors implement this via a two-phase training paradigm .

1. Phase One: Fundamental Memorization. Before deployment, the policy undergoes reinforcement learning (RL) using Group Sequence Policy Optimization (GSPO). The goal here is "how" to memorize. The model is optimized against a multi-objective reward function. This function enforces factual accuracy, non-redundancy, and proper formatting. Crucially, the authors introduce a "richness reward" to prevent the policy from hacking the objective. Without this, the policy might generate extremely short, safe, but uninformative summaries .

2. Phase Two: Task-Focused Adaptation. Once deployed, the agent must learn "what" to memorize. To avoid the high cost of full-parameter fine-tuning (updating all model weights) during inference, the authors use a parameter-efficient approach. They insert a lightweight, additive trainable vector adapter (only 2,048 parameters) at a single transformer layer. Using Direct Preference Optimization (DPO), the model is tuned on candidate memories. A reward model compares these candidates based on recent, observed tasks. This turns sparse task feedback into dense preference signals.

The architecture relies on persistent identities. It tags faces with global IDs and overlays ASR (Automatic Speech Recognition) transcripts. This ensures the policy maintains coherence across long video streams without constant re-parsing of raw pixels.

Numbers

The authors test TaskMem on reformulated streaming versions of VideoMME, EgoLife, and EgoTempo. Built on a Qwen3-VL-30B-A3B backbone, the paper reports significant gains. TaskMem improves VQA (Video Question Answering) accuracy by 6.3%, 7.0%, and 5.3% on these benchmarks, respectively. These deltas represent a measurable increase in the agent's ability to solve tasks using only its generated memory.

The efficiency of the Phase Two update is vital for production. The authors found that the adapter direction converges early. By scaling the weights of a 10-step adapter to a target norm, they achieved performance comparable to a 40-step run. This reduces training data and time requirements by approximately 75% (see Table 1). Furthermore, the ablation study in shows that adapter placement matters.

Inserting the adapter in shallow or middle layers (specifically layer 22) is more effective than deep layers.

What's Missing

While the results are compelling, there are gaps for practitioners to consider.

• Fine-grained Visual Blind Spots: Error analysis in Tables 16 and 17 reveals that the policy still struggles with granular visual evidence. Examples include specific numbers on small objects or subtle text. If your agent's success depends on reading tiny serial numbers, TaskMem may not suffice.
• Relational Complexity: The authors note difficulties in integrating complex temporal and relational structures. As shown in the "hard case" analysis [Table 18], the model can capture local visual cues. However, it sometimes fails to weave them into a coherent semantic narrative of how things relate over time.
• Reward Model Dependency: Phase Two logic hinges on the quality of the task-relevance reward model (powered by GPT-4o or Gemini). There is a risk of a feedback loop. The agent might learn to generate memories that please the evaluator model rather than memories that truly assist a physical executor.

Should You Prototype This

Yes, specifically if you are building embodied agents or long-horizon video assistants.

The use of a 2,048-parameter adapter is a win for production. It allows for task-specific personalization without the OOM (Out of Memory) risks of full-parameter fine-tuning. The ability to update the agent's focus online using only a handful of recent tasks is highly viable. This is useful where continuous retraining is not an option.

However, do not expect it to solve low-level perception errors. If your bottleneck is the base MLLM's inability to see a specific object, TaskMem will only help you "remember" that you missed it. Use this to manage the strategy of memory, not to compensate for poor vision. Code is reportedly available; see the project page at https://taskmem.github.io/.

Figures from the paper