Just-in-Time Memory: Learning to Curate Task-Adaptive Memory for LLM Agents

Just-in-Time Memory. Learning to Curate Task-Adaptive Memory for LLM Agents. By Yefan Zhou and four others, from Salesforce AI Research. Published on the twenty-third of September. No code repository is linked to it.

Agentic memory systems reuse past experience to improve future performance, but most existing designs curate memory at write time. Once a task is completed, its trajectory is distilled into a fixed artifact, such as a reflection, a workflow, a skill or a reasoning strategy, and later retrieved by similarity. That forces the system to decide what to keep before the future query is known. Information is discarded irreversibly, and what remains is a query-independent summary that has to serve many possible downstream tasks.

Learning such a curator is difficult for a second reason. The value of a storage decision may only become apparent when a relevant query arrives, potentially many tasks later. That makes it a long-horizon credit-assignment problem. The curator makes a choice now and is only told much later whether it was a good one, if it is told at all.

Just-in-Time Memory, or JitMem, moves curation to read time. Raw trajectories are retained as they are. When a new task arrives, a memory curator takes the retrieved traces together with that task, and synthesizes a compact payload adapted to the immediate need. Nothing is thrown away in advance. The summarizing still happens, but only once it is known what the summary is for.

The same timing is what makes the curator trainable. The payload is consumed on the task it was built for, so the curator can be trained directly from immediate task success. That avoids delayed utility signals, and it avoids the need to group related tasks artificially. The contrast with write-time curation is direct. There, the payoff of a decision arrives many tasks later. Here, it arrives on the same task, in the same episode. Did the agent succeed with this payload or not? That answer is available at once.

JitMem is evaluated on ALFWorld, WebShop and tau squared bench. It consistently outperforms agents with no memory, as well as heuristic and learned write-time memory methods. Over the strongest baseline, it improves the absolute success rate by sixteen point two points on ALFWorld, sixteen point three on WebShop, and three point nine on tau squared bench.

Even an untrained curator is already competitive with those baselines or better, which points to task-adaptive curation at read time as a major source of the gain in itself. Training the curator then compounds the improvement. So the design choice and the training each contribute. Moving curation to read time does much of the work before any learning, and learning adds more on top.

The shift is from summarizing for every possible future to summarizing for the one that has arrived. Waiting until read time keeps the raw record intact, and it turns a credit-assignment problem spread over many tasks into one that resolves on the task at hand.