Technical Summary: Process Rewards with Learned Reliability (BETAPRM)

1. Problem Statement

Process Reward Models (PRMs) provide step-level feedback for reasoning tasks, guiding candidate selection and policy optimization. However, existing PRMs typically output a single point estimate (e.g., a scalar probability) for each reasoning step. This interface presents two critical limitations:

1. Lack of Uncertainty Quantification: A single scalar score cannot distinguish between a confident prediction and an uncertain one. Downstream methods are forced to treat imperfect step-level predictions as reliable decision signals, even when the model's confidence is low.
2. Noisy Finite-Sample Supervision: Standard PRM training often relies on Monte Carlo continuations to generate supervision. If $N$ continuations are sampled from a prefix and $K$ succeed, the empirical ratio $\hat{q} = K/N$ is treated as a point target for regression. However, $\hat{q}$ is merely a noisy estimate of the true prefix success probability. Repeated sampling from the same prefix could yield different $K$ values due to randomness, yet standard training forces the model to fit this specific noisy ratio rather than the underlying distribution of possible outcomes.

Consequently, downstream applications cannot distinguish between high-reward steps that are reliably correct and those that are high-reward but uncertain, leading to suboptimal test-time scaling and computation allocation.

2. Methodology: BETAPRM

The authors propose BETAPRM, a distributional Process Reward Model that predicts both the step-level success probability and the reliability of that prediction.

2.1 Beta-Binomial Count Model

Instead of regressing to a point target $\hat{q} = K/N$, BETAPRM models the prefix success probability $q_t$ as a latent variable following a Beta distribution:
$$q_t \sim \text{Beta}(\alpha_t, \beta_t)$$
The observed number of successful continuations $K_t$ out of $N$ trials is modeled as a Binomial draw conditioned on $q_t$:
$$K_t | q_t \sim \text{Binomial}(N, q_t)$$
By marginalizing out the latent $q_t$, the model learns a Beta-Binomial distribution over the observed counts. This allows the model to assign likelihood to the observed count $K_t$ rather than forcing a fit to the ratio $K_t/N$.

2.2 Parameterization

The Beta distribution is reparameterized by its mean $\mu_t$ and concentration $\kappa_t$:

• Mean ($\mu_t$): Represents the expected success probability, serving as the standard scalar PRM score.
• Concentration ($\kappa_t$): Controls the sharpness of the belief. A high $\kappa_t$ indicates a concentrated (reliable) belief, while a low $\kappa_t$ indicates a flattened (uncertain) belief that can explain a wider range of Monte Carlo observations.

The model architecture uses a standard language model backbone. The mean $\mu_t$ is derived from the softmax probability of a "Yes" reward token. The concentration $\kappa_t$ is predicted by a separate lightweight linear head applied to the hidden state at the process marker, ensuring the reward and reliability channels are distinct.

2.3 Training Objective

The training loss consists of two components:

1. Beta-Binomial Likelihood ($L_{\text{Beta-Binomial}}$): Maximizes the likelihood of the observed count $K_t$ given the predicted Beta parameters. This encourages the model to explain the observed Monte Carlo data distribution.
2. Auxiliary Regularization ($L_{\text{reg}}$): To prevent the concentration parameter from becoming uncalibrated, an auxiliary loss penalizes high concentration ($\kappa_t$) when the predicted mean $\mu_t$ disagrees with the observed ratio $K_t/N$. Crucially, a stop-gradient operation is applied to $\mu_t$ in this term to ensure the regularizer calibrates $\kappa_t$ without pulling the reward mean toward the noisy ratio.

$$L = L_{\text{Beta-Binomial}} + \lambda_{\text{reg}} L_{\text{reg}}$$

3. Application: Adaptive Computation Allocation (ACA)

Leveraging the learned reliability signal, the authors introduce Adaptive Computation Allocation (ACA) for PRM-guided Best-of-N reasoning. Unlike fixed-budget Best-of-N, which samples a constant number of candidates regardless of difficulty, ACA dynamically allocates compute:

• Risk-Adjusted Scoring: Candidates are scored using a risk-adjusted metric $S(y) = \frac{1}{T}\sum (\mu_t - \lambda \sigma_t)$, where $\sigma_t$ is the standard deviation derived from the Beta distribution. This penalizes candidates with high uncertainty.
• Early Stopping: ACA samples candidates in progressive batches. It stops early if the lower confidence bound (LCB) of the current best candidate exceeds the upper confidence bound (UCB) of all competitors, indicating the winner is reliably ahead.
• Uncertainty-Guided Repair: If the stopping criterion is not met, ACA allocates the next batch of compute to repair the most promising uncertain prefix (the non-winner with the highest UCB), rather than generating entirely new candidates from scratch.

4. Experimental Results

The authors evaluated BETAPRM across four backbone models (InternVL2.5-8B, InternVL3-8B, InternVL3-14B, Qwen2.5-VL-7B) and four reasoning benchmarks (MathVision, OlympiadBench, MathVerse, MathVista).

• Best-of-N Selection: BETAPRM improved final-answer accuracy over standard PRMs across all backbones. For example, on InternVL2.5-8B, BETAPRM achieved an average gain of +3.37 points over standard PRM.
• Error Detection: BETAPRM preserved standard step-level error detection capabilities, achieving comparable or slightly better micro-F1 scores on VisualProcessBench compared to standard PRMs.
• Accuracy-Token Tradeoff: In the ACA setting (fixed max budget $N=16$), BETAPRM significantly improved the efficiency of inference.
  • It reduced token usage by up to 33.57% (on MathVerse with InternVL2.5-8B) while simultaneously improving final-answer accuracy compared to fixed-budget Best-of-16.
  • Ablation studies confirmed that the learned uncertainty signal from BETAPRM was superior to using a standard PRM with proxy uncertainty or reward-only allocation, as the latter failed to identify prefixes where additional computation was necessary despite high scores.

5. Significance and Claims

The paper claims that BETAPRM addresses a fundamental mismatch in current PRM interfaces: the inability to distinguish between reliable and uncertain predictions. By modeling the prefix success probability as a distribution rather than a point estimate, BETAPRM:

1. Enables Trustworthy Decision Making: Downstream algorithms can now distinguish confident rewards from uncertain ones, preventing unreliable scores from dictating decisions.
2. Improves Test-Time Efficiency: The learned reliability signal allows for adaptive computation allocation, reducing inference costs (tokens) without sacrificing, and often improving, accuracy.
3. Better Utilization of Supervision: The Beta-Binomial objective effectively utilizes the finite-sample nature of Monte Carlo supervision, treating the count $K$ as a stochastic observation rather than a deterministic label.

The authors conclude that turning scalar process rewards into reliability-aware signals is a crucial step for scaling test-time compute and improving the robustness of reasoning systems.