Bayesian belief approach improves reliability of language model rewards
Process Reward Modeling with Reliability- and Conservatism-Aware Bayesian Belief
Machine Learning
Summary
Large language models rely on reward scores to know how well they are reasoning, but these scores can be wrong for two main reasons: limited training data and differences between training and testing conditions. The authors propose a new method, BayesianPRM, which treats these reward scores like uncertain guesses and updates its beliefs carefully to be more reliable and cautious. Their approach uses a set of reward guesses and adjusts belief weights to avoid overestimating good outcomes. Tests show this method provides more accurate reward estimates and better performance when scaling language models.
What this means in practice
- •For machine learning engineers: Improve accuracy of reward estimates during language model deployment to enhance model reasoning performance.
- •For ai system operators: Better calibrate reward signals to prevent overestimation when applying AI models in real-world tasks involving multimodal reasoning.
Authors
Zhongjian Qiao, Zhixiang Zhou, Jiafei Lyu, Peisong Wang, Qiaosheng Zhang, Siyang Gao, Shuang Qiu
Abstract
Process Reward Models (PRMs) play a central role in test-time scaling for large language model (LLM) reasoning. However, PRMs may provide inaccurate process rewards for test-time scaling, which we identify as arising from two sources of uncertainty: (1) \textit{PRM uncertainty} that arises from insufficient supervision data and imperfect PRM training; and (2) \textit{policy uncertainty} that results from the mismatch between the policy to collect supervision data and the policy for test-time scaling, leading to systematic reward overestimation. Previous studies fail to jointly consider both uncertainties. To bridge this gap, we propose BayesianPRM, which treats the process reward as a random variable within a reward hypothesis space and maintains a Bayesian belief over it. Since exact Bayesian inference over the continuous hypothesis space is generally intractable, we discretize the reward hypothesis space with a reward ensemble and introduce two complementary belief updates, each targeting one source of uncertainty: (1) to address PRM uncertainty, we introduce \textit{reliability-aware belief updating} that measures how well each ensemble member explains the supervision data and updates the belief with an amortized belief network optimized by the evidence lower bound (ELBO); and (2) to tackle policy uncertainty, we propose \textit{conservatism-aware belief calibration} which further reweights the belief toward lower rewards to suppress overestimation. These two updates yield a reliability- and conservatism-aware Bayesian belief, whose expectation determines the process reward. Experiments on multimodal reasoning benchmarks indicate that BayesianPRM consistently yields more accurate process reward estimates while enhancing test-time scaling performance.