ToolPRM: Fine-Grained Inference Scaling of Structured Outputs for Function Calling
ToolPRM uses fine-grained step scoring to improve structured function calling inference, outperforming coarse reward models.
Key Findings
Methodology
This paper introduces ToolPRM, a fine-grained process reward model that decomposes function calls into multiple decision steps—selecting function names, parameters, and values—and trains a reward predictor for each. Using function masking, rollout collection, and step-level annotations, a high-quality dataset is built. The approach models the call process as a sequence of state transitions, each supervised with binary correctness labels. Guided by ToolPRM, a beam search strategy explores multiple trajectories, applying the principle of 'explore more but retain less' to avoid early irreversible errors. Experimental results show significant improvements over outcome and coarse reward models across multiple benchmarks, especially on smaller models, demonstrating robustness and efficiency.
Key Results
- ToolPRM achieves a step accuracy of 99.11% and trajectory accuracy of 99.38% on the training dataset, outperforming ORM and C-PRM. On BFCL and ToolAlpaca benchmarks, it significantly surpasses other inference strategies, especially with smaller models like Hammer2.1-1.5B, where performance approaches larger models. The approach effectively mitigates early JSON errors by pruning incorrect partial trajectories early, leading to higher overall success rates.
- In ablation studies, fine-grained reward modeling consistently yields lower loss (0.0286) and higher accuracy metrics, validating its superiority over coarse-grained methods. The experiments confirm that expanding beam width while aggressively pruning based on step supervision enhances inference quality, especially in resource-constrained scenarios.
- The results demonstrate that integrating ToolPRM with smaller models can match or exceed the performance of larger models, offering practical benefits for edge deployment and real-time applications. The approach also generalizes well across different datasets and benchmarks, confirming its robustness.
Significance
This work addresses a fundamental challenge in structured output inference: early irreversible errors severely limit performance. By introducing a fine-grained reward mechanism and a principled search strategy, it advances the state-of-the-art in function calling tasks. The methodology enhances both accuracy and robustness, making large language models more reliable in real-world applications like API automation, virtual assistants, and data integration. The principle of 'explore more but retain less' offers a new paradigm for structured inference, emphasizing early error avoidance over broad exploration. This has broad implications for designing future AI systems that require precise multi-step reasoning, especially in safety-critical domains. The approach also opens avenues for further research into granular supervision and dynamic search strategies, promising to push the boundaries of what large models can achieve in complex, structured tasks.
Technical Contribution
The core technical contribution lies in developing ToolPRM, a fine-grained process reward model that decomposes function calls into multiple decision points and trains a binary classifier for each step. This contrasts with traditional outcome reward models that evaluate only the final result. The model leverages a state transition framework, defining explicit states for function name selection, parameter identification, and value filling, with supervision at each step. The training employs annotated trajectories with step-level correctness labels, enabling the reward model to guide inference via step-wise scoring. The integration with beam search, guided by the 'explore more but retain less' principle, allows for efficient pruning of invalid trajectories, significantly improving inference accuracy, especially for smaller models, and reducing error propagation.
Novelty
This study introduces the first fine-grained intra-call reward modeling framework tailored for structured function calling, breaking away from prior coarse-grained reward approaches. Its novelty includes the explicit decomposition of function calls into interpretable steps, the use of step-level supervision, and the state transition-based modeling of the call process. The 'explore more but retain less' inference principle is a new strategy specifically designed for structured outputs, addressing the unrecoverability of early errors. These innovations collectively enable more precise and robust inference, setting a new standard for structured reasoning in large language models.
Limitations
- The approach relies heavily on high-quality, step-level annotated datasets, which are costly and time-consuming to produce, limiting scalability in some domains.
- Early JSON errors, once occurred, remain difficult to recover from, potentially causing significant performance drops in highly error-prone environments.
- Computational costs increase with larger beam widths and detailed supervision, which may hinder real-time deployment on resource-constrained devices.
Future Work
Future directions include developing semi-supervised or unsupervised methods to reduce annotation costs, exploring adaptive beam strategies to balance exploration and pruning dynamically, and extending the framework to other structured tasks like multi-modal reasoning or multi-turn dialogues. Enhancing robustness against noisy data and integrating with continual learning setups are also promising avenues.
AI Executive Summary
Large language models (LLMs) have demonstrated remarkable capabilities in natural language understanding and generation, especially in function calling tasks that enable interaction with external systems. However, existing inference strategies largely treat function calls as monolithic units, which limits their ability to handle errors early in the process. This paper introduces ToolPRM, a novel fine-grained process reward model that decomposes each function call into multiple decision steps—such as selecting function names, parameters, and values—and trains a reward predictor for each step. Using function masking, rollout collection, and step-level annotations, the authors build a high-quality dataset to supervise the reward model. The core idea is to model the function call process as a sequence of state transitions, each supervised with binary correctness labels, enabling the model to evaluate and guide each decision accurately.
The inference strategy guided by ToolPRM adopts the principle of 'explore more but retain less,' expanding the search space with larger beam widths while aggressively pruning invalid trajectories based on step-wise scores. This approach effectively prevents the accumulation of early errors, which are often unrecoverable in structured outputs like JSON-formatted function calls. Extensive experiments on multiple benchmarks, including BFCL and ToolAlpaca, demonstrate that ToolPRM consistently outperforms outcome and coarse reward models, especially on smaller models such as Hammer2.1-1.5B, achieving performance comparable to larger models.
The significance of this work lies in its ability to enhance the robustness and accuracy of structured output inference, addressing a long-standing challenge in multi-step reasoning tasks. By focusing on fine-grained supervision and strategic pruning, the method offers a scalable solution that is particularly suitable for edge deployment and real-time applications. The research opens new avenues for granular supervision in complex reasoning tasks, promising to improve the reliability of AI systems in practical scenarios. Future work will explore reducing annotation costs, adaptive search strategies, and extending the framework to other structured domains, further advancing the frontier of AI reasoning capabilities.
Deep Dive
Abstract
Large language models (LLMs) excel at function calling, but inference scaling has been explored mainly for unstructured generation. We propose an inference-scaling framework for structured outputs that combines fine-grained beam search with \textbf{ToolPRM}, a process reward model scoring each intra-call decision (function name and argument filling). We build the first fine-grained intra-call supervision dataset via function masking, rollout collection, and step-level annotation. ToolPRM outperforms outcome and coarse-grained reward models in predictive accuracy and yields consistent test-time gains on multiple function-calling benchmarks. We further show that structured generation follows ``\textbf{explore more but retain less}'', since early JSON errors are unrecoverable.