Training an LLM to call tools reliably requires datasets that pair user queries with correct tool-use chains. Producing that data at scale has been slow and expensive. A team of researchers from Google, the University of Tokyo, RIKEN AIP, and Tohoku University introduce ToolGrad. The research work inverts the usual pipeline: build a verified tool chain first, then write the query. Gemma-3 models fine-tuned on 500 samples of the resulting data reach scores that sit alongside frontier proprietary models on the Berkeley Function Calling Leaderboard.
Is it deployable? Yes. The code is Apache-2.0, the ToolGrad-500 dataset and the 1B, 4B, and 12B models are on Hugging Face, and there is a PyPI package.
The problem with query-first generation
Prior pipelines such as ToolBench and ToolACE follow a query-first recipe. The system samples a pool of APIs, asks an LLM to invent a plausible user instruction, and then dispatches a depth-first search (DFS) agent to find a tool-use path that satisfies it. The search has no guarantee of success. When it dead-ends, the compute spent on exploration is wasted, and the sample is discarded. The paper frames this as distilling valuable trajectories from a complex and often failing agent exploration, which is inherently inefficient.
ToolGrad reverses the order. It first constructs a ground-truth tool-use chain by actually executing APIs, then annotates that chain with a matching user query. An explicit, working chain is far less ambiguous than a hypothetical prompt, so the chain-to-query step takes a single LLM call.
Four modules in a loop
Each iteration runs four modules in sequence:
- API Proposer narrows a sampled set of APIs down to a few candidates that could extend the current workflow.
- API Executors run those candidates in parallel and produce detailed execution reports.
- API Selector reviews the reports, picks the single best-performing call, and appends it to the workflow. Its directional feedback is the textual gradient.
- LLM Updater rewrites the synthetic user query and AI response so they match the new API set.
Repeating the loop yields one sample: a user query, a verified API workflow, and the final response. The repository’s default configuration runs 10 iterations over 50 sampled APIs per workflow.
The research team evaluated data generation on the ToolBench API database, which contains 16,000+ real-world APIs, and compared ToolGrad against ToolBench’s DFS-based query-first approach. According to the research paper:
- Pass rate rose from 63.8% (DFS) to 99.8% (ToolGrad).
- Ground-truth tool uses per sample rose from 2.1 to 3.4, meaning longer chains.
- Tool-use steps per sample fell from 34.3 to 20.0.
- LLM invocations per sample fell slightly, from 64.5 to 63.9.
The 0.2% failure case occurred when the agent could not get a successful response from 3 selected APIs across all 10 iterations and saved an empty sample.
Interactive explainer
Credit: Source link
