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
Folding the tool-call speculator into the agent itself and reusing one KV cache, joint RL lifts next-call Hit@1 from 44.1 to 61.2 on Qwen3-4B without hurting task success.
LLM agents burn a lot of wall-clock time waiting on tool calls. They fire a search request, then sit idle until the result comes back before they can reason again. Tool-call speculation tries to hide that latency: while the agent is still reasoning, predict the next tool call it will make, pre-execute it, and have the answer ready by the time the agent actually issues it. The catch is that the prediction has to be right.
Existing speculators are usually a separate draft model or a cache of historical traces, and they don't line up with how the deployed agent actually behaves. The call they predict is often not the call the agent would really issue next. The authors call this the speculator-agent gap.
The core observation is that the target agent itself is the best predictor of its own next call. So the authors fold agent and speculator into one model, a self-speculating agent. One model, two modes: agent mode solves the task and issues tool calls as usual; speculator mode branches off the shared trajectory prefix to predict the next call, fully reusing the prefix KV cache already computed, with no second copy.
To train this dual mode without degrading the agent, they use joint agent-speculator RL. Speculation targets come from the agent's own rollouts: the call it actually issued next. Agent and speculator updates alternate on a 4:8 step ratio, with optimizer state reset at each switch, and an SFT warmup before RL. The speculator reward is tool-name match plus a macro token-F1 over the arguments.
Across multi-hop search QA and tau-bench conversational tool use, Hit@1 (the share where the first prediction is correct) climbs sharply:
| Model | Start (SFT) | After joint RL | Task success |
| Qwen3-4B | 44.1 | 61.2 | 26.6 to 27.7 |
| Qwen3.5-4B | 48.9 | 66.3 | 49.2 to 50.6 |
On HotpotQA, Qwen3-4B's Hit@1 goes from 23.2 to 55.2; on tau-bench Airline, 35.6 to 68.0. The memory saving is structural: on MuSiQue, self-speculation costs 8.3s and 8.70GB, whereas bolting on an external 0.6B draft model costs 27.4s and 10.71GB, and a 1.7B draft costs 33.1s and 12.76GB. Self-speculation needs no second set of parameters or KV cache, which is its root advantage over external drafts.
For anyone building agent inference engines, this route drops the full memory and scheduling cost of a standalone draft model, at the price of training one extra prediction mode into the same 4B model. It is speculative decoding's idea applied to the tool-call setting rather than another leaderboard trick. It slots directly into existing inference frameworks that already keep a prefix cache.
The authors draw three boundary lines. First, it assumes speculative calls don't mutate external state; tools that place orders, edit databases, or send messages can't be blindly pre-run and would need a dry-run mode. Second, it's only tested on search QA and tau-bench; code execution, web navigation, long-running workflows, and multi-agent setups are untouched. Third, all training is at 4B scale; larger models may have different optimization dynamics.
One gap the paper leaves unquantified: it reports Hit@1, but how much hit-rate improvement actually buys back wall-clock time depends heavily on the latency distribution of the tools, and no end-to-end latency curve is given.