Adaptive Test-Time Reasoning via Reward-Guided Dual-Phase Search (2509.25420v1)

Abstract: 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.

• The paper introduces DREAM, a dual-phase search framework that decouples planning and execution to enhance efficiency in complex reasoning tasks.
• The paper employs adaptive budget allocation and rollout-based reward models, achieving up to 20% accuracy gains on math benchmarks and significant improvements in code generation.
• The paper demonstrates robust generalization across both in-distribution and out-of-distribution datasets, indicating scalable and transferable performance.

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

Motivation and Problem Formulation

The paper introduces DREAM, a dual-phase test-time scaling framework for LLM reasoning that explicitly decomposes the reasoning process into planning and execution phases, each with dedicated search and reward modeling. Existing tree-based search methods for LLMs, such as PRM and MCTS variants, typically treat planning and execution as a unified pipeline, evaluating plan–execution pairs as atomic units. This approach is suboptimal: correct plans may be discarded due to execution errors, and flawed plans may waste computation on their executions. Furthermore, most prior methods allocate a fixed sampling budget per reasoning step, ignoring the heterogeneity in step difficulty within and across problems.

Figure 1: An example of reasoning with plan and execution as a single unit versus searched separately.

The paper argues that complex tasks like math problem solving and code generation naturally involve distinct cognitive phases—high-level strategic planning and low-level execution. By searching and pruning plans and executions separately, computation can be allocated more efficiently, and promising plans can be paired with multiple execution attempts, increasing the probability of correct solutions.

Dual-Phase Search Framework

DREAM formalizes reasoning in a plan–execution format, where each step is decomposed into a plan (subgoal formulation) and an execution (concrete calculation or implementation). Few-shot prompting is used to induce LLMs to generate solutions in this format.

Figure 2: Plan–execution Format for Reasoning in Math and Code Tasks.

The dual-phase search algorithm extends standard beam search by introducing two separate phases per reasoning step:

• Planning Phase: $N_1$ candidate plans are sampled and scored by a planning reward model ($\text{PRM}_{\text{plan}}$). The top $n_1$ plans are selected.
• Execution Phase: For each selected plan, $N_2$ candidate executions are generated and scored by an execution reward model ($\text{PRM}_{\text{exec}}$). The top $n_2$ executions are retained for expansion.

This separation allows early pruning of weak plans and multiple execution attempts for promising plans, mitigating the risk of discarding good strategies due to execution errors.

Figure 3: Workflow of Standard Beam Search and Dual-Phase Search (with budget allocation).

Adaptive Budget Allocation

To further improve efficiency, DREAM incorporates a dynamic budget allocation mechanism. Sampling in both phases follows a two-threshold rule:

• Early Stopping: If $n_1$ (planning) or $n_2$ (execution) candidates exceed a high reward threshold, sampling is terminated early.
• Additional Sampling: If no candidate exceeds a lower threshold after exhausting the budget, additional samples are generated.

This adaptive allocation enables early stopping on easy steps and reallocates computation to harder steps, aligning compute with step-level difficulty.

Reward Model Construction

Reward models for both planning and execution are trained using rollout-based annotation. For each intermediate plan or execution, multiple continuations are generated; if any lead to a correct final answer, the step is labeled positive. The reward model is fine-tuned as a next-token prediction task, outputting a binary label (“+” or “–”) at the final position. The reward is computed as the softmax probability of the “+” token.

$$\text{Reward}(x) = \frac{\exp(\ell_{+}(x))}{\exp(\ell_{+}(x)) + \exp(\ell_{-}(x))}$$

where $\ell_{+}(x)$ and $\ell_{-}(x)$ are the logits for the “+” and “–” tokens, respectively.

Experimental Results

Math Reasoning

DREAM is evaluated on GSM8K and MATH benchmarks, using Qwen2.5-32B-Instruct as the reward model and various LLMs for reasoning. The method is compared against majority vote, standard beam search, and REBASE. All methods use the plan–execution format and the same reward model for fairness.

Figure 4: Accuracy vs Tokens (log) on GSM8K/MATH datasets

Key findings:

• Tree-based search methods (including DREAM) achieve a substantially better accuracy–efficiency trade-off than majority vote, with up to 20% higher accuracy on MATH.
• DREAM outperforms standard beam search: Independent sampling and selection of planning and execution yields higher-quality trajectories, especially as token budget increases.
• Dynamic budget allocation (DREAM(+)) provides further gains, with up to 2% improvement in accuracy at comparable token budgets on GSM8K. The benefit is more pronounced when step difficulty is variable.
• Reward model generalization: DREAM maintains strong performance even when the reward model is trained on different LLMs than those used for evaluation, indicating robust generalization.

Code Generation

DREAM is integrated with CodeTree and compared against Reflexion and standard CodeTree on MBPP and HumanEval benchmarks.

• CodeTree+DREAM achieves up to 10% higher accuracy than CodeTree at similar token budgets.
• Budget allocation further improves accuracy–efficiency trade-off.
• Reflexion performs well at low token budgets but saturates earlier than CodeTree+DREAM as budget increases.
• Reward model generalizes to unseen datasets (e.g., HumanEval not in training data).

Generalization to Out-of-Distribution Datasets

DREAM is evaluated on AMC23 and ASDiv, which are out-of-distribution for the reward model. DREAM consistently achieves higher accuracy than majority vote, with up to 30% improvement on AMC23, demonstrating strong transferability.

Ablation Studies

• Reward Model Size: Larger reward models (Qwen2.5-32B-Instruct) outperform smaller ones (Qwen2.5-7B-Instruct), but even 7B models provide substantial benefits over majority voting.
• Critic Agent in Code Generation: Removing the critic agent improves efficiency without sacrificing accuracy.
• Synergy of Dual-Phase Search and Budget Allocation: Budget allocation improves accuracy–efficiency trade-off for both beam search and DREAM, but gains are larger for DREAM, indicating complementary effects.

  Figure 5: Accuracy-Tokens Frontier across methods.

Implications and Future Directions

The dual-phase search paradigm introduced by DREAM provides a principled approach to decomposing reasoning into planning and execution, enabling finer-grained control over search and resource allocation. The rollout-based reward modeling framework yields robust supervision signals that generalize across models and datasets. The adaptive budget allocation mechanism aligns compute with step-level difficulty, improving efficiency.

Practically, DREAM can be applied to any reasoning task expressible in a plan–execution format and amenable to tree-based search. The framework is compatible with various LLMs and reward model architectures. The demonstrated generalization of reward models suggests that transfer learning and cross-task adaptation are feasible.

Theoretically, the separation of planning and execution aligns with cognitive models of problem solving and may inform future research on modular reasoning architectures. The synergy between dual-phase search and dynamic budget allocation highlights the importance of adaptive computation in LLM inference.

Future work may explore:

• Extending dual-phase search to multi-agent or hierarchical reasoning settings.
• Developing more sophisticated reward models, e.g., leveraging uncertainty estimation or meta-learning.
• Integrating dual-phase search with distributed or parallel inference frameworks for large-scale deployment.
• Investigating the limits of reward model generalization and transfer across domains.

Conclusion

DREAM introduces a dual-phase test-time scaling framework for LLM reasoning, explicitly separating planning and execution and equipping each phase with dedicated reward models and adaptive budget allocation. Empirical results demonstrate consistent improvements in accuracy and efficiency over standard beam search and prior PRM-based methods, with strong generalization to out-of-distribution tasks. The dual-phase paradigm and adaptive allocation mechanisms provide a robust foundation for future advances in scalable, efficient, and interpretable LLM reasoning.