Policy Guided Tree Search for Enhanced LLM Reasoning

We thank the reviewer for their thoughtful assessment of our work and the recommendation for acceptance. We appreciate the positive feedback on our method's performance, efficiency, and adaptive policy design. Below we address the questions and suggestions raised:

On Quantifying "Overthinking" in LLMs

The reviewer rightly suggests quantitative metrics on reasoning step counts to support our claim about reducing overthinking. Our analysis shows that:

• For StrategyQA, PGTS uses on average 5.97 reasoning steps per problem vs. MCTS's 50.38 and CoT's 4.75 steps.
• For GSM8K problems, PGTS uses on average 10.73 reasoning steps per problem vs. MCTS's 91.97 and CoT's 7.35 steps.
• For MATH problems, PGTS uses 61.64 steps on average vs. MCTS's 220.24 steps.
• For Auqa problems, PGTS uses 17.46 reasoning steps vs. MCTS's 64.20 steps.

Importantly, PGTS's "terminate" action effectively stops the reasoning process once sufficient evidence has been generated, reducing steps compared to MCTS while improving accuracy.

These metrics confirm that PGTS achieves better accuracy while using fewer reasoning steps, directly addressing the overthinking problem.

On Generalization Without Dataset-Specific Tuning

We'd like to clarify our approach to hyperparameter tuning. While we did train separate policies for each benchmark, we deliberately used the same:

• Network architecture (2-layer GPS)
• Action costs (0.1, 0.2, 0.5, 0.0 for expand, branch, backtrack, terminate)
• PPO hyperparameters
• State/action representations

This demonstrates that our framework generalizes across diverse reasoning tasks without requiring task-specific architectural changes or cost function designs. Unlike RAP, which requires careful reward engineering for each task domain, PGTS maintains a unified approach across tasks. We agree that greater domain-generalizability is an important direction for future work and will add this discussion to the paper.

On Additional Baseline Comparisons

We thank the reviewer for suggesting Tree-of-Thought (ToT) and Algorithm-of-Thought (AoT) as additional baselines. However, we note several challenges with direct comparisons:

• ToT requires task-specific design for proposal and evaluation prompts, which is not a standardized process and introduces variability that would complicate fair comparisons.
• AoT requires a strong model to generate complete, executable code as noted in Section 7 of their paper, making it particularly challenging to implement consistently across our diverse benchmarks.

PGTS addresses these limitations by learning a task-agnostic policy that doesn't require manual prompt engineering or code generation capabilities.

On Qualitative Evaluation

We appreciate the reviewer's suggestion regarding human evaluation of reasoning chains. We acknowledge that this type of qualitative assessment would provide valuable insights about the actual quality and coherence of reasoning paths beyond accuracy metrics. While we haven't conducted a formal human evaluation in the current work, we agree this would be an important direction for future work. Such evaluation could assess factors like logical coherence, correctness of intermediate steps, and overall reasoning quality across different methods. We will add this consideration to our discussion of future work in the revised paper.

On Evaluation on Open-ended Tasks

We acknowledge the reviewer's point about evaluating on more open-ended reasoning tasks. We agree that our current approach requires training a policy for each task domain, which is a limitation for truly open-ended generalization.

For future work, we plan to explore meta-learning and multi-task reinforcement learning approaches to train a single policy across diverse reasoning tasks, which could enable zero-shot or few-shot adaptation to new domains. This aligns with active research directions in RL literature on cross-task generalization and transferable policies. By training on a diverse set of reasoning tasks simultaneously, we believe PGTS could develop more general reasoning strategies that transfer to unseen problems without task-specific fine-tuning.

Conclusion

We thank the reviewer for their thoughtful feedback, which has helped us strengthen our claims with additional quantitative and qualitative evidence. We believe PGTS makes significant contributions to structured reasoning in LLMs by learning an adaptive search policy, reducing overthinking, and improving efficiency across diverse reasoning tasks.

We thank the reviewer for their positive assessment of our work and the constructive feedback. We address the key points below:

Missing References

We appreciate the suggestion to discuss the works by Zelikman et al. (STAR) and Zhang et al. (REST-MCTS*). These papers indeed share conceptual connections with our approach:

• STAR bootstraps reasoning by selecting and refining high-quality reasoning chains, focusing on improving reasoning through iterative training.
• REST-MCTS* combines tree search with self-training, using process rewards to guide exploration.

Our approach differs in that PGTS trains a policy to guide exploration during inference rather than improving the reasoning capabilities of the LLM itself. The learned policy enables dynamic decisions about exploration paths without modifying the underlying LLM. We will include these references and discuss their relationship to our work in the revised paper.

Action Importance

Regarding $C(a)$ and action importance, we want to clarify that the cost values should not be interpreted as reflecting the relative importance of each action. Rather, they serve as soft constraints to discourage unnecessary branching and backtracking, promoting exploration efficiency.

The branching and backtracking actions are fundamental to effective reasoning under the PGTS framework. Without these exploration capabilities, PGTS would be functionally equivalent to a standard Chain-of-Thought (CoT) approach - limited to generating a single linear reasoning path. These structured exploration actions constitute the core distinction between PGTS and CoT, enabling more robust reasoning through systematic exploration of alternative paths and recovery from suboptimal reasoning steps.

Examples of branching and backtracking are shown in Figures 5-8 in the Appendix, including Figure 5 (branching at node 4), Figure 6 (backtracking at node 3), Figure 7 (backtracking at node 5), and Figure 8 (branching after reaching an answer).

Policy Function Compared to PRM

The reviewer correctly identifies that our policy functions similarly to a Process Reward Model (PRM) at inference time. However, there are important distinctions:

1. State representation: Our policy operates on the tree structure rather than just the current reasoning path, enabling informed backtracking and branching decisions based on the full exploration history.

2. Action space: Traditional PRMs evaluate individual steps, while our policy chooses between structured exploration actions (expand, branch, backtrack, terminate). Critically, PRMs don't inherently determine when to backtrack or how many steps to backtrack - these strategic decisions require understanding the global reasoning context and are naturally captured in our policy's action space.

This design allows PGTS to more efficiently guide the reasoning process while maintaining the flexibility to explore diverse paths as needed.

Blocksworld Training Set Sampling

Regarding the Blocksworld training data, we use their split-v1 for training and split-v2 for testing. For clarification:

1. We used the original problem generator from Valmeekam et al. with the same configuration parameters.
2. Following RAP's methodology, we generated distinct training (split-v1) and testing (split-v2) sets.

This approach ensures a fair comparison with existing methods while maintaining the challenging nature of the planning task.These details are already included in appendix.

Comparison with Other Training-Based Approaches

We agree that comparing with other training-based approaches would strengthen our evaluation. While our original experiments focused on comparing with zero-shot methods (CoT and MCTS), there are two key distinctions when comparing PGTS with other training-based methods:

• Training Data Efficiency: PGTS requires only problem statements and answers for training (not annotated reasoning chains), making it more practical than methods requiring extensive reasoning annotations. Our approach achieves strong performance with only 1,000 examples.
• Model Preservation: PGTS keeps the LLM parameters unchanged, maintaining its general capabilities while learning to guide its reasoning process.

Our preliminary experiments show that directly fine-tuning LLaMA3.1-8B on the same limited dataset (1,000 examples) with CoT annotations produces negligible improvements over the base model.

We appreciate the reviewer's valuable feedback and believe addressing these points will further strengthen our paper.

We thank the reviewer for their thoughtful and constructive feedback. We address the key points raised below:

Fair Comparison of Computational Budget

The reviewer raises an excellent point about fair comparisons between methods.

However, on most datasets, the token usage ratio is much more favorable, and PGTS often outperforms even CoT with self-consistency:

• On AQUA: PGTS achieves 60.63% vs CoT (SC4)'s 53.94% while using only 1.42× tokens compared to CoT (less than generating 4 samples)
• On StrategyQA: PGTS achieves 77.70% vs CoT (SC4)'s 77.10% while using only 1.58× tokens compared to CoT
• On GPQA: PGTS achieves 18.69% vs CoT (SC4)'s 15.66% with modest token increase
• On Blocksworld (4): PGTS achieves 27.63% vs CoT (SC4)'s 22.37% with efficient token usage

More importantly, our primary contribution is improving efficiency compared to existing tree search methods like MCTS. Across all datasets, PGTS uses approximately 1/3 of the tokens required by MCTS while achieving comparable or better performance. This represents a significant efficiency gain in the space of tree search approaches for reasoning.

PGTS effectively occupies a middle ground in the accuracy-efficiency tradeoff space:

1. CoT methods: Lowest token usage but limited exploration capability
2. PGTS: Moderate token usage with adaptive exploration capability
3. MCTS: Highest token usage with exhaustive exploration capability

For applications requiring the exploration benefits of tree search but facing computational constraints, PGTS offers a valuable alternative that wasn't previously available. It provides a new point on the Pareto frontier of reasoning methods.

In the revised paper, we will clarify these comparisons, present dataset-specific efficiency analyses, and more precisely position PGTS relative to both CoT and MCTS baselines.

Technical Details and Algorithm Clarity

The PGTS approach follows standard reinforcement learning methodology with a Graph Neural Network policy. Section 2.3.4 describes our policy architecture using GPS layers to process the tree-structured state, though we can clarify specific implementation details:

• Policy network: Two GPS layers with 256 hidden units each (~1M parameters)
• Input: Tree structure with node features (768d hidden states from the LLM's final layer) and edge features (1d immediate rewards)
• Output: Categorical distribution over D+2 actions (expand, branch, D-1 backtrack depths, terminate)

For the algorithm, Section 2.3 provides the conceptual framework while Algorithm 1 in the Appendix details the training procedure. We've included our complete implementation in the supplementary materials, and we will release the code publicly upon publication for full reproducibility.

We believe the core algorithmic concepts are clearly presented in the current manuscript, but we'll ensure the revised version includes additional implementation details and more comprehensive pseudocode for the inference process to improve clarity.

Reward Signal for Termination

For the terminate action, $R(s_d)$ evaluates the final reasoning chain quality using task-specific measures (exact match for math, binary correctness for yes/no tasks, goal achievement for planning). These are used only during policy training, not inference.

Wall-time Efficiency

LLM inference dominates computational cost (~99% of total time). PGTS policy inference adds minimal overhead (~5-10ms per decision) compared to LLM generation time (100-500ms). While CoT with self-consistency can be fully parallelized, the sequential nature of tree search methods introduces some latency that's offset by their improved reasoning capabilities. However, both approaches can be batched across multiple reasoning problems simultaneously. At larger batch sizes processing multiple queries simultaneously, these differences become less significant as GPU utilization improves for all methods. The additional wall-time cost of PGTS compared to CoT is a reasonable tradeoff for applications requiring more reliable reasoning, especially considering its substantial efficiency advantage over other tree search methods like MCTS.

Fast Convergence and Broader Applicability

PGTS converges quickly because it learns exploration strategies (not domain reasoning) over a small action space using an architecture well-suited to tree structures. Beyond LLMs, PGTS could enhance planning in robotics, program synthesis, or game playing where structured exploration with intermediate feedback is possible.

We appreciate the reviewer's encouraging words about the research direction and will address these concerns in the revised version to improve clarity and rigor.

We thank the reviewer for their thoughtful and constructive feedback. We address the key points raised below:

PGTS Performance on Complex Tasks like GPQA

The reviewer correctly notes that PGTS underperforms MCTS on GPQA (27.78% vs. 34.85%). This performance gap can be attributed to three primary factors:

1. Task Complexity: GPQA consists of graduate-level questions across diverse domains requiring specialized knowledge, making it particularly challenging for a small policy without extensive world knowledge.
2. Limited Training Data: GPQA's diversity and complexity means that our limited training data (250 examples) may not provide sufficient coverage across all domains for the policy to learn effective exploration strategies.
3. Task-Specific Reward Design: Our current reward design uses the same structure across all tasks, with the intermediate reward simply being the likelihood of each reasoning step, which may not be optimal for GPQA's unique characteristics.

We believe this is not a fundamental limitation of PGTS but rather an opportunity for task-specific enhancement. We plan to address this in future work through:

• Developing more sophisticated reward mechanisms tailored to complex reasoning tasks
• Exploring synthetic data generation to improve training diversity and coverage

Training with More Examples

The plateau observed in Figure 3 occurs primarily for simpler reasoning tasks where 1,000 examples provide sufficient coverage of reasoning patterns. For complex tasks like GPQA, we expect the learning curve would continue to improve with substantially more training examples, especially if they provide better domain coverage.

GPQA presents a particular challenge due to its limited training data availability (we were restricted to 250 examples) and the exceptional breadth of graduate-level topics it covers. This limitation likely contributes to the performance gap compared to MCTS on this specific benchmark.

To address this challenge in future work, we plan to explore synthetic data generation techniques to expand training diversity.

Action Definitions Beyond Sentence Splitting

The reviewer raises an important point about action granularity. We chose sentence-level segmentation for simplicity and consistency across diverse tasks. However, we agree that semantically meaningful action definitions could further improve performance, and our PGTS framework can easily adapt to different action deefinitions.

We plan to explore hierarchical reasoning structures and task-specific action definitions in future work. This would allow PGTS to operate at multiple levels of abstraction, potentially enhancing both efficiency and performance.

Clarifications on Algorithm 1

Thank you for noting the undefined symbols in Algorithm 1. We will revise the algorithm description to clearly define:

• H: The entropy of the policy outputs, which is used as regularization to encourage exploration.
• $\omega$: The entropy regularization coefficient (set to 0.01 in our experiments)
• $\eta_{\pi}$ and $\eta_{V}$: Learning rates for the policy and value networks (both set to 3e-4)

We will include these definitions in the final version.

Ablation on Cost Function Settings

We acknowledge the lack of ablation analysis on the cost function $C(a)$. In our preliminary experiments, we found that performance is relatively robust to cost settings within reasonable ranges, with the relative ordering (backtrack > branch > expand > terminate) being more important than absolute values.

We'd be happy to include a more detailed analysis of cost function sensitivity in the final version.

We appreciate the reviewer's thoughtful comments and will incorporate these suggestions in the final version.