Table of Contents
Fetching ...

Optimal Aggregation of LLM and PRM Signals for Efficient Test-Time Scaling

Peng Kuang, Yanli Wang, Xiaoyu Han, Yaowenqi Liu, Kaidi Xu, Haohan Wang

TL;DR

This work reframes test-time scaling as a MAP problem and shows that the optimal strategy to fuse LLM outputs with PRM verification is a weighted majority vote that combines a PRM-derived signal with an LLM reliability term $q_M$. Crucially, the optimal weights are model-pair dependent and can be negative, meaning low-quality PRM guidance actively penalizes corresponding answers. The authors propose practical calibration methods—non-parametric KDE and parametric forms (Logit WV and Linear WV)—to approximate these weights using a one-time calibration set, enabling efficient test-time scaling. Across 35 LLM-PRM pairs on the MATH benchmarks, calibrated weighting consistently outperforms baselines and reduces computation by up to around 37% (MATH) or 21% (MATH500) while achieving higher accuracy, illustrating that smarter aggregation can outperform mere scaling.

Abstract

Process reward models (PRMs) are a cornerstone of test-time scaling (TTS), designed to verify and select the best responses from large language models (LLMs). However, this promise is challenged by recent benchmarks where simple majority voting, which ignores PRM signals, occasionally outperforms standard PRM-based selection. This raises a critical question: How can we effectively utilize verification signals from PRMs for TTS? To address this, we start by developing a theoretical framework for optimally combining signals from both the LLM and the PRM. Our framework reveals that the optimal strategy is a weighted aggregation of responses, a strategy whose effectiveness hinges on estimating weights that capture the complex interplay between the models. Based on our theoretical results, we empirically show that these optimal weighting functions differ significantly across LLM-PRM pairs and, notably, often assign substantial negative weights. Motivated by these insights, we propose efficient pre-computation methods to calibrate these weighting functions. Extensive experiments across 5 LLMs and 7 PRMs demonstrate that our calibration method significantly boosts the TTS efficiency, surpassing the performance of vanilla weighted majority voting while using only $21.3\%$ of the computation. Ultimately, our work demonstrates that investing in a more intelligent aggregation strategy can be a more convincing path to performance gains than simply scaling test-time computation.

Optimal Aggregation of LLM and PRM Signals for Efficient Test-Time Scaling

TL;DR

This work reframes test-time scaling as a MAP problem and shows that the optimal strategy to fuse LLM outputs with PRM verification is a weighted majority vote that combines a PRM-derived signal with an LLM reliability term . Crucially, the optimal weights are model-pair dependent and can be negative, meaning low-quality PRM guidance actively penalizes corresponding answers. The authors propose practical calibration methods—non-parametric KDE and parametric forms (Logit WV and Linear WV)—to approximate these weights using a one-time calibration set, enabling efficient test-time scaling. Across 35 LLM-PRM pairs on the MATH benchmarks, calibrated weighting consistently outperforms baselines and reduces computation by up to around 37% (MATH) or 21% (MATH500) while achieving higher accuracy, illustrating that smarter aggregation can outperform mere scaling.

Abstract

Process reward models (PRMs) are a cornerstone of test-time scaling (TTS), designed to verify and select the best responses from large language models (LLMs). However, this promise is challenged by recent benchmarks where simple majority voting, which ignores PRM signals, occasionally outperforms standard PRM-based selection. This raises a critical question: How can we effectively utilize verification signals from PRMs for TTS? To address this, we start by developing a theoretical framework for optimally combining signals from both the LLM and the PRM. Our framework reveals that the optimal strategy is a weighted aggregation of responses, a strategy whose effectiveness hinges on estimating weights that capture the complex interplay between the models. Based on our theoretical results, we empirically show that these optimal weighting functions differ significantly across LLM-PRM pairs and, notably, often assign substantial negative weights. Motivated by these insights, we propose efficient pre-computation methods to calibrate these weighting functions. Extensive experiments across 5 LLMs and 7 PRMs demonstrate that our calibration method significantly boosts the TTS efficiency, surpassing the performance of vanilla weighted majority voting while using only of the computation. Ultimately, our work demonstrates that investing in a more intelligent aggregation strategy can be a more convincing path to performance gains than simply scaling test-time computation.
Paper Structure (18 sections, 1 theorem, 15 equations, 4 figures, 2 tables)

This paper contains 18 sections, 1 theorem, 15 equations, 4 figures, 2 tables.

Key Result

Theorem 3.2

Under the assumptions above, maximizing the log-likelihood is equivalent to maximizing the score:

Figures (4)

  • Figure 1: The PRM score distributions and optimal weighing functions on 6 combinations of LLM-PRM pairs. Left y-axis: the probability density of the PRM scores. Right y-axis: the optimal weights $w^*(p)$ learned via KDE for different LLM-PRM pairs. Note their model-dependent nature and the presence of negative weights for low PRM scores.
  • Figure 2: The performance of various scaling methods averaged across all LLM and PRM combinations. The computation efficiency improvement of the Logit WV compared to the best-performing baseline, Vanilla WV, is marked in red. Left: On the MATH dataset. Right: On the MATH500 dataset.
  • Figure 3: Left and Middle: The grid search result of the offset parameter $b$ for both Logit WV and Linear WV, where the optimal value is marked with vertical lines. The consistently positive optimal value across LLMs demonstrates the necessity of negative weights. Right: The performance of the calibration methods when we scale the calibration set size. The performance can be further improved with larger calibration sets.
  • Figure 4: Left:The comparison of the dataset-wise estimated PRM score weighting function and the per-question estimated optimal PRM score weighting function. A large variance among questions can be seen. Right: The mean absolute error of estimated $\hat{q}_M$ compared to the true $q_M$.

Theorems & Definitions (2)

  • Theorem 3.2: Optimal Aggregation Score
  • proof