Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adaptive Test-Time Reasoning via Reward-Guided Dual-Phase Search

Yingqian Cui, Zhenwei Dai, Pengfei He, Bing He, Hui Liu, Xianfeng Tang, Jingying Zeng, Suhang Wang, Yue Xing, Jiliang Tang, Benoit Dumoulin

TL;DR

This work addresses inefficiencies in test-time reasoning by separating planning and execution into two distinct search phases guided by phase-specific reward models. The DREAM framework introduces adaptive budget allocation to stop easy steps early and reallocate resources to harder steps, improving both accuracy and efficiency on math reasoning and code generation tasks. Empirical results show DREAM and its budgeted variant outperform baselines like standard beam search and other PRM-based methods, with strong generalization to out-of-distribution datasets. The approach highlights that explicit two-phase reasoning and dynamic resource management can substantially enhance complex problem solving in LLMs.

Abstract

Large Language Models (LLMs) have achieved significant advances in reasoning tasks. A key approach is tree-based search with verifiers, which expand candidate reasoning paths and use reward models to guide pruning and selection. Although effective in improving accuracy, these methods are not optimal in terms of efficiency: they perform simple decomposition on the reasoning process, but ignore the planning-execution nature of tasks such as math reasoning or code generation. This results in inefficient exploration of reasoning process. To address this, we propose a dual-phase test-time scaling framework that explicitly separates reasoning into planning and execution, and performs search over the two phases individually. Specifically, we decompose reasoning trajectories and develop reward models for each phase, enabling the search to explore and prune plans and executions separately. We further introduce a dynamic budget allocation mechanism that adaptively redistributes sampling effort based on reward feedback, allowing early stopping on confident steps and reallocation of computation to more challenging parts of the reasoning process. Experiments on both mathematical reasoning and code generation benchmarks demonstrate that our approach consistently improves accuracy while reducing redundant computation.

Adaptive Test-Time Reasoning via Reward-Guided Dual-Phase Search

TL;DR

This work addresses inefficiencies in test-time reasoning by separating planning and execution into two distinct search phases guided by phase-specific reward models. The DREAM framework introduces adaptive budget allocation to stop easy steps early and reallocate resources to harder steps, improving both accuracy and efficiency on math reasoning and code generation tasks. Empirical results show DREAM and its budgeted variant outperform baselines like standard beam search and other PRM-based methods, with strong generalization to out-of-distribution datasets. The approach highlights that explicit two-phase reasoning and dynamic resource management can substantially enhance complex problem solving in LLMs.

Abstract

Large Language Models (LLMs) have achieved significant advances in reasoning tasks. A key approach is tree-based search with verifiers, which expand candidate reasoning paths and use reward models to guide pruning and selection. Although effective in improving accuracy, these methods are not optimal in terms of efficiency: they perform simple decomposition on the reasoning process, but ignore the planning-execution nature of tasks such as math reasoning or code generation. This results in inefficient exploration of reasoning process. To address this, we propose a dual-phase test-time scaling framework that explicitly separates reasoning into planning and execution, and performs search over the two phases individually. Specifically, we decompose reasoning trajectories and develop reward models for each phase, enabling the search to explore and prune plans and executions separately. We further introduce a dynamic budget allocation mechanism that adaptively redistributes sampling effort based on reward feedback, allowing early stopping on confident steps and reallocation of computation to more challenging parts of the reasoning process. Experiments on both mathematical reasoning and code generation benchmarks demonstrate that our approach consistently improves accuracy while reducing redundant computation.

Paper Structure

This paper contains 21 sections, 1 equation, 6 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related works
  3. Proposed Method
  4. Dual-phase search over planning and execution
  5. Construction of the Reward Models
  6. Experiments
  7. Setup
  8. Main Results.
  9. Math Reasoning
  10. Code Generation
  11. Generalization of the Math Reward Model to Unseen Datasets
  12. Ablation Studies
  13. Reward Model Size
  14. Conclusion
  15. Appendix
  16. ...and 6 more sections

Figures (6)

  • Figure 1: An example of reasoning with plan and execution as a single unit versus searched separately.
  • Figure 2: Plan–execution Format for Reasoning in Math and Code Tasks.
  • Figure 3: Workflow of Standard Beam Search and Dual-Phase Search (with budget allocation).
  • Figure 4: Accuracy vs Tokens (log) on GSM8K/MATH datasets
  • Figure 5: Accuracy vs Tokens (log) on MBPP/HumanEval datasets
  • ...and 1 more figures