Reward Model Generalization for Compute-Aware Test-Time Reasoning

TL;DR

This paper introduces a theoretical framework analyzing how the generalization error of reward models impacts compute efficiency and reasoning accuracy in test-time reasoning for large language models, and proposes a dynamic search method called CATS.

Contribution

It provides a PAC-Bayes based theoretical analysis of test-time reward model generalization and introduces CATS, a compute-aware search method that improves reasoning performance under compute constraints.

Findings

CATS outperforms existing methods on MATH and AIME benchmarks.

Lower reward model generalization error reduces sample requirements for correct answers.

Theoretical bounds link reward model error to compute efficiency.

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.