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
Salesforce把agent记忆从写时蒸馏改成读时按当前任务整理原始轨迹。JitMem在ALFWorld、WebShop、τ²-bench上比最强基线分别高16.2、16.3、3.9个成功率点，未训练的读时整理已经能打平或超过写时方法。
Agent要跨任务复用自己的经历。主流做法是任务一结束就把轨迹蒸馏成反思、技能、工作流或推理策略，存进库，下次按相似度取出来。蒸馏发生在未来任务还不知道的时候。一条轨迹可能同时教「先加热再放冰箱」和「确认沙发上能放报纸」，写时只能留下一种摘要，丢掉的细节以后补不回来。写时策展器的奖励还要等很多步之后才出现，不得不靠人为把相关任务捆在一起制造信号。
JitMem把整理推迟到读的那一刻。库里只存原始轨迹，当前任务已知以后，再由策展器根据检索到的轨迹和题目合成一份短payload，立刻交给执行器。
四步循环。BM25只按任务描述检索top-k原始轨迹（默认k=3）。策展器读题目和这些轨迹，写出一份自然语言简报：哪些经历相关、什么策略有效、当前任务要注意什么。冻结的执行器把payload前置进prompt去做事。成功与否由执行器当judge，通过才把本轮原始轨迹追加进库。payload本身不入库。
训练用GRPO。每个训练任务从固定银行里取出轨迹，策展器采样一组payload，冻结执行器跑出环境真奖励（ALFWorld/τ²为0/1，WebShop为连续分），组内减均值当优势。没有延迟回报，也不需要任务分组。训练银行是无记忆执行器在训练集上跑一遍、用真标签留下的成功轨迹，训练期间冻结，避免检索内容跟着策略漂。测试时银行从空开始，按批次更新，多种任务顺序取平均。
策展器从关思考模式的Qwen3-8B初始化，100步、lr 1e-6、batch 32、group 8，训练执行器用Qwen3-8B。测试执行器换Gemini-2.5-Pro或GPT-5.4时，策展器不再训。
Qwen3-8B执行器上，训练后的JitMem在ALFWorld成功率77.4，对SkillOS的61.2为+16.2，对无记忆的47.9更高；WebShop成功率32.8，对SkillOS的16.5为+16.3，无记忆只有9.8。同一套Gemini执行器加Gemini零样本策展时，WebShop上JitMem-gemini为61.0，SkillOS-gemini为41.0，读时按任务整理本身就贡献一大截。τ²-bench用GPT-5.4执行器，JitMem-gpt微平均75.6，比ReasoningBank的71.7高3.9，增益几乎全在Telecom（72.6对61.6）；Airline和Retail上记忆方法普遍抬不动无记忆基线。
| 设置 | ALFWorld SR | WebShop SR |
| 无记忆（Qwen3-8B） | 47.9 | 9.8 |
| SkillOS（训过） | 61.2 | 16.5 |
| JitMem（训过） | 77.4 | 32.8 |
未训练的JitMem-base在ALFWorld已到60.5，接近训过的SkillOS。用Qwen训好的策展器接到GPT-5.4执行器，ALFWorld 86.7，和用GPT-5.4当训练执行器的88.1只差1.4。相对写时方法，执行器输入token少50.3%–56.3%，步数少28.4%–31.4%；相对无记忆，步数从17.8降到11.6。消融：拿掉任务条件，训练后最多掉11.4（ALFWorld）和10.4（WebShop）；检索强制为空，掉到未训练水平以下（最多-14.8/-15.2）；写时先蒸馏再存，WebShop掉6.8–8.2。测试银行预热100条几乎没帮助。
记忆系统的关键变量经常被写成「存什么格式」，这篇把它改成「什么时候整形」。读时整形让同一条轨迹服务不同下游任务，也让策展器的强化学习变成单步题。对要上线的agent记忆，未训练的读时策展已经值得先换上；再花一次GRPO，payload更短、步数更少。策展器和执行器解耦，一个小策展器可以喂多个大执行器。
τ²上Airline/Retail没有增益，提醒读时整理擅长的是程序性指南，不是查一条事实。
作者列了三条：BM25在库变大后可能成为瓶颈；每任务多一次LLM调用；payload格式按基准手写。此外，测试银行从空开始有冷启动，只是预热也补不了多少；训练银行和测试银行分布有意错开，分阶段刷新只再涨2.8 SR（Qwen）。τ²没有标准训练集，只报了零样本策展。WebShop上Qwen执行器绝对成功率仍低（32.8），增益大、天花板不高。Judge用执行器自己当过滤器，错判成功会污染库。