Speculate While You Reason: Teaching Agents to Predict Their Next Tool Call via Joint Agent-Speculator RL
Jiabao Ji, Yujian Liu, Li An, Rohit Jain, Gungor Polatkan, Siyu Zhu, Shiyu Chang
cs.AI
2026-07-28
把工具调用预测器合并进Agent自身、复用同一份KV cache,联合RL训练后Qwen3-4B的下一调用预判命中率从44.1升到61.2,任务成功率不掉。
LLM agent解题时,很大一块墙钟时间花在等工具返回上。发了个搜索请求,模型干等结果回来才能继续推理。tool-call speculation(工具调用投机)想藏掉这段延迟:趁agent还在推理,提前猜它会发的下一个工具调用、预先跑起来,等agent真发出来时结果已经备好。前提是猜得准。
现有的speculator通常是单独的draft model或缓存的历史trace,跟实际部署的agent行为对不齐,它猜的下一个调用常常不是agent真正会发的那一个。这篇把它叫speculator-agent gap。
核心观察是被预测的那个agent自己就是最好的下一个调用预测器。作者把agent和speculator塞进同一个模型,做成self-speculating agent(自投机agent)。同一个模型两种模式:agent模式正常解题、发工具调用;speculator模式从共享的trajectory前缀出发,提前预测下一个调用,且完全复用前面已算好的prefix KV cache,不另开一份。
要训出这种双模式又不拖累agent本身,他们提了joint agent-speculator RL(联合RL)。先从agent自己的rollout里抽出它真实会发的下一个调用,作为speculator的训练目标;然后按4:8的步数比交替更新agent和speculator,每次切换重置优化器状态,并用SFT做warmup。speculator的奖励是工具名匹配加参数的macro token-F1。
在多跳搜索QA和τ-bench对话式工具调用两类任务上,Hit@1(第一个预测就命中的比例)显著上升:
| 模型 | 起始(SFT) | 联合RL后 | 任务成功率变化 |
| Qwen3-4B | 44.1 | 61.2 | 26.6 → 27.7 |
| Qwen3.5-4B | 48.9 | 66.3 | 49.2 → 50.6 |
分项看,Qwen3-4B在HotpotQA上Hit@1从23.2到55.2,在τ-bench Airline上从35.6到68.0。省内存是结构性的:在MuSiQue上,自投机只花8.3秒、8.70GB;挂一个外部0.6B draft model要27.4秒、10.71GB,挂1.7B要33.1秒、12.76GB。自投机不需要第二套参数和KV cache,这是它相对外部draft model的根本优势。
对做agent推理引擎的人,这条路省掉了一个独立draft model的全部显存和调度成本,代价只是同一个4B模型多训出一种预测模式。它属于speculative decoding思想在工具调用场景的落地。工程上能直接接到现有带prefix cache的推理框架里。
作者自己划了三道边界。第一,它假设投机执行的调用不改变外部状态,下单、改数据库、发消息这类会mutate环境的工具不能盲跑,得配套dry-run模式。第二,只测了搜索QA和τ-bench两类任务,代码执行、网页浏览、长流程、多agent协作都没碰。第三,训练实验都在4B规模,更大的模型优化动态可能不同。
还有一个没被量化的点:论文报的是Hit@1命中率,但命中率提升多少能换回多少实际墙钟时间,在不同工具延迟分布下差别很大,文中没给端到端的延迟收益曲线。