Technical Summary: Speculate While You Reason

Problem Statement

Large Language Model (LLM) agents often incur substantial wall-clock latency while waiting for external tool call results (e.g., web search, database queries, API calls). A promising mitigation strategy is tool-call speculation, where a system predicts and pre-executes an agent's next tool call in parallel with the agent's reasoning. If the prediction matches the agent's eventual call exactly, the pre-executed result is reused, effectively hiding the latency.

However, existing speculation methods suffer from a speculator–agent gap. Current approaches rely on:

1. External draft models: Smaller LLMs that approximate a generic tool-using assistant but fail to align with the specific behavior of the deployed agent.
2. Cached traces: Historical tool call logs that may correspond to different user queries, reasoning states, or outdated environments.

These methods introduce system overhead (additional model weights, separate KV caches) and often fail to predict the specific next structured call the target agent will make, limiting the utility of speculation.

Methodology

The paper proposes a Self-Speculating Agent, a single model trained to operate in two distinct modes without requiring separate parameters or KV caches:

1. Agent Mode: The model solves tasks by interleaving reasoning, tool calls, observations, and final answers.
2. Speculator Mode: Given an intermediate trajectory and a fixed "speculation suffix," the model directly predicts the next structured tool call (tool name and arguments).

Joint Agent–Speculator Reinforcement Learning (RL)

To train a single policy ($\pi_\theta$) for both objectives without degrading task performance, the authors propose a Joint Agent–Speculator RL framework:

• On-Policy Data Generation: Speculation targets are derived directly from the agent's own current rollouts. For every tool-call turn $t$ in a sampled trajectory, the prefix $h_t$ is used to construct a speculator input, and the agent's actual next call $a_t$ serves as the ground truth target.
• Reward Formulation:
  • Agent Reward ($R_{ag}$): Based on the final task success (e.g., answer correctness).
  • Speculation Reward ($R_{sp}$): A name-conditioned argument reward. It requires an exact match of the tool name ($R_{name}$) and calculates a macro token-F1 score for the argument dictionary ($R_{args}$). The final reward is the product $R_{sp} = R_{name} \cdot R_{args}$.
• Alternating Optimization: The training alternates between agent updates and speculator updates (e.g., a 4:8 schedule). This prevents the two modes from interfering with each other's optimization landscapes.
• Stabilization Techniques:
  1. SFT Warmup: A preliminary supervised fine-tuning phase on a mix of successful trajectories and speculation examples to establish well-formed behavior in both modes.
  2. Optimizer Resets: The optimizer state is reset whenever the training schedule switches between agent and speculator modes to prevent momentum carryover from one objective to the other.
  3. Penalty for Repetition: The agent reward includes a penalty for repeated, meaningless tool calls to prevent collapse into unproductive loops.

Key Contributions

1. Identification of the Speculator–Agent Gap: The authors demonstrate that off-the-shelf smaller draft models often fail to match the next-call distribution of the target agent, while the target agent itself is a stronger predictor of its own behavior.
2. Self-Speculating Agent Architecture: A unified model design that switches between task-solving and next-call prediction modes, fully reusing the prefix KV cache and eliminating the need for auxiliary models or historical trace caches.
3. Joint Agent–Speculator RL: A training methodology that derives speculation targets from the agent's own rollouts and alternates updates to preserve task-solving performance while improving speculation accuracy.

Experimental Results

The method was evaluated on Agentic SearchQA (HotpotQA, MuSiQue, BrowseComp-Plus) and Conversational Tool-Use (τ-bench Airline and Retail) tasks using Qwen3-4B and Qwen3.5-4B backbones.

• Speculation Performance: The joint RL method significantly improved the Hit@1 score (exact match of tool name and arguments).
  • Qwen3-4B: Improved from 44.1 (SFT baseline) to 61.2.
  • Qwen3.5-4B: Improved from 48.9 (SFT baseline) to 66.3.
  • These results outperformed external speculators (smaller same-family models) and the base agent prompted to speculate its own calls.
• Task Performance: Crucially, the downstream task success rates remained stable or slightly improved (e.g., Qwen3-4B task success moved from 26.6 to 27.7), confirming that the dual-mode training did not degrade the agent's primary task-solving capabilities.
• Efficiency: Self-speculation reduced serving overhead by avoiding the need for a second model's weights and a separate KV cache, while maintaining the prefix cache reuse.

Significance and Claims

The paper claims that tool-call speculation can be folded directly into the agent itself. By unifying the agent and speculator within a single model and training them jointly via on-policy RL, the system achieves high-accuracy next-call prediction without the system overhead or alignment issues associated with external draft models.

The authors position this work as an initial step toward speculative tool use, noting that the method is particularly well-suited for read-only tools (search, retrieval, database lookups) where incorrect pre-executions can be safely discarded. They acknowledge limitations regarding tools that mutate external state (e.g., placing orders), which would require additional safeguards like dry-run modes or transactional rollback, and note that evaluation is currently limited to 4B-scale models and specific task families.