Reward Model Generalization for Compute-Aware Test-Time Reasoning

• URL: http://arxiv.org/abs/2505.18065v1
• Date: Fri, 23 May 2025 16:12:12 GMT
• Title: Reward Model Generalization for Compute-Aware Test-Time Reasoning
• Authors: Zeen Song, Wenwen Qiang, Siyu Zhao, Changwen Zheng, Gang Hua,
• Abstract summary: External test-time reasoning enhances large language models (LLMs) by decoupling generation and selection.<n>A central challenge in this setting is test-time compute optimality (TCO), i.e., how to maximize answer accuracy under a fixed inference budget.<n>We analyze how the generalization error of the PRM affects compute efficiency and reasoning performance.<n>Motivated by this analysis, we propose Compute-Aware Tree Search (CATS), an actor-critic framework that dynamically controls search behavior.
• Score: 21.05692631562457
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: External test-time reasoning enhances large language models (LLMs) by decoupling generation and selection. At inference time, the model generates multiple reasoning paths, and an auxiliary process reward model (PRM) is used to score and select the best one. A central challenge in this setting is test-time compute optimality (TCO), i.e., how to maximize answer accuracy under a fixed inference budget. In this work, we establish a theoretical framework to analyze how the generalization error of the PRM affects compute efficiency and reasoning performance. Leveraging PAC-Bayes theory, we derive generalization bounds and show that a lower generalization error of PRM leads to fewer samples required to find correct answers. Motivated by this analysis, we propose Compute-Aware Tree Search (CATS), an actor-critic framework that dynamically controls search behavior. The actor outputs sampling hyperparameters based on reward distributions and sparsity statistics, while the critic estimates their utility to guide budget allocation. Experiments on the MATH and AIME benchmarks with various LLMs and PRMs demonstrate that CATS consistently outperforms other external TTS methods, validating our theoretical predictions.

Related papers

• Accelerating LLM Reasoning via Early Rejection with Partial Reward Modeling [12.835376812101323]
  We introduce the hypothesis that PRMs are also Partial Reward Models.<n>This allows for principled early rejection based on intermediate token-level signals.<n>On math reasoning benchmarks, our method achieves up to 1.4$times$-9$times$ reduction in inference FLOPs without degrading final performance.
  arXiv  Detail & Related papers  (2025-08-04T00:58:56Z)
• From Mathematical Reasoning to Code: Generalization of Process Reward Models in Test-Time Scaling [32.72867198629561]
  We investigate the interplay between pre-training and reward model training FLOPs to assess their influence on PRM efficiency and accuracy.<n>Our findings indicate that PRMs trained on mathematical datasets exhibit performance comparable to those tailored for code generation.
  arXiv  Detail & Related papers  (2025-05-24T12:44:15Z)
• DGRO: Enhancing LLM Reasoning via Exploration-Exploitation Control and Reward Variance Management [18.953750405635393]
  Decoupled Group Reward Optimization (DGRO) is a general RL algorithm for Large Language Models (LLMs) reasoning.<n>We show that DGRO achieves state-of-the-art performance on the Logic dataset with an average accuracy of 96.9%, and demonstrates strong generalization across mathematical benchmarks.
  arXiv  Detail & Related papers  (2025-05-19T10:44:49Z)
• Supervised Optimism Correction: Be Confident When LLMs Are Sure [91.7459076316849]
  We establish a novel theoretical connection between supervised fine-tuning and offline reinforcement learning.<n>We show that the widely used beam search method suffers from unacceptable over-optimism.<n>We propose Supervised Optimism Correction, which introduces a simple yet effective auxiliary loss for token-level $Q$-value estimations.
  arXiv  Detail & Related papers  (2025-04-10T07:50:03Z)
• The First Few Tokens Are All You Need: An Efficient and Effective Unsupervised Prefix Fine-Tuning Method for Reasoning Models [69.798277882245]
  We introduce Unsupervised Prefix Fine-Tuning (UPFT) to enhance large language models' reasoning efficiency.<n>UPFT removes the need for labeled data or exhaustive sampling.<n> Experiments show that UPFT matches the performance of supervised methods.
  arXiv  Detail & Related papers  (2025-03-04T18:56:03Z)
• BRiTE: Bootstrapping Reinforced Thinking Process to Enhance Language Model Reasoning [78.63421517563056]
  Large Language Models (LLMs) have demonstrated remarkable capabilities in complex reasoning tasks.<n>We present a unified probabilistic framework that formalizes LLM reasoning through a novel graphical model.<n>We introduce the Bootstrapping Reinforced Thinking Process (BRiTE) algorithm, which works in two steps.
  arXiv  Detail & Related papers  (2025-01-31T02:39:07Z)
• The Lessons of Developing Process Reward Models in Mathematical Reasoning [62.165534879284735]
  Process Reward Models (PRMs) aim to identify and mitigate intermediate errors in the reasoning processes.<n>We develop a consensus filtering mechanism that effectively integrates Monte Carlo (MC) estimation with Large Language Models (LLMs)<n>We release a new state-of-the-art PRM that outperforms existing open-source alternatives.
  arXiv  Detail & Related papers  (2025-01-13T13:10:16Z)
• Process Reward Model with Q-Value Rankings [18.907163177605607]
  Process Reward Modeling (PRM) is critical for complex reasoning and decision-making tasks.<n>We introduce the Process Q-value Model (PQM), a novel framework that redefines PRM in the context of a Markov Decision Process.<n>PQM optimize Q-value rankings based on a novel comparative loss function, enhancing the model's ability to capture the intricate dynamics among sequential decisions.
  arXiv  Detail & Related papers  (2024-10-15T05:10:34Z)
• LLaMA-Berry: Pairwise Optimization for O1-like Olympiad-Level Mathematical Reasoning [56.273799410256075]
  The framework combines Monte Carlo Tree Search (MCTS) with iterative Self-Refine to optimize the reasoning path. The framework has been tested on general and advanced benchmarks, showing superior performance in terms of search efficiency and problem-solving capability.
  arXiv  Detail & Related papers  (2024-10-03T18:12:29Z)
• Steamroller Problems: An Evaluation of LLM Reasoning Capability with Automated Theorem Prover Strategies [0.18416014644193066]
  We evaluate the performance of GPT4, GPT3.5 Turbo and Google's recent Gemini model on problems from a steamroller domain. We found that the models' performance when using the ATP reasoning strategies was comparable to one-shot chain of thought.
  arXiv  Detail & Related papers  (2024-07-17T22:49:23Z)
• Variational Inference of Parameters in Opinion Dynamics Models [9.51311391391997]
  This work uses variational inference to estimate the parameters of an opinion dynamics ABM. We transform the inference process into an optimization problem suitable for automatic differentiation. Our approach estimates both macroscopic (bounded confidence intervals and backfire thresholds) and microscopic ($200$ categorical, agent-level roles) more accurately than simulation-based and MCMC methods.
  arXiv  Detail & Related papers  (2024-03-08T14:45:18Z)
• Self-Evaluation Guided Beam Search for Reasoning [61.523627290397556]
  We introduce a stepwise self-evaluation mechanism to guide and calibrate the reasoning process of Large Language Model (LLM) We propose a decoding algorithm integrating the self-evaluation guidance via beam search. Our approach surpasses the corresponding Codex-backboned baselines in few-shot accuracy by $6.34%$, $9.56%$, and $5.46%$ on the GSM8K, AQuA, and StrategyQA.
  arXiv  Detail & Related papers  (2023-05-01T02:37:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.