# Reward Model Generalization for Compute-Aware Test-Time Reasoning

arXiv: https://arxiv.org/abs/2505.18065 · HTML: https://mlsanity.com/paper/2505.18065/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 14 of its 48 references that are in the map.
- First posted on arXiv 23 May 2025.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 14 of 14)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Solving math word problems with process- and outcome-based feedback](https://mlsanity.com/paper/2211.14275/) (2022)
- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [Self-critiquing models for assisting human evaluators](https://mlsanity.com/paper/2206.05802/) (2022)
- [RLHF Workflow: From Reward Modeling to Online RLHF](https://mlsanity.com/paper/2405.07863/) (2024)
- [Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning](https://mlsanity.com/paper/2405.00451/) (2024)
- [AlphaMath Almost Zero: Process Supervision without Process](https://mlsanity.com/paper/2405.03553/) (2024)
- [Accessing GPT-4 level Mathematical Olympiad Solutions via Monte Carlo Tree Self-refine with LLaMa-3 8B](https://mlsanity.com/paper/2406.07394/) (2024)
- [LLaMA-Berry: Pairwise Optimization for O1-like Olympiad-Level Mathematical Reasoning](https://mlsanity.com/paper/2410.02884/) (2024)
- [Enhancing LLM Reasoning with Reward-guided Tree Search](https://mlsanity.com/paper/2411.11694/) (2024)
- [LiteSearch: Efficacious Tree Search for LLM](https://mlsanity.com/paper/2407.00320/) (2024)
- [BoostStep: Boosting mathematical capability of Large Language Models via improved single-step reasoning](https://mlsanity.com/paper/2501.03226/) (2025)

## Abstract (arXiv)

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.