Just-in-Time Memory: Learning to Curate Task-Adaptive Memory for LLM Agents
Yefan Zhou, Yang Li, Zeyu Leo Liu, Semih Yavuz, Shafiq Joty
cs.AI
2026-09-23
JitMem stores raw traces and curates a task-adaptive payload at read time, gaining 16.2, 16.3 and 3.9 success points on ALFWorld, WebShop and τ²-bench.
Agents should reuse their own past work. The default is to distill a finished trajectory into a reflection, skill, workflow, or strategy, store that artifact, and retrieve it by similarity. Distillation happens before the next task is known. One household trace may teach a heat-then-fridge routine and a sofa-placement check; write-time memory keeps one summary and throws the rest away. Learning a write-time curator is also awkward: the value of a storage decision may show up many tasks later, so methods invent task grouping to manufacture a delayed signal.
JitMem delays shaping until read time. The bank stores raw traces. Once the new task is known, a curator writes a short payload from those traces and the query, and the executor consumes it on the same task.
Four steps. BM25 retrieves top-k raw traces by task text only (default k=3). The curator reads the query and those traces and writes a briefing: which episodes matter, which tactics worked, what to watch on this task. A frozen executor prepends the payload and acts. An executor-as-judge gates insertion of the new raw trace. The payload is ephemeral.
Training is GRPO. For each train task the curator samples a group of payloads against a frozen bank; the frozen executor returns the environment's true reward (binary on ALFWorld and τ², continuous on WebShop); advantages are group-mean-centered. No delayed return, no task grouping. The train bank is collected once from the bare executor with ground-truth success labels and then frozen. At test the bank starts empty, updates per batch, and scores average over task orders.
The curator starts from Qwen3-8B with thinking off, 100 GRPO steps, lr 1e-6, batch 32, group 8, Qwen3-8B as the training executor. The same curator is reused with Gemini-2.5-Pro or GPT-5.4 executors without retraining.
With a Qwen3-8B executor, trained JitMem reaches 77.4 SR on ALFWorld versus 61.2 for SkillOS (+16.2) and 47.9 with no memory. On WebShop it reaches 32.8 versus 16.5 for SkillOS and 9.8 with no memory. With Gemini as both zero-shot curator and executor, JitMem-gemini hits 61.0 SR on WebShop versus 41.0 for SkillOS, so task-adaptive read-time shaping is already a large share of the gain. On τ²-bench with GPT-5.4, JitMem-gpt micro-averages 75.6 versus 71.7 for ReasoningBank (+3.9), almost entirely from Telecom (72.6 vs 61.6). Airline and Retail do not move beyond the no-memory agent.
| Setting | ALFWorld SR | WebShop SR |
| No memory (Qwen3-8B) | 47.9 | 9.8 |
| SkillOS (trained) | 61.2 | 16.5 |
| JitMem (trained) | 77.4 | 32.8 |
Untrained JitMem-base already reaches 60.5 on ALFWorld, next to trained SkillOS. A curator trained with Qwen as executor transfers to GPT-5.4 at 86.7 ALFWorld SR, 1.4 below a curator trained with GPT-5.4 (88.1). Versus write-time methods, executor input tokens fall 50.3%–56.3% and steps 28.4%–31.4%; versus no memory, steps fall from 17.8 to 11.6. Ablations: dropping task conditioning costs up to 11.4 / 10.4 SR after RL; emptying retrieval drops below the untrained curator (up to -14.8 / -15.2); write-time distillation before storage costs 6.8–8.2 SR on WebShop. Warm-starting the test bank with 100 traces does almost nothing.
Memory design is often framed as which format to store. This paper moves the lever to when to shape. Read-time shaping lets one trace serve many downstream queries and turns curator RL into a single-step problem. For a deployed agent memory, an untrained read-time curator is already worth swapping in; a GRPO pass then shortens payloads and trajectories. Decoupling curator from executor means one small curator can feed several large executors.
Null gains on Airline/Retail are a boundary: this helps procedural guidance more than fact lookup.
The authors list a simple BM25 retriever, an extra LLM call per task, and hand-written payload schemas. The empty test bank is a cold start that warm-start barely fixes. Train and test banks are deliberately shifted; a staged refresh adds only 2.8 SR for Qwen. τ² has no standard train split, so only zero-shot curators are reported. Absolute WebShop SR with Qwen remains low (32.8): large gains, low ceiling. The executor-as-judge can pollute the bank with false successes.