Command Palette
Search for a command to run...
Spéculer tout en raisonnant : apprendre aux agents à prédire leur prochain appel d’outil par apprentissage par renforcement conjoint agent–spéculateur
Spéculer tout en raisonnant : apprendre aux agents à prédire leur prochain appel d’outil par apprentissage par renforcement conjoint agent–spéculateur
Jiabao Ji Yujian Liu Li An Rohit Jain Gungor Polatkan Siyu Zhu Shiyu Chang
Résumé
Les agents fondés sur de grands modèles de langue passent souvent un temps d’horloge considérable à attendre les résultats des appels d’outils. La spéculation d’appel d’outil peut masquer cette latence en prédisant et en pré-exécutant le prochain appel d’outil de l’agent si la prédiction correspond à l’appel finalement effectué, mais les spéculateurs existants sont généralement des modèles d’ébauche séparés ou des traces en cache qui sont mal alignés avec le comportement propre de l’agent déployé. Nous identifions cet écart entre spéculateur et agent et montrons que l’agent cible lui-même constitue un bon spéculateur du prochain appel. Cela conduit à une conception plus simple : unifier l’agent et le spéculateur au sein du même modèle. Dans cet article, nous présentons l’agent auto-spéculant, un modèle unique qui, d’une part, résout des tâches en mode agent et, d’autre part, prédit son prochain appel d’outil à partir de trajectoires partielles en mode spéculateur, en réutilisant intégralement le cache KV de préfixe. Pour permettre ce mode dual sans dégrader les performances, nous proposons une méthode d’apprentissage par renforcement conjoint agent–spéculateur, qui dérive les cibles de spéculation des propres déploiements de l’agent et alterne les mises à jour de l’agent et du spéculateur. Sur des tâches de recherche agentique de questions-réponses et des tâches agentiques conversationnelles avec utilisation d’outils, notre méthode améliore le Hit@1 moyen du prochain appel d’outil de 44,1 à 61,2 pour Qwen3-4B et de 48,9 à 66,3 pour Qwen3.5-4B, tout en préservant le succès des tâches agentiques.
One-sentence Summary
Researchers from UC Santa Barbara and LinkedIn Inc. introduce a self-speculating agent that unifies the agent and speculator into a single model, trained via joint agent-speculator reinforcement learning to predict its own next tool call, thereby closing the speculator-agent gap and boosting next-call Hit@1 from 44.1 to 61.2 for Qwen3-4B and from 48.9 to 66.3 for Qwen3.5-4B without hurting task success.
Key Contributions
- The paper identifies a speculator–agent gap where external draft models or cached traces often fail to align with the agent’s own next tool call, adding serving overhead.
- It introduces a self-speculating agent that unifies task-solving and next-call prediction within a single model, reusing prefix KV cache to avoid separate model overhead.
- It proposes a joint agent-speculator reinforcement learning method that trains the speculator mode on the agent’s own rollouts via alternating updates, achieving next-call Hit@1 improvements from 44.1 to 61.2 (Qwen3-4B) and 48.9 to 66.3 (Qwen3.5-4B) while preserving downstream task success.
Introduction
Large language model agents increasingly rely on external tools such as search engines, code interpreters, and APIs, but waiting for tool responses often dominates end-to-end latency. Tool-call speculation can hide this latency by pre-executing likely future calls, yet existing speculators are built outside the agent, using smaller draft models or cached traces. These external approaches suffer from a speculator-agent gap: the speculator may predict a different call than the deployed agent, leading to wasted work, and they introduce additional model-serving overhead. The authors address this by training the agent itself to speculate on its own next tool call. They introduce a self-speculating agent that operates in dual modes, task solving and call prediction, and propose joint agent-speculator reinforcement learning to improve speculation accuracy without degrading task performance.
Method
The authors propose training a single tool-calling LLM agent to serve simultaneously as the task-solving policy and as its own next-call speculator. The core insight is that the same model can generate full reasoning-and-action trajectories in agent mode, while also predicting upcoming tool calls from partial prefixes in speculator mode, simply by appending a fixed speculation suffix to the input context. This design keeps speculation fully on-policy with the deployed agent, since the targets for speculation are drawn from the agent’s own current rollouts.
The task reward Rag(τ) is evaluated on the final answer y against the ground-truth answer y∗. For speculation, each tool-call turn t provides a prefix ht containing the query and all preceding reasoning, calls, and observations. The speculator input is formed by concatenating ht with a fixed suffix ssp, and the policy is asked to emit a predicted call a^t. The target is the agent’s actual next call at from the same trajectory. The speculation reward Rsp is defined as the product of a tool-name match indicator Rname and a name-conditioned argument reward Rargs. The argument score is computed via macro token-F1 over the keys of the target argument dictionary, giving equal weight to each key and avoiding dominance by long argument values. This construction mirrors the exact-match reuse requirement for speculative execution: a wrong tool name cannot reuse a pre-executed result, while a correct name can still earn partial credit for close argument values.
The joint agent–speculator RL procedure couples data collection but separates policy updates. At each RL iteration, a batch of queries is sampled and the current policy generates Gag agent trajectories. These rollouts provide both standard agent advantages (computed via DAPO-style normalization of Rag within each query group) and on-policy speculation supervision derived from the same trajectories. A fixed schedule m(k) selects whether the optimizer step uses the agent objective or the speculator objective. In the reported configuration, the schedule follows a repeated 4:8 pattern: four consecutive agent updates followed by eight consecutive speculator updates. Both modes share the same DAPO policy optimization but differ in input context, rollout structure, and reward function.
To stabilize training of the shared policy, three components are employed. First, a short SFT warmup is performed on a mixture of successful agent trajectories and corresponding intermediate-trajectory speculation examples, giving the model well-formed behavior in both modes before RL begins. Second, optimizer state is reset whenever the schedule switches between agent and speculator updates, preventing momentum and adaptive statistics from one mode’s objective from interfering with the next. Third, the agent reward includes a penalty for repeated meaningless tool calls, discouraging the failure mode where the agent continues calling tools without making task progress.
Experiment
The evaluation measures whether an off-the-shelf draft model can match a deployed agent’s next tool call under the same online context, using Hit@1 exact match as the primary metric. External speculators add substantial serving overhead and consistently underperform the target agent’s own self-speculation, confirming a speculator–agent gap. Joint agent-speculator RL improves self-speculation accuracy while preserving end-task performance, and ablation shows that SFT warmup, optimizer reset, and alternating update schedules are all critical for stable training. Cross-domain evaluation reveals that call-level prediction transfers across tasks, but matched-domain RL is needed to maintain task success.
Self-speculation by the target agent consistently achieves higher tool-call prediction accuracy than external smaller draft models from the same family, while also requiring less GPU memory and lower latency. External speculators add substantial serving overhead in memory and time, and their lower Hit@1 rates make the extra cost hard to justify. These results motivate training the agent itself as the speculator rather than deploying a separate model. Self-speculation with the target 4B model yields the highest Hit@1 across both model families and domains, outperforming all smaller external speculators. Adding an external speculator increases GPU memory and serving time substantially, even for the smallest draft models, compared to self-speculation. Within the same model family, larger external draft models improve Hit@1 over smaller ones but still fall short of the target agent predicting its own next call. The agent is the strongest off-the-shelf predictor of its own next tool call, indicating that family similarity alone does not recover the agent's call distribution.
Joint agent-speculator RL substantially improves next-call prediction accuracy over smaller external draft models and an SFT warmup baseline, while keeping task success stable. The base 4B agent already outperforms the best smaller speculator, and RL further lifts Hit@1 with consistent gains across all benchmarks, showing no trade-off between speculation quality and end-task performance. The base Qwen3-4B agent reaches 29.1 average Hit@1, far above the strongest external speculator at 18.6. After RL, average Hit@1 rises to 61.2 while task success changes only from 26.6 to 27.7, preserving agent capability.
When a self-speculator trained on one interaction family is evaluated on the other, next-call prediction accuracy (Hit@1) consistently improves, while end-task success declines. For example, a SearchQA-trained model raises Airline and Retail Hit@1 but lowers task success, and a ToolScale-trained model similarly boosts HotpotQA and MuSiQue Hit@1 at the cost of success. This suggests that call-level behaviors transfer across domains, but domain-specific reasoning remains necessary for task completion. Cross-domain transfer improves next-call Hit@1 across all tested benchmarks, regardless of whether the model was trained on SearchQA or ToolScale. Task success drops after transfer: the SearchQA-trained model loses success on Airline and Retail, while the ToolScale-trained model loses success on HotpotQA and MuSiQue. The self-speculator learns transferable tool-call argument matching, but successful task completion still depends on domain-specific interaction strategies.
Increasing the number of consecutive speculator updates per agent update substantially improves both next-call prediction accuracy and task success under a fixed iteration budget. The 4:8 schedule yields the strongest results, indicating that the speculator benefits from extended training blocks while periodic agent updates remain necessary to maintain task performance. The 1:1 schedule performs worst, achieving only 31.8 average Hit@1 and 10.3 average task success across the three search datasets. Lengthening the speculator block relative to the agent block consistently raises both metrics, with the 4:8 schedule reaching 55.2 average Hit@1 and 26.1 average task success. On BCP, task success remains low across all schedules, peaking at 8.1 under the 4:8 schedule, while HotpotQA and MuSiQue show larger gains in both Hit@1 and success.
These experiments evaluate self-speculation, where the agent predicts its own next tool call, against external smaller draft models. Self-speculation consistently achieves higher accuracy with lower overhead, and joint agent-speculator reinforcement learning further improves next-call prediction while keeping task success stable. Cross-domain transfer shows that call-level behaviors transfer across domains, but task completion still requires domain-specific reasoning. An imbalanced update schedule with more consecutive speculator updates per agent update yields the best trade-off between prediction accuracy and task performance.