ToolPRM: Fine-Grained Inference Scaling of Structured Outputs for Function Calling

AI summary

Overview

Research area: Large language model inference scaling and process reward modeling, applied specifically to structured output generation for function calling / tool use.

Technical level: Intermediate. The paper assumes familiarity with large language models, beam search, reward models, and the distinction between outcome-level and process-level supervision, but its core ideas are explained concretely enough for readers with a general machine learning background.

Scope: The paper introduces ToolPRM, a fine-grained process reward model that scores each intra-call decision in a function call (function name selection, parameter name selection, and argument value filling), and pairs it with a beam search procedure that follows the principle "explore more but retain less."

Authors are affiliated with Shanghai Jiao Tong University, Longshine AI Research, and Shanghai Innovation Institute. The paper is arXiv:2510.14703v2 [cs.AI].

What This Paper Is About

Inference scaling — spending extra computation at test time to search over multiple candidate outputs — has been extensively studied for unstructured generation such as math reasoning, but is underexplored for structured outputs like function calls. Existing work on function calling typically scores an entire generated call as one monolithic unit, ignoring that a call is really a sequence of smaller decisions (which function, which parameters, which values). The paper's goal is to build a process reward model that supervises those individual decisions, and to show that this fine-grained supervision both predicts correctness better and produces better final outputs when used to guide search.

Key Contributions

1. A fine-grained intra-call process supervision dataset. The authors built what they describe as the first annotated dataset for fine-grained, intra-call reward modeling of structured function calling, including annotation scripts. It is derived from xlam-function-calling-60k and xlam-irrelevance-7.5k, with an average of 6.25 step labels per sample across a total of 192,061 samples before masking augmentation.

2. ToolPRM, a fine-grained process reward model. ToolPRM decomposes each function call into five states and five hierarchical step-label types, and is trained generatively to predict a binary correctness label ("+" or "-") for each state-action pair. It is intended to work with different backbone policy models.

3. An inference scaling principle for structured outputs: "explore more but retain less." The paper argues that because errors in structured output (such as malformed JSON) are effectively unrecoverable, the right strategy is to widen beam search exploration while aggressively pruning the number of retained candidate beams — the opposite of what works for unstructured reasoning.

4. Empirical validation across multiple benchmarks and policy model sizes. ToolPRM is evaluated as a reward model and as a search guide on BFCL, ToolAlpaca, and (in the appendix) API-Bank.

Main Findings

• Finer-grained reward models predict correctness better. On the authors' held-out data, ToolPRM achieved loss 0.0286, step-level accuracy 99.11%, and trajectory-level accuracy 99.38%, compared with C-PRM (0.0371, 98.87%, 99.06%) and ORM (0.0536, 98.39%, 98.39%).

• ToolPRM-guided beam search gives the most consistent test-time gains. With Hammer2.1 as the policy model, ToolPRM improved the unweighted BFCL average over the base model in all three tested sizes: 1.5B (82.79 to 85.61), 3B (86.86 to 88.88), and 7B (88.65 to 89.52).

• Other inference scaling strategies were unstable. Token-level beam search, majority voting, and Best-of-N (scored by an ORM) sometimes performed worse than the base policy model. For example, token-level beam search reduced Hammer2.1-3B's BFCL average from 86.86 to 82.50. The authors attribute this to non-greedy sampling introducing small errors that invalidate an entire structured trajectory.

• Gains are larger for smaller policy models. The paper reports that Hammer2.1-1.5B with ToolPRM reaches performance comparable to the baseline 3B model, Hammer2.1-3B with ToolPRM reaches a level on par with the baseline 7B model, and Hammer2.1-7B with ToolPRM outperforms several larger models including Qwen2.5-32B-Instruct.

• Increasing beam width helps more than increasing the number of retained beams. In the ToolAlpaca budget analysis with N=4 and M=4 variants, scaling the beam width produced consistent improvement, while scaling the number of retained candidates yielded less gain and in some instances caused noticeable accuracy degradation.

• Self-reflection is a weak substitute for intra-call pruning. Qwen2.5-7B-Instruct with self-reflection reached 86.92 BFCL average (versus 86.46 without) and Qwen2.5-32B-Instruct with reflection reached 88.04 (versus 87.21), both below Hammer2.1-7B with ToolPRM at 89.52.

• Function masking helps. Training ToolPRM without function masking gave an 84.90 BFCL average with Hammer2.1-1.5B, versus 85.61 with masking. Masking is applied to only a subset of training samples and is described as both an anti-memorization and regularization mechanism.

• ToolPRM outperformed a constrained decoding baseline. Against FANTASE with a RoBERTa reranker trained on the same dataset, ToolPRM scored higher on all four BFCL splits (85.61 versus 84.27 average on Hammer2.1-1.5B).

• Generalization beyond the main benchmarks. On the cleaned API-Bank test set (314 tool-use dialogues, 753 API calls), ToolPRM improved over the base model on both L1 (93.26 to 94.35 average) and L2 (76.46 to 78.11 average).

Methodology in Plain English

Step 1 — Break a function call into decisions. Instead of judging a whole call as good or bad, the authors split it into five states: the starting state with the query and available tools; choosing a function name; choosing a parameter name; filling in that parameter's value; and termination. Each transition between states is labeled.

Step 2 — Build training data. They took xlam-function-calling-60k and xlam-irrelevance-7.5k, and applied function masking — replacing function names and parameter identifiers with random strings so the model has to read descriptions rather than memorize names. Hammer2.1-3b and Hammer2.1-7b were used as policy models to generate rollouts. Each generated call (usually JSON) was then annotated with five hierarchical label types: <FUNC_NAME>, <ARG_VALUE>, <PARAM_FINISH>, <FUNC_FINISH>, and <TOTAL_FINISH>, each carrying a binary label determined by exact match against ground truths. The authors note that some labels appear redundant (e.g., <ARG_VALUE> and <PARAM_FINISH>) but that the hierarchical redundancy improved generalization in their experiments.

Step 3 — Train the reward model. Because the backbone is itself a language model, ToolPRM is trained generatively to output a "+" or "-" token for a given state-action pair, minimizing a standard negative log-likelihood loss. The reward model backbone is Hammer2.1-3b, trained with SFT for 5 epochs.

Step 4 — Search with the reward model. At each step, ToolPRM scores each candidate using s = e^(s+) / (e^(s+) + e^(s−)), where s+ and s− are logits on the "+" and "−" tokens. The top-N highest-scoring candidates are retained, and each retained candidate can expand into M subsequent steps. N is the number of beams and M is the beam width. The paper's principle is to raise M (explore more) while keeping N small (retain less).

Step 5 — Evaluate. Reward model quality is measured with loss, step-level accuracy, and trajectory-level accuracy. End-to-end function calling is measured on BFCL using AST-based accuracy across the Simple, Multiple, Parallel, and Multiple Parallel splits, and on ToolAlpaca using F1 for API selection and parameter value assignment.

Setup details: Experiments ran on NVIDIA 8xH100 GPU clusters. Training used the Adam optimizer, batch size 1024, learning rate 1e-3 with a warmup ratio of 0.008 followed by linear decay, and weight decay 1e-5. Beam search used temperature 0.8, with N and M selected from {1, 2, 4, 8, 16}. The dataset's function masking expands its size 4 to 5 times, which the authors say surpasses typical Math PRM datasets such as OpenAI's prm800k.

Why This Matters

Impact on research. The paper makes a conceptual argument that inference scaling is not transferable unchanged between unstructured and structured generation. It identifies unrecoverability of early structured errors as the reason why retaining many candidate trajectories can hurt, and it supplies both a labeled supervision resource and a training recipe for intra-call process rewards. It also offers a direct empirical rebuttal to the idea that stronger reasoning models can self-correct structural errors post hoc.

Real-world applications.

• Agentic assistants that must emit valid JSON tool calls — where one malformed argument can break an entire downstream pipeline.
• On-device or edge deployment, since the paper specifically notes that ToolPRM's larger gains on small models (1.5B and 3B) suit edge environments where function calling is frequently deployed.
• Enterprise API orchestration, where a smaller locally hosted model plus ToolPRM could substitute for a larger hosted model at inference time.
• Tool-use evaluation and filtering pipelines, where a fine-grained reward model could score and rank candidate tool calls without executing them.

Industry relevance. The paper reports that a 7B policy model guided by ToolPRM can beat significantly larger models on BFCL, which maps directly onto cost and latency trade-offs in production function-calling systems. Function-calling benchmarks used here (BFCL, ToolAlpaca, API-Bank) are widely referenced in industry evaluation, and the released dataset and annotation scripts lower the barrier for building similar reward models.

Future Directions

1. Adaptive calibration of exploration versus retention. The conclusion states explicitly that the optimal trade-off in "explore more but retain less" is not yet dynamically adjustable, and proposes future work on calibrating exploration and retention based on input complexity or ToolPRM-derived confidence.

2. Beyond discretized, step-wise decision modeling. The limitations section notes that the framework assumes a discretized step-wise view of decision making, which may not capture implicit reasoning or latent uncertainty.

3. Guarantees on global optimality. The authors acknowledge that rewarding intermediate structure and consistency cannot guarantee the globally optimal final tool choice or argument specification in every case.

4. Sensitivity to design choices. The limitations also flag that the masking design and state definitions are additional modeling components whose choices may affect behavior and require careful implementation — an invitation to systematically study those choices.

Target Audience

This paper is most useful for researchers and engineers working on LLM agents, tool use, and function-calling systems; for practitioners interested in test-time compute scaling and process reward models; and for teams building on-device or cost-constrained function-calling deployments. Readers studying inference scaling for structured or constrained decoding, and those building annotated supervision datasets for reward modeling, will also find the dataset construction pipeline and the "explore more but retain less" argument directly relevant.

Authors’ 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.