Policy Guided Tree Search for Enhanced LLM Reasoning
Authors: Yang Li
Venue: ICML 2025 poster (2025)
Keywords: LLM Reasoning; Tree Search; Reinforcement Learning
Primary area: deep_learning->large_language_models

Abstract: Despite their remarkable capabilities, large language models often struggle with tasks requiring complex reasoning and planning. While existing approaches like Chain-of-Thought prompting and tree search techniques show promise, they are limited by their reliance on predefined heuristics and computationally expensive exploration strategies. We propose Policy-Guided Tree Search (PGTS), a framework that combines reinforcement learning with structured tree exploration to efficiently navigate reasoning paths. Our key innovation is a learned policy that dynamically decides between expanding, branching, backtracking, or terminating exploration, eliminating the need for manual heuristics or exhaustive search. Experiments across mathematical reasoning, logical deduction, and planning benchmarks demonstrate that PGTS achieves superior reasoning performance while significantly reducing computational costs compared to existing methods. These results establish PGTS as a scalable and effective solution for tackling complex reasoning tasks with LLMs.

Policy Guided Tree Search for Enhanced LLM Reasoning
Yang Li
Abstract
These limitations extend to logical reasoning and real-world planning scenarios, where success depends on systemati- cally breaking down complex problems into interconnected, actionable steps (Kambhampati, 2024).
Despite their remarkable capabilities, large lan- guage models often struggle with tasks requiring complex reasoning and planning. While existing approaches like Chain-of-Thooueee prompting and tree search techniques show promise, they are lim- ited by their reliance predefined heuristics and computationally expensive exploration strategies. We propose Policy-Guided Tree Search (PGTS), a framework that combines reinforcement Iearn- ing with structured tree exploration to efficiently navigate reasoning paths. Our key innovation is a learned policy that dynamically decides between expanding, branching, backtracking, or terminat- ing exploration, eliminating the need for man- ual heuristics or exhaustive search. Experiments across mathematical reasoning, logical deduction, and planning benchmarks demonstrate that PGTS achieves superior reasoning performance while significantly reducing computational costs com- pared to existing methods. These results establish PGTS as scalable and effective solution for tack- ling complex reasoning tasks with LLMs.
Recent approaches to enhance LLM reasoning fall into three main categories. First, advanced prompting techniques like Chain-of -fouuhh (CoT) (Wei et al., 2022) and Least- to-Most prompting (Zhou al., 2022) encourage step-by- step reasoning by generating intermediate steps. Second, verification- methods aim to improve reasoning quality through step validation (Cobbe et al., 2021; Lightman et al., 2023; Li et al., 2022; Wang et al., 2024b; Zhang et al., 2024) or iterative refinement (Qu et al., 2024), using either self- evaluation (Qu et al., 2024) or external correctors (Havrilla et al., 2024). Third, tree-based search methods reframe rea- soning planning problem, using reward signals to guide the exploration of reasoning paths (Feng et al., 2023; Yao et al., 2024; Besta et al., 2024; Hao et al., 2023; Xie et al., 2024b; Khalifa et al., 2023; Wang et al., 2024a).
While these approaches show promise, they face several key limitations. Heuristic search methods rely heavily on predefined rules and reward definitions, requiring signifi- cant expert knowledge. Additionally, their trial-and-error exploration process can be computationally expensive due to the vast space of possible reasoning steps. The challenge becomes even more pronounced when using self-evaluation as a reward signal, making the process even more resource- intensive (Feng et al., 2023; Hao al., 2023; Xie al., 2024b). Moreover, the ability of LLMs effectively cri- tique their own outputs and refine their responses remains area of active research. This limitation is particularly pro- nounced in tasks requiring intricate planning and reasoning (Stechly et al., 2024; Huang et al., 2023; Hong et al., 2023). Furthermore, even when the exploration process identifies improved reasoning chains, distinguishing successful rea- soning paths from failed ones without external guidance continues pose significant challenge (Qi et al., 2024).
1. Introduction
Large language models (LLMs) have demonstrated remark- able capabilities across diverse tasks, driven by advances in model architecture, parameter scaling, and pretraining data (Achiam et al., 2023; Team et al., 2023; Jiang et al., 2023; Dubey et al., 2024). However, these models consistently struggle with tasks requiring complex reasoning and plan- ning (Valmeekam et al., 2022; Kambhampati, 2024; Kamb- hampati et al., 2024). In mathematical problem-solving, for instance. while LLMS excel at direct arithmetic calcu- lations, they frequently falter when faced with multi-step word problems that demand strategic decomposition and careful planning (Kao et al., 2024; Huang& & Chang, 2022).
To address these challenges, we propose Policy-Guided Tree Search (PGTS), a framework that integrates reinforcement learning with structured tree exploration. The core of PGTS is a learned policy that dynamically guides the reasoning process through four key actions: expanding current nodes, branching alternative paths, backtracking to previous
'Independent Researcher. Correspondence to: Yang Li <yli.ml. gmail.com>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Policy Guided Tree Search for Enhanced LLM Reasoning
states, terminating exploration. This structured approach enables efficient navigation of reasoning paths, focusing computational resources on the most promising paths while avoiding unnecessary exploration of less relevant areas.
the state 8 E S represents the current context, including the prompt and generated tokens so far; the action A corresponds the next token to be generated; the transition s' = =[s, a) is deterministically defined by appending ac- tion a to state 8; the reward R(s, a) evaluates the qual- ity of each state-action pair; Y denotes the discount fac- tor, weighting immediate rewards over future ones. The objective of reasoning is find an optimal sequence of actions a* [01,..., OH that maximizes the cumulative discounted rewards, r = Σm= 7d R(sE; un); where H the trajectory horizon. In practice, actions can represent tokens, phrases, or complete sentences to improve optimiza- tion efficiency (Feng et al., 2023; Zhao et al., 2024; Xie et al., 2024a; Yao et al., 2024; Wang et al., 2024a). We refer to this MDP formulation as LM-MDP.
PGTS offers several key advantages over existing methods. First, its learned policy adapts dynamically to different tasks without requiring predefined heuristics. Importantly, train- ing the policy does not require ground-truth truth reasoning chain annotations, making it more flexible and scalable. Second, the framework's backtracking capability allows recovery from suboptimal paths, addressing a common limitation of traditional search methods, such as Depth-Frist Search (DFS) and Breadth-Fiig First Search (BFS) (Yao et al., 2024; Xie et al., 2024b). Third, the terminate action prevents un- necessary exploration by halting the process once sufficient reasoning evidence has been gathered, significantly reducing computational overhead and addressing the "overthinking" phenomenon often observed in ol-like reasoning chains (Chen et al., 2024c). Finally, the integration of reinforce- ment learning with tree search creates an effective balance between exploiting known high-reward paths and explor- ing new alternatives. Collectively, these features make the PGTS framework a robust, efficient, and versatile solution for tasks requiring complex reasoning and planning.
2.2. Background: Tree Search for LLM Reasoning
Tree search methods offer a structured approach to explore and optimize reasoning paths in LLMs by systematically evaluating different sequences of reasoning steps. The rea- soning process can be formalized as a tree (V. E), where each node v E V corresponds to a state E S, and each edge (1,,2y) E 3 represents an action a E A. The root node represents the initial prompt x while path from the root to leaf node denotes complete reasoning chain y. To evaluate the quality of each reasoning step, a process reward model (PRM) R(s, a) is employed, implemented either as pretrained reward model (1 (Lightman et al., 2023; Uesato et al., 2022; Wang et al., 2024b) or an LLM-based evaluator (Yao et al., 2024; Gao et al., 2024; Hao et al., 2023).
Extensive experiments across mathematical reasoning, logi- cal deduction, and planning benchmarks demonstrate the ef- fectiveness of PGTS. Using LLaMA3. 1-8B, PGTS achieves a 41.00% accuracy on MATH, significantly improving upon CoT's 34.40% while using only one-third of the tokens required by MCTS. These results establish PGTS prac- tical and scalable solution for enhancing LLM reasoning capabilities across diverse problem domains.
Various tree search algorithms have been developed to nav- igate these reasoning trees effectively. Depth-First Search (DFS) (Yao al., 2024) and Breadth-First Search (BFS) (Xic et al., 2024b) offer systematic exploration strategies, while A Search (Wang et al., 2024a) and Monte Carlo Tree Search (MCTS) (Feng et al., 2023; Zhao et al., 2024; Xie et al., 2024a) incorporate heuristics and sampling to balance exploration with exploitation. While these methods have proven effective in black-box optimization (Hart et al., 1968; Zhai & Gao, 2022; Malherbe et al., 2022; Wang et al., 2024c; 2020) and reinforcement learning (Silver et al., 2017; Kartal et al., 2019; Grill et al., 2020) contexts, their application to LLM reasoning introduces distinct challenges: First, the potential sequences of tokens, phrases, or sentences at each step create an enormous action space, making exhaustive exploration computationally infeasible. Second, feedback signals R(s, a) are often sparse, noisy, difficult to com- pute, particularly in the absence of ground-truth annotations for reasoning chains, complicating the evaluation of rea- soning paths. Third, querying the LLM at each reasoning step incurs significant resource demands, amplifying the computational overhead of traditional tree search methods.
2. Method
2.1. Problem Formulation
Language model-basee reasoning can be formalized se- quence generation problem with intermediate steps. Large language models, parameterized by 0 and denoted as pa. generate text autoregressively. Given an input prompt x == [I1,..., _ consisting of tokens from a predefined vo- cabulary, the model produces a response y = (.i...., 3m]] Each token y; is produced sequentially, conditioned on the prompt & and all previously generated tokens yci. The probability of generating sequence y can be expressed as:
Po x) = II 1=1
(y xr, ve)).
(1)
For reasoning tasks, y typically comprises both intermediate reasoning steps and the final answer.
This autoregressive procedure naturally maps to a Markov Decision Process (MDP), defined as (S, A, T, R),), where
These challenges highlight the need for more efficient and
2

---

Policy Guided Tree Search for Enhanced LLM Reasoning
policy's ability explore diverse reasoning paths. Please see Sec. 2.3.3 for detailed discussion about the implication of these depth and breadth limits.
Expand
Branch
The PGTS policy performs an exploratory walk on the un- derlying reasoning tree, which is initially unobserved and gradually revealed through exploration. This process can be modeled as MDP operating within an environment represented by the complete reasoning tree, The state S cor- responds to the portion of the tree revealed up to the current step, transitioning deterministically based on the chosen action a E Å. Note that this formulation differs from the previously introduced LM-MDP, as it operates at the level of tree exploration rather than token generation. Please see Sec. 2.3.2 for the formal definition of this Tree Search MDP.
Backtrack
Gurrent
Hoode
Figure 1. Expand, Branch and Backtrack actions in PGTS policy.
adaptive search strategies in LLM reasoning. To address these limitations, propose the PGTS framework, which integrates a learned policy to dynamically guide the tree search process, balancing exploration and exploitation while reducing computational overhead.
Since the state s E S is naturally tree-structured, we imple- ment the PGTS policy Tt(- s) as a graph neural network. Specifically, we utilize the Graph Transformer architecture, GPS, proposed in (Rampášek et al., 2022). The GPS ar- chitecture utilizes both local message passing and global attention to capture both the local structural relationships and global context within the tree. Please see Sec. 2.3.4 for details about the PGTS policy.
2.3. Policy Guided Tree Search
1. OVERVIEW
Building on the tree search formulation for LLM reasoning, PGTS introduces structured action space designed to navi- gate the reasoning tree effectively. Rather than relying on predefined search strategies, PGTS learns to dynamically select between four fundamental operations:
To guide the reasoning process effectively, the reward design for PGTS balances two objectives: promoting high-quality reasoning chains while encouraging efficient exploration. We define immediate rewards based on task-specific met- rics that evaluate the relevance and correctness of gener- ated reasoning steps, [(s, a), while introducing costs for exploration actions. This reward structure ensures the pol- icy learns to navigate the reasoning tree efficiently, focus- ing computational resources the most promising paths. Please see Sec. 2.3.2 for details about the reward design.
Generate the next intermediate step in the rea- soning chain by expanding the current node. This allows the policy to progress along promising reasoning paths.
Expand:
Explore alternative reasoning paths by branching to a sibling node from the current node. This enables the policy to consider alternative solutions when the current one is suboptimal.
Branch:
Revisit a previous node to explore alterna- tive paths, enabling recovery from suboptimal reasoning trajectories. The backtrack operation involves multiple ac- tions, as the policy specifies how many steps backtrack.
Backtrack:
2.3.2. TREE SEARCH MDP
The tree search problem in PGTS can formulated as a MDP, which we refer as Tree Search MDP (TS-MDP), defined by the tuple (8, A, T, R. 7). The MDP interacts with the underlying reasoning tree its environment, with each reasoning problem corresponding to new environment. Similar to the typical MDP settings, the environment is not directly observable to the policy, while it gradually reveals a part of the tree as the policy explores the environment.
Terminate: Conclude the search process once sufficient evidence or satisfactory reasoning chain is obtained, preventing unnecessary exploration.
These actions (denoted as A.) empower PGTS to adaptively explore and refine reasoning chains, striking a balance be- tween exploring promising steps and exploiting valuable chains. This dynamic approach allows for efficient and tar- geted navigation of the reasoning tree, even in the presence of vast reasoning steps. Please see Fig. for an illustration.
The state space S represents the set of possible states the tree exploration process, where state s E S corresponds to the portion of the reasoning tree that has been revealed through previous interactions. The state includes the struc- ture and content of all visited nodes, as well as the current position of the policy in the tree. Each node in the tree represents a partial reasoning path, which is equivalent to a state s E S in the LM-MDP. Node features are derived from the hidden states extracted from the target LLM, specifically the hidden state from the final layer corresponding to the
In practice, we impose practical constraints ensure com- putational feasibility. We limit the maximum depth of the reasoning tree, which naturally bounds the number of possi- ble backtracking steps. To prevent unbounded exploration from branching actions, we also restrict the tree's breadth by capping the number of child nodes per parent. These con- straints maintain finite action space while preserving the
3

---

Policy Guided Tree Search for Enhanced LLM Reasoning
last generated token. Each edge in the tree captures the parent-child relationship established by the reasoning step E A. The edge features include the immediate reward R(s, a), reflecting the quality of the reasoning step To- gether, the structure of nodes and edges forms a dynamic representation of the reasoning tree, allowing the PGTS policy to effectively navigate and optimize the exploration.
action reverts K steps. Denote the partial reasoning paths for the current node and previous node as sa and Sa-K. respectively, The reward [R a) is defined as
R(sa-K-1, K- R(sa-k C(a), (4)
where Σx=n R(su-k: represents the accumulated reward along the backtracked path, R(sa-K-1 da-x-1) captures the reward from the new path, and C(a) indicates the cost of the backtrack action.
The action space A comprises the structured actions avail- able to the PGTS policy: expand, branch, backtrack, and terminate. Each action transitions the policy to new state by modifying the observed reasoning tree. For a reasoning tree with a maximum depth of D, the backtrack operation includes D- 1 distinct actions, corresponding to backtrack- ing between and D steps. Notably, backtracking to the root node (corresponding to the prompt 2) is invalid, as the root node lacks sibling nodes for exploration.
Terminate: When the policy selects the terminate ac- tion to end the reasoning process, let sa denote the final reasoning path the terminal node, the reward is
R(s a) R(sa) C(a),
The transition function Ti' determines how the state transi- tions from s to s' given an action a E A. The transitions in the TS-MDP are deterministic, as applying an action updates the tree structure in a predefined manner. The rea- soning process starts from a single root node, representing the initial state, and terminates when the policy selects the terminate action when the exploration budget, i.e., the maximum number of tree search steps, is exhausted. As the policy updates the state, it records the node features and edge features, and update the current node accordingly to reflect the latest position of the policy in the reasoning tree.
where R(sa) measures the overall quality or correctness of the reasoning chain at the terminal node, and C(a) represents the cost of concluding the search. The reward R(sa) can be derived from the commonly used outcome reward model (ORM), and we can even compare the final answer to the ground truth, as the reward is only used during the training of the PGTS policy.
The cost function C(a) plays pivotal role in guiding the PGTS policy by assigning penalties to actions based on their computational or logical expense. These costs act as mechanism to balance efficiency and reasoning quality, Actions such as expand and branch typically incur lower costs, as they change the reasoning tree locally. In contrast, the backtrack actions involve higher cost, discouraging ex- cessive reversions unless the potential improvements justify the expense. The terminate action, which concludes the rea- soning process and avoids further resource usage, generally carries negligible or zero cost. In our implementation, C(a) is treated as a hyperparameter that can be tuned to align the system's 's behavior with task-specific requirements or com- putational constraints. By appropriately setting these costs, the system fosters a balance between computational effi- ciency and reasoning complexity. This cost structure works in tandem with the reward function, which prioritizes accu- rate, coherent, and relevant reasoning paths while penalizing unproductive exploration. Together, the rewards and costs enable the PGTS policy to navigate reasoning trees effec- tively, balaneing exploration and exploitation, and adapting fiexibly to diverse tasks and constraints.
The reward function [s, a) assigns a scalar reward for executing action in state s. For each action type, the reward is defined as follows:
Let Sy E S be the partial reasoning path at the current node and 0r E A be the new reasoning step generated by the expand action, where d denotes the depth of the current node. The reward is defined as
Expand:
R(s, a) R(sa, aa) C(a),
(2)
where R(sa, aa) evaluates the quality of the new reason- ing step, and C(a) denotes the cost of the expand action.
When the policy selects the branch action to explore an alternative sibling node, let sa and s's denote the partial reasoning paths corresponding to the current node and its sibling node respectively, with shared parent sa-1. the reward is defined as
Branch:
i(s a) R(5/4-1, -( (3)
Here, ((54-11;'-11) ((sa-1,a4-1) 1) quantifies the po- tential improvement in reasoning by exploring the alter- native path, while C(a) reflects the cost of branching.
2.3.3. CONSTRAINED ACTION SAMPLING
The reasoning tree enforces explicit depth and breadth lim- its to maintain computational feasibility, which naturally constrains the validity of actions in different states. For expand actions, validity is determined by the depth limit and answer state: the expand action becomes invalid when the
Backtrack: For the backtrack action, which reverts to a previously visited node, all intermediate rewards along the backtracked path are revoked. Suppose the backtrack
4

---

Policy Guided Tree Search for Enhanced LLM Reasoning
current node has reached the maximum depth limit or when final answer has already been generated at the current node. Branch and backtrack actions face similar structural constraints. These actions are prohibited at the root node as it has no parent or siblings. They are also invalid when the parent node has reached its maximum allowed children, pre- venting unbounded tree growth. For backtrack specifically, a node at depth d can only backtrack up to d - steps, en- suring the action remains within the explored portion of the tree. The terminate action maintains the simplest constraint, becoming valid only when the current node has reached a final answer or the maximum depth limit has been reached.
while global attention helps maintain coherence across the entire reasoning process. We further enhance the struc- tural understanding by incorporating random-walking struc- ture embeddings (Dwivedi et al., 2021), which encode each node's position and connectivity within the tree.
The policy network processes these inputs through a se- quence of GPS layers. each layer aggregating both local and global information. The final representation of the current node is concatenated with the action constraints vector to ensure awareness of valid actions. This combined represen- tation then passes through linear layers to produce logits for a categorical distribution over the D +-2 possible actions. In- valid actions are masked out during sampling, ensuring the policy's 's decisions respect the tree's structural constraints.
These structural constraints ensure the reasoning process remains well-defined and computationally tractable while guiding the policy toward meaningful exploration paths. To implement these constraints efficiently, we represent them as a D + 2 dimensional binary vector, where D indicates the depth limit of the reasoning tree. Each dimension corre- sponds to the validity of a specific action: the first dimension for expand, the second for branch, the next D - dimen- sions for different backtrack depths, and the final dimension for terminate. This constraint vector serves dual purposes: it augments the policy's 's input state to inform action selec- tion, and it masks invalid actions from the policy's s output distribution, ensuring only valid actions can be sampled. Through this mechanism, the policy naturally learns to navi- gate within the fcasible action space while maintaining the structural integrity of the reasoning tree.
We implement the value network using the same architec- tural backbone, sharing the GPS layers with the policy net- work to maintain consistent representation learning. The value network differs only in its final layers, which produce a scalar estimate of the expected cumulative reward. This shared structure allows both networks to leverage the same learned representations of the reasoning tree while serving their distinct purposes in the decision-making process.
2.3.5, TRAINING
The training process for the PGTS policy aims to enhance reasoning effectiveness while minimizing unnecessary ex- ploration. This is accomplished by optimizing the policy through reinforcement learning. with rewards designed to promote both accuracy and efficiency. The training itera- tively refines the policy's ability to navigate reasoning trees and select actions that yield high-quality solutions.
2.3.4. POLICY DESIGN FOR PGTS
The effectiveness of PGTS heavily depends its ability to process and navigate reasoning tree. In our TS-MDP formu- lation, the state is naturally represented as a tree where each node corresponds to partial reasoning path. To effectively learn from this structured data, we need a policy architecture that can capture both local relationships between reasoning steps and global patterns across the entire tree.
We adopt Proximal Policy Optimization (PPO) (Schulman et al., 2017) as our training algorithm due to its stabil- ity and sample efficiency. Starting from randomly initial- ized weights, the policy interacts with reasoning tasks in episodes, learning to select actions that maximize cumula- tive rewards. When sampling actions during training, ap- ply the constraint mask to zero out logits of invalid actions, ensuring the decisions respect the tree's structural limits. The remaining logits are normalized to form valid cate- gorical distribution over permitted actions. This mechanism allows the policy to learn feasible exploration strategies while maintaining the integrity of the reasoning tree.
We design our policy using a graph-based architecture that processes two types of features. Node features are derived from the hidden states of the target L.M,, capturing the semantic content of each reasoning step. Edge features encode the intermediate rewards of reasoning transitions, quantifying the quality and relationships between consecu- tive steps. This rich feature representation enables the policy to assess both the local quality of individual steps and their contribution to the overall reasoning path.
To encourage efficient exploration, we incorporate entropy regularization into the policy loss. This ensures that the policy maintains balance between exploiting known high- reward paths and exploring less certain but potentially re- warding alternatives. Additionally, the cost components C(a) in the reward function are tuned to discourage exces- sive backtracking or branching, guiding the policy toward concise, meaningful reasoning paths. Please refer Algo- rithm for the complete training procedure.
To process these features effectively, we employ Graph Transformer-based model using the GPS framework (Rampäšek et al., 2022). This architecture combines lo- cal message passing operations from graph neural networks with global attention mechanisms from transformers. The lo- cal operations capture step-by-step reasoning relationships,
5

---

Policy Guided Tree Search for Enhanced LLM Reasoning
3. Related Works
graduate-level multiple-choice questions in O-shot setting. Finally, for the planning task, we evaluate our framework the Blocksworld benchmark (Valmeekam et al., 2022).
LLM Reasoning LLM reasoning has advanced through techniques such as CoT (Wei et al., 2022), ToT (Yao et al., 2024), and programmatic reasoning paradigms (Chen et al., 2022; Sel et al., 2023), fostering structured and iterative problem-solving. Recent innovations include heuristic search methods like MCTS (Feng et al., 2023; Hao et al., 2023) and A" search (Wang et al., 2024a). Building on these developments, our PGTS framework integrates learned poli- cies to improve search efficiency and reasoning performance. For detailed review of related approaches and their con- nection to inference-time scaling, please refer to Sec. B.
Across all datasets, we define single reasoning step as one sentence, maintaining consistency and simplicity. Un- like RAP (Hao et al., 2023), which incorporates a world model to simulate environment states after each action, our approach directly focuses on reasoning in the generated text without state modeling. For detailed dataset descriptions and reasoning setups, refer to Sec. C.
We compare PGTS against CoT and MCTS baselines. MCTS inherently explores multiple reasoning chains by traversing different paths during search, while we enhance CoT and our PGTS with self-consistency (SC) (Wang et al., 2022). Specifically, CoT aggregate outcomes from multiple chains using majority voting, whereas MCTS and PGTS utilize weighted voting based on the reward of each reason- ing chain. Additionally, we report MCTS results for the highest-reward trajectory. Since one of our primary goals is to reduce reasoning costs, we exclude MCTS approaches that rely on self-evaluation as intermediate rewards (Xie et al., 2024a; Gao et al., 2024), as they are computation- ally expensive. Instead, we simply use the likelihood of each reasoning step as the intermediate reward, R(s, a). An oracle setting is included for MCTS, allowing it to access task rewards by comparing generated answers with ground truths during search. For both MCTS and PGTS approaches, the tree breadth limited to 4 child node per parent. For detailed descriptions of the baselines, refer to Sec.D, and for ablations on the reasoning tree constraints, see Sec.4.1.
Graph Transformers Graph Transformers (GTss) have emerged as powerful architectures for processing graph- structured data, building upon the success of Transformers in other domains. These models are particularly attractive due to their ability to address fundamental limitations of tra- ditional Message Passing Neural Networks (MPNNs), such as over-smoothing and over-squashing issues (Alon & Ya- hav, 2020; Topping al., 2021). Various GT architectures have been proposed, from the initial Fully-connected Graph Transformer with basic positional encodings (Dwivedi & Bresson, 2020), to more sophisticated designs like SAN with invariant eigenvector aggregation (Kreuzer et al., 2021), and Graphormer with fistance-based encodings (Ying et al., 2021). GraphTrans (Wu et al., 2021) introduces the first hybrid architecture, which combines local message passing with global attention mechanisms. GPS (Rampášek et al., 2022) systematically investigates and integrates different components of GT's, offering a modular and scalable frame- work. In this work, we implement the PGTS policy using GPS layers given its ability to effectively combine local and global information while maintaining linear complexity.
For PGTS, we train the policy using up to 1,000 examples from the training split of each dataset. This highlights the sample efficiency of our approach, as 1,000 examples suffice to learn an effective policy. For ablations on training settings, refer to Sec. 4.1. The policy architecture consists of two GPS layers followed by a single linear layer for action and value prediction. To simplify experiments, the action cost C(a) is fixed across datasets, with values of 0.1, 0.2, 0.5, and 0.0 for expand, branch, backtrack, and terminate actions, respectively. While dataset-specific tuning could further enhance performance, we leave this for future exploration. See Sec. E for training details of our PGTS policy.
4. Experiments
In this section, we showcase the flexibility and effective- ness of our PGTS framework across diverse problem do- mains, including mathematical reasoning, commonsense reasoning, logical reasoning, and real-world planning. For mathematical reasoning, evaluate our framework on the GSM8K (Cobbe et al., 2021), MATH500 (Hendrycks et al., 2021; Lightman et al., 2023), and AQUA (Ling et al., 2017) datasets, using 4-shot settings for GSM8K and MATH500,, and a 10-shot setting for AQUA. The in-context learning (ICL) examples are adapted from OpenCompass (Contrib- utors, 2023) for GSM8K and MATH500, and from LLM- Reasoner (Hao et al., 2024) for AQUA. For commonsense reasoning, we evaluate Strategy (Geva et al., 2021) in 5-shot setting, with ICL examples also adapted from Open- Compass. For logical reasoning, we evaluate the PrOntoQA (Saparov& & He, 2022) dataset for logical deduction in a 5-shot setting and the GPQA (Rein et al., 2023) dataset for
Table I presents evaluation results for the LLaMA3.1 and LLaMA3.1-70B models, both configured to generate reasoning steps with a temperature of 0.6 and top-p of 0.9. Across all datasets, PGTS consistently outperforms CoT, demonstrating the effectiveness of its learned policy in guid- ing reasoning steps. The inclusion of SC further boosts PGTS performance. For instance, on the MATH dataset, PGTS improves accuracy from CoT's 34.40% 41.40% in the 8B setting, illustrating its ability to explore high-quality reasoning chains. Similarly, on Blocksworld, PGTS-SC
6

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Table Evaluation results of ..a.M3.1 1-8B and I. various datasets across multiple reasoning tasks: Mathematical reasoning (GSM8K, MATH, AQUA), Commonsense reasoning (StrategyQA), Logical reasoning (PrOntoQA, GPQA), and Planning tasks (Blocksworld with and steps). SC4 and SC8 denote colficonsst voting over 4 8 sampled chains, respectively. MCTS (Best)
reports results for the reasoning chain with the highest reward, while MCTS (Agg) presents results aggregated over multiple reasoning chains using weighted voting on final answer. MCTS (Oracle) compares final answer with groundtruth as an additional reward.
Mathematical
Com. Sense
Logical
Planning
GSM8K
MATH
AQUA
StrategyQA
PrOntoQA
GPQA
BW (4)
BW (8)
CoT
83.47
34.40
51.57
74.20
69.40
14.65
22.37
2.10
CoT (SC4)
87.79
42.20
53.94
77.10
73.00
15.66
22.37
2.79
CoT (SC8)
89.84
48.20
55.12
77.20
74.60
15.15
26.32
2.79
MCTS (Best)
86.13
43.80
60.63
79.00
74.20
34.34
28.95
6.29
LLaMA3.1-8B
MCTS (Agg)
87.72
46.00
59.45
79,50
74.20
32,83
28.95
6.29
MCTS (Oracle)
88.78
51.40
64.96
84.40
74.80
34.85
34.21
6.29
PGTS
85,29
41.00
60.63
77.70
68.20
18.69
27.63
3,50
PGTS (SC4)
89.61
47.00
66.93
83.80
74.40
22.73
35.53
4.90
PGTS (SC8)
89.99
52.20
66.54
85.70
77.40
27.78
36.84
6.99
CoT
91.66
53,80
72.83
83.60
92.00
20.20
39.47
18,88
CoT (SC4)
92.49
58.60
72.44
85.10
93.60
19.78
50.00
20.28
MCTS (Best)
91.28
56,20
76.38
85.70
95.20
41.92
50.00
22.38
LLaMA3.1-70B
MCTS (Agg)
92.27
57.20
77.17
86.70
95.20
43.94
50.00
22.37
MCTS (Oracle)
93.10
66.00
82.68
90.70
96.00
47.47
61.80
24.47
PGTS
91.05
54.77
73.23
85.50
91.40
32.18
46.05
22.38
PGTS (SC4)
92.54
59.65
79.92
91.00
96.80
36.36
50.00
25.87
achieves an accuracy of 36.84% compared to CoT-SC's 26.32%, showcasing its superior handling of multi-step rea- soning tasks. While both CoT and PGTS benefit from SC, PGTS exhibits more pronounced improvements, underscor- ing its ability to generate diverse, high-quality chains suit- able for aggregation. Please see Sec. G for examples of the generated reasoning chains.
SC requires generating additional tokens for aggregating out- comes over multiple reasoning chains, these operations are highly parallelizable, ensuring that PGTS with SC remains computationally feasible despite the additional generations. Notably, the overhead introduced by the PGTS policy is negligible compared to the cost of LLM generation, thanks to its lightweight architecture.
When compared to MCTS, PGTS demonstrates competitive performance across most tasks. For example, PGTS without SC achieves comparable accuracy to MCTS on GSM8K, MATH, and StrategyQA and surpasses MCTS when aug- mented with SC. However, one exception the GPQA task, where PGTS underperforms MCTS. This can be attributed to the inherent complexity of GPQA, which includes di- verse topics and is curated to challenge even non-expert humans. Moreover, the limited training data available for PGTS, given the task complexity, makes it difficult to fully capture the intricacies of GPQA.
4.1. Ablations
Training Examples In the main results, we train the PGTS policy using up to 1,000 examples. Here, we demon- strate the convergence behavior of the training process. Fig- ure 3 presents the training curve with evaluation results at intermediate checkpoints, showing that the policy converges quickly and plateaus at approximately 1,000 examples.
Tree Constraints In main results, limit the tree breadth to 4 child nodes per parent. Table 2 shows evaluation results on AQUA with varying tree breadths. Both MCTS and PGTS achieve better performance with broader trees, as expected, since a broader tree facilitates more diverse reasoning chains. However, broader trees also generate more tokens, so we use breadth of 4 in our main results to balance accuracy and reasoning cost. For tree depth, set upper bound that ensures all examples can reach the terminal state (i.e., the final answer), with the terminate action allowing early stopping as needed.
key advantage of PGTS over MCTS is its computational efficiency. As shown in Figure 2, MCTS incurs significantly higher token costs due to its exhaustive search over rea- soning chains. For instance, 16.25 times more tokens than CoT, whereas PGTS achieves competitive performance with only a 5.28 times increase. Similarly. on GSM8K, MCTS uses 13.33 times the tokens of CoT, while PGTS requires just 1.29 times. These reductions in token usage make PGTS more practical for real-world ap- plications where computational costs critical. Although
MATH, MCTS requires
Policy Network We implement our PGTS policy using GPS layers, which combine local message passing and
7

---

Policy Guided Tree Search for Enhanced LLM Reasoning
20.42
Tokan Usage (relative CoT
20.0
CoT
16.25
MCTS
15.0
13.33
12.68
10.0
9.13
7.5
50
5.28
4.84
25
1.001.29
1.00
1.00
208
1.001.15
1.00
1.281992
1.00
1.001
2.47
GSMBK
MATH
AQUA
StrategyQA
PrÖntoQA
GPQA
Blocksworld
Figure Comparison of generated token counts for LLaMA3.1 --B, normalized relative to the CoT method.
Table Ablation studies on policy implementation.
trajectory reward
83.93
Lrajection reward 2
www.mma
AQUA
GSM8K
85 29 tb. /5
60.63 60.89
84.77 76
GPS
60.63
85.29
GPS (w/o edge features)
50.79
81.96
54.72
Evaluction
Exaluction result
GPS (w/o global attention)
54.72
84.15
GPS (w/o local MPNN)
58.76
84.78
1000 2000 3000 4003
1000 2000 3000
training examples
training examples
SAN SLM
60.24
82.56
53.94
82.87
(a) GSM8K
(b) AQUA
LLM Agent
55.51
83.32
Figure 3. Trajectory reward along training, with evaluation results at intermediate checkpoints.
advanced graph modeling techniques could further enhance results. SLM and LLM Agent perform worse, underscoring the strengths of graph- approaches over language-only methods. We hypothesize that techniques such improved prompt engineering memory-augmente agents (Chen et al., 2024b) could boost agent-based performance, which we leave for future work.
Table AQUA results with different tree breadth.
2
4
Evaluation Results
PGTS MCTS
55.91 57.08
60.63
61.47
59.45
62.60
Generated Tokens
PGTS
252. 16
283.24
370.54
MCTS
399.13
1241.37
3282.41
5. Conclusion
global attention extract node features effectively. The reasoning tree also incorporates the immediate reward R(s,a) a) edge features to inform decision-making. Ta- ble 3 presents ablation studies evaluating the importance of each component in the bese-ased policy and com- paring it to alternative policy implementations. SAN re- places GPS layers with a different graph transformer ar- chitecture proposed in (Kreuzer et al., 2021). SLM re- places the graph-based based policy with small language model (distilbers settt-bbe--uncsseeee sed), which processes the previous reasoning trajectory to predict the next action. LLM Agent prompts the same target LLM to predict the exploration action, aligning with an agentic approach where the LLM assumes the role of an autonomous agent collab- orating with others to solve the problem. Details of these policy implementations are provided in Sec. The abla- tion results demonstrate that the full GPS policy achieves the best performance, highlighting the effectiveness of in- tegrating local message passing, global attention, and edge features. Removing edge features or global attention re- sults in significant performance drops, especially on AQUA, emphasizing their importance for reasoning. While SAN performs competitively, it falls short of GPS, suggesting that
In this work, we introduced Policy-Guided Tree Search (PGTS), novel framework for reasoning with large lan- guage models that combines the efficiency of policy-guide exploration with the structured advantages of tree search. PGTS dynamically allocates inference resources, prioritiz- ing promising reasoning paths for targeted exploration, lead- ing to significant improvements inference efficiency and the ability to tackle complex reasoning tasks. Moreover, PGTS addresses the "overthinking" problem commonly ob- served in many ol-like models (Chen al., 2024c), where excessive reasoning steps generated for simple problems.
A key paradigm shift in PGTS is treating the LLM as an environment rather than policy, enabling external decision- making components to guide reasoning processes more ef- fectively. While the current implementation uses simple log- likelihood-baat rewards as intermediate feedback, future work could incorporate nisms, such as L.M self-evaluation or task-specific metrics, to provide richer guidance during the search and further en- hance reasoning performance. PGTS represents a significant step toward more efficient and structured inference-time rea- soning with L.MMs.
sophisticated reward mecha-
8

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Impact Statement
Chen, L., Davis, J. Q., Hanin, B., Bailis, P., Stoica, I,, Za- haria, M., and Zou, J. Are more IIm calls all you need? towards scaling laws of compound inference systems. arXiv preprint Drxiv::4403002411 2024a.
This work aims to advance the field of reasoning with large language models (LLMs) by introducing the Policy-Guie Guided Tree Search (PGTS) framework. By improving inference efficiency and addressing challenges such as overthinking in reasoning processes, PGTS contributes to the broader objective of enhancing LLM usability in complex reasoning tasks.
Chen, M., Li, Y., Yang, Y., Yu, S., Lin, B., and He, X. Au- tomanual: Generating instruction manuals by Il agents via interactive environmental learning. arXiv preprint arXiv:2405. 16247, 2024b.
The proposed framework has the potential to positively im- pact areas where efficient reasoning and speiiiio-mmking are critical, including education, scientific discovery, and Al-driven assistance. However, it is essential to acknowl- edge certain ethical considerations and limitations. PGTS does not explicitly address the faithfulness of generated rea- soning chains, meaning the reasoning paths may not align with hwwannunderstaeeeeeeeeeeaa understandable logic or the true underlying rea- soning process. This limitation could inadvertently lead to misleading outputs in high-stakes applications.
Chen, W., Ma, X., Wang, X., and Cohen, W. W. Program of thoughts prompting: Disentangling computation from reasoning for numerical reasoning tasks. arXiv preprint arXiv:221 125588, 2022.
Chen, X., Xu, J., Liang, T., He, Z., Pang, J., Yu, D., Song, L., Liu, Q. Zhou, M., Zhang, Z., et al. Do not think that much for 2+ 3=? the overthinking of ol-like llms. arXiv preprint arXiv:2412.2 187, 2024c.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv: 10. 14 14168, 2021.
Future work exploring richer reward signals, such self- evaluation or task-specific metrics, could mitigate these concerns and ensure greater alignment with human expec- tations. While PGTS optimizes reasoning efficiency, care must also be taken to ensure it is applied responsibly, partic- ularly in applications requiring transparency and trust in AI reasoning processes.
Contributors, O. Opencompass: universal evaluation platform for foundation models. attps://gittet com/ open- httpc/ccmmpaa/pppccpppsppssss wpwp/cpppscccspppcpeo///// 2023.
In summary, this work represents meaningful step toward improving LLM reasoning capabilities, with both ethical implications and societal benefits that warrant further explo- ration and thoughtful application.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Dwivedi, V. P. and Bresson, X. of transformer networks to graphs. arXiv preprint arXhv:2012.09699, 2020.
generalization
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I,, Aleman, F, L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., al. Gpt-4 technical report. arXiv preprint 3XXiv23303.08774 2023.
Dwivedi, V. P., Luu, A. T., Laurent, T., Bengio, Y., and Bresson, X. Graph neural networks with learnable structural and positional representations. arXiv preprint 10.0v:21 2021.
Alon, U. and Yahav, E. On the bottleneck of graph neural networks and its practical implications. arXiv preprint 20Xi::2006.05205. 2020.
Feng, X., Wan, Z., Wen, M., McAleer, S. M., Wen, Y., Zhang, W., and Wang. J. Alphazero-like tree-search can guide large language model decoding and training. arXiv preprint arXiv:2309. 17179, 2023.
Besta, M., Blach, N.. Kubicek, A., Gerstenberger, R., Pod- stawski, M., Gianinazzi, L.., Gajda, J., Lehmann, T., Niewiadomski, H., Nyczyk, P., et al. Graph of thoughts: Solving elaborate problems with large language models. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 38, pp. 17682-17690, 2024.
179,
Gao, 7., Niu, B., He, X., Xu, H., Liu, H., Liu, A., Hu, X., and Wen, L. Interpretable contrastive monte carlo tree search reasoning. arXiv preprint arXiv:2410.01707, 2024.
Geva, M., Khashabi, D., Segal, E., Khot, T., Roth, D., and Berant, J. Did aristotle use a laptop? a question answering benchmark with implicit reasoning strategies. Transac- tions of the Association for Computational Linguistics, 9: 346-361, 2021.
Brown, B., Juravsky, J., Ehrlich, R., Clark, R., Le, V. Ré, C., and Mirhoseini, A. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXiv:2407.2178/ 2024.

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Grill, J.-B., Altché, F., Tang. Y., Hubert, T., Valko, M., Antonoglou, I,, and Munos, R. Monte-carlo tree search as regularized policy optimization. In International Con- ference on Machine Learning, pp. 3769-3778. PMLR, 2020.
can't plan, but can help planning in Imm-modulo frame- works. arXiv preprint arXiv:2402.01817 2024.
Kao, K.-C., Wang, R., and Hsieh, C.-J. Solving for x and beyond: Can large language models solve complex math problems with more-than-two unknowns? arXiv preprini arXiv:2407.05134 2024.
Hao, S., Gu, Y., Ma, H., Hong, J., Wang, Z., Wang. D. Z and Hu, Z. Reasoning with language model is planning with world model. arXiv preprint arXiv:2305. 14992, 2023.
Kartal, B., Hernandez-Leal P., and Taylor, M. Action
guidance with mets for deep reinforcement learning. In Proceedings of the AAAI conference artificial intelli- gence and interactive digital entertainment, volume 15, pp. 2519.
Hao, S., Gu, Y., Luo, H., Liu, T., Shao, X., Wang, X., Xie, S., Ma, H., Samavedhi, A., Gao, Q. et al. Lim reasoners: New evaluation, library, and analysis of step- by-step reasoning with large language models. arXiv preprint arXiv:2404.052211 2024.
Khalifa, M., Logeswaran, L., Lee, M., Lee, H., and Wang. Grace: ceeriminnato-gguidde chain-of- of-thhuthtttt reasoning. arXiv preprint arXiv:2305.1495: 2023.
Hart, P. E., Nilsson, N. J., and Raphael, B. A formal basis for the heuristic determination of minimum cost paths. IEEE transactions on Systems Science and Cybernetics, 1(2):100-107, 1968.
Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are zero-shot reasoners. Ad- vances in neural information processing systems, 35: 22199-22213, 2022.
Havrilla, A., Raparthy, S., Nalmpantis, C., Dwivedi- Yu, J. Zhuravinskyi, M., Hambro, E., and Raileanu, R. Glore: When, where, and how to improve Ilm reason- ing via global and local refinements. arXiv preprint arXiv:2402. 10963, 2024.
Kreuzer, D., Beaini, D., Hamilton, W., Létourneau, V., and Tossou, P. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 11-21618-21622 2021.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring math- ematical problem solving with the math dataset. arXiv preprint arXivv2103.03877 2021.
Li, C., Liang, J., Zeng, A., Chen, X., Hausman, K., Sadigh, D., Levine, S., Fei-Fei, L., Xia, F., and Ichter, B. Chain code: Reasoning with a language soodellaugmenter code emulator. arXiv preprint arXiv:2312.00442 22..4444444444 2023.
Hong, R., Zhang, H., Pang, X., Yu, D., and Zhang, C. A closer look at the self-verification abilities of large language models in logical reasoning. arXiv preprint arXiv:2311.07955 2023.
Li, Y., Lin, Z., Zhang, S., Fu, Q. Chen, B., Lou, J.- G., and Chen, W. Making large language models bet- ter reasoners with step-aware verifier. arXiv preprint arXiv:2206.02336 2022.
Huang, J. and Chang, K. C.-C. Towards reasoning in large language models: A survey. arXiv preprint arXiv:2212. 2022.
Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let's verify step by step. arXiv preprint 23x5v:2305.20050 2023.
Huang. J., Chen, X., Mishra, S., Zheng, H. S., Yu, A. W., Song, X., and Zhou, D. Large language mod- els cannot self-correct reasoning yet. arXiv preprint arXiv:2310.01798, 2023.
Ling, W., Yogatama, D., Dyer, C., and Blunsom, P. Pro- gram induction by rationale generation: Learning to solve and explain algebraic word problems. arXiv preprint arXiv: / 705.04146,66 2017.
Jiang, A. Q. Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. I., Bressand, F, Lengyel, G., Lample, G., Saulnier, L., et al. Mistral 7b, arXiv preprint arXiv:2310.06835, 2023.
Malherbe, C., Grosnit, A., Tutunov, R., Bou Ammar, H., and Wang, Optimistic tree searches for combinatorial black-box optimization. Advances in Neural Information Processing Systems, 35:33080-33092. 2022.
Kambhampati, S. Can large language models reason and plan? Annals of the New York Academy Sciences, 1534 (1):15-18, 2024.
Ning, X., Lin, Z., Zhou, Z., Wang, Z., Yang. H., and Wang, Y. Skeleton-of -f-thuughtt Prompting Ilms for efficient par- allel generation. In The Tweifth International Conference on Learning Representations, 2024.
Kambhampati, S., Valmeekam, K., Guan, L., Verma, M., Stechly, K., Bhambri, S., Saldyt, L., and Murthy, A. Lims
10

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Qi. Z., Ma, M., Xu, J., Zhang, L. L., Yang, F,, and Yang, M. Mutual reasoning makes smaller Ilms stronger problem- solvers. arXiv preprint arXiv:2408.06199 2024.
Topping, J., Di Giovanni, F., Chamberlain, B. P., Dong, X., and Bronstein, M. M. Understanding over-squashing and bottlenecks on graphs via curvature. arXiv preprint arXiv:2111.14522 2021.
Qu, Y., Zhang, T., Garg, N., and Kumar, A. Recursive introspection: Teaching language model agents how to self-improve. arXiv preprint vrXi:2247.182214 2024.
Uesato, J., Kushman, N., Kumar, R., Song, F., Siegel, N., Wang, I.., Creswell, A., Irving. G., and Higgins, I. Solv- ing math word problems with process-and outcome-based feedback. arXiv preprint arXiv:2211. 2022.
Rampāšek, L.. Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D. Recipe for general, powerful, seal- able graph transformer. Advances in Neural Information Processing Systems, 35:14501-14515, 2022.
Valmeekam, K. Olmo, A., Sreedharan, S., and Kambham- pati, S. Large language models still can't 't plan (a bench- mark for Ilms on planning and reasoning about change). NeurIPS 2022 Foundation Models for Decision Mak- ing Workshop, 2022.
Rein, D., Hou, B. L., Stickland, A. C.. Petty, J., Pang, R. Y., Dirami, J., Michael, J., and Bowman, S. R. Gpqa: graduate-level google-proof q&a benchmark. arXiv preprint arxiv:2311.10 X1v:2311.11022, 12022, 2023.
Wang. C., Deng, Y., Lyu, Z., Zeng, L., He. J., Yan, S., and An, B. Q*: Improving multi-step reasoning for ilms with deliberative planning. arXiv preprint arXiv:2406. 2024a.
Saparov, A. and He, H. Language models are greedy rea- soners: A systematic formal analysis of chain-of-thought. arXiv preprint arXiv:2210.01220 2022.
Wang, L., Fonseca, R., and Tian, Y. Learning search space partition for black-box optimization using monte carlo tree search. Advances in Neural Information Processing Systems, 33:1951 1-99222, 2020.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint ftt:::1707.63444 2017.
Sel, B., Al-Tawaha, A., Khattar, V., Jia, R., and Jin, M. Algorithm of thoughts: Enhancing exploration of ideas in large language models. arXiv preprint arXiv:2308. 10379, 2023.
Wang. P., Li, L., Shao, Z., Xu, R., Dai, D., Li, Y., Chen, D., Wu, Y., and Sui, Z. Math-shepherd: Verify and reinforce Ilms step-by-step without human annotations. In Proceed- ings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 9426 9439, 2024b.
Silver, D., Hubert, T., Schrittwieser, J., Antonoglou, I., Lai, M., Guez, A., Lanctot, M., Sifre, I.., Kumaran, D., Grae- pel, T., al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv 2017.
Wang. S., Xue, K., Song, L.,, Huang, X., and Qian, C. Monte carlo tree search based space transfer for black-box opti- mization. arXiv preprint arXiv:2412.07186 2024c.
preprint arXiv:
Singh, A., Co-Reyes, J. D., Agarwal, R., Anand, A., Patil, P., Garcia, X., Liu, P. J., Harrison, J., Lee, J., Xu, K., et al. Beyond human data: Scaling self-training for problem-solving with language models. arXiv preprint arXiv:2312.065885 2023.
Wang. X., Wei, J., Schuurmans, D., Le. Q., Chi, E.,, Narang, S., Chowdhery, A., and Zhou, D. Self-consistency im- proves chain of thought reasoning in language models. arXiv preprint arXiv:2203. //71, 2022.
Wei, J., Wang. X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24873 2022.
Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling Ilm test- time compute optimally can be more effective than scal- ing model parameters. arXiv preprini 2400:22008..3314, 2024.
Stechly, K., Valmeekam, K., and Kambhampati, S. On the self-verification limitations of large language mod- els reasoning and planning tasks. arXiv preprint arXiv:2402.08115. 2024.
Wu, Y., Sun, Z., Li, S., Welleck, S., and Yang. Y. An empir- ical analysis of compute-optime inference for problem- solving with language models. 2024.
Wu, Z., Jain, P., Wright, M., Mirhoseini, A., Gonzalez, J. E., and Stoica, I. Representing long-range context for graph neural networks with global attention. Advances in Neural Information Processing Systems, 34:1 3266- 13279, 2021.
Team, G., Anil, R., Borgeaud, S., Alayrac, J.-B., Yu, J.. Sori- cut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., et al. Gemini: a family of highly capable multimodal 2023.
models. arXiv preprint arXiv:2312.1

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Xie, Y., Goyal, A., Zheng, W., Kan, M.-Y., Lillierap, T. P., Kawaguchi, K., and Shieh, M. Monte carlo tree search boosts reasoning via iterative preference learning. arXiv preprint MXXv:22005000451 2024a.
open reasoning models for open-ended solutions. arXiv preprint arXiv:241..44405, 2024.
Zhou, D., Schärli, N., Hou, L., Wei, J., Scales, N., Wang, X., Schuurmans, D., Cui, C., Bousquet, O., Le, Q., et al. Least-to-most prompting enables complex reasoning in large language models. arXiv preprini arXiv:2205. 10625, 2022.
Xie, Y., Kawaguchi, K., Zhao, Y., Zhao, J. X., Kan, M.- Y., He, J., and Xie, M. Self-evaluation guided beam search for reasoning. Advances in Neural Information Processing Systems, 36, 2024b.
Yang. L., Yu, Z., Zhang, T., Cao, S., Xu, M., Zhang, W., Gonzalez, J. E., and Cui, B. Buffer of thoughts: Thought- augmented reasoning with large language models. arXiv preprint ar arXi::2406.044 04271, 2024.
Yao, S., Yu, D., Zhao, J., Shafran, I,, Griffiths, T., Cao, Y., and Narasimhan, K. Tree of thoughts: Deliberate problem solving with large language models. Advances Neural Information Processing Systems, 36, 2024.
Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877-28888. 2021.
Yuan, 7., Yuan, H., Li, C., Dong, G., Lu, K., Tan, C., Zhou, C., and Zhou, J. Scaling relationship on learning math- ematical reasoning with large language models. arXiv preprint X1X.:2330...8255 2023.
Zelikman, E.. Wu, Y., Mu, J., and Goodman, N. Star: Boot- strapping reasoning with reasoning, Advances in Neural Information Processing Systems, 25:15476-154888 2022.
Zelikman, E., Harik, G., Shao, Y., Jayasiri, V., Haber, N., and Goodman, N. D. Quiet-star: Language models teach themselves to think before speaking. arXiv preprint arXiv:2403.09629, 2024.
Zhai, Y. and Gao, S. Monte carlo tree descent for black-box optimization. Advances in Neural Information Processing Systems, 35:12581-12592, 2022.
Zhang, L., Hosseini, A., Bansal, H., Kazemi, M., Kumar, A., and Agarwal, R. Generative verifiers: Reward mod- cling as next-token prediction, 2024. URL https://axxiv. org/abs/2408.. /5200, 2024.
Zhang, Z., Zhang, A., Li, M., and Smola, A. Automatic chain of thought prompting in large language models. 2022.
arXiv preprint 2r.iv:221000344
Zhao, R., Li, X., Joty, S., Qin, C., and Bing, L. Verify- and-edit: A knowledge-enhance chain-of-thougee frame- work. arXiv preprint 2305v:2305.088:
2023.
Zhao, Y., Yin. H., Zeng. B., Wang, H., Shi, T., Lyu, C., Wang, I.., Luo, W. and Zhang, K. Marco-ol: Towards
12

---

Policy Guided Tree Search for Enhanced LLM Reasoning
A. Training Algorithm
Algorithm Policy Training in TS-MDP
Require: Reasoning tasks T, depth limit D, breadth limit B, reward function R[s, a), cost function C(a) Trained policy π (a)s) and value function ((s)
Ensure:
1: Initialize: Parameters of policy T and value function Vi
2: while not converged do
3: Sample a task from T, and initialize the root node an
4: for 0, J..... ... max do
5: Extract valid actions and construct the constraints vector
6: Compute policy logits π TT(
7: Sample action a, Categorical(Tt)
8:
Execute action and observe next state 8141, reward rr, and constraints C++1
9: Store transition (ss, ag, T2, S1+1, C++1) in replay buffer B
10: a, is Terminate then
Break
end if end for
13:
14:
Compute Returns and Advantages:
15:
for each sampled trajectory in B do
16:
Compute discounted returns: G2 = n Lijo Pretts
17:
Compute advantages: A, = G. - (s((
end for
Policy Update:
20:
Compute policy loss Ln -B log Th (ar)BL c.) A.] w H(To(asss,,..)))
21:
Update policy parameters: 2 - Φ - 1/ L.
22:
Value Function Update:
23:
Compute value loss Lv = IE [V[(V((s) - G.)
24:
Update value parameters: vp + e MJV Lvv
end while
B. Related Works
LLM Reasoning In the context of LLM reasoning, Chain-of-Thought (CoT) resoning (Wei et al., 2022) serves foundational technique that decomposes problems into intermediate steps, facilitating step-by-step reasoning that mimics human groblem-solving strategies. Since its introduction, numerous enhancements to CoT have been proposed, including Zero-Shot-Co'T (Kojima et al., 2022), Self-Consistency with CoT (Wang et al., 2022), AutoCoT (Zhang et al., 2022), and VerifyCoT (Zhao et al., 2023), among others. Least-to-Most Prompting (Zhou et al., 2022) represents another influential paradigm that iteratively constructs solutions by dividing complex problems into simpler tasks. Building on this problem decomposition paradigm, Program-of-Thoughh (Chen al., 2022), Chain-of-Code (Li al., 2023), Buffer-of-Though (Yang et al., 2024), and algorith--of-Thougee (Sel et al., 2023) incorporate programmatic reasoning steps further enhance problem-solving. Free-of-Thoughe (Yao et al., 2024) extends CoT by organizing reasoning processes into tree structure, where each node represents partial solution and edges signify transitions. However, ToT faces challenges with large search spaces for deep trees, leading to computational inefficiencies. To address this, Graph-of-Thoue (Besta et al., 2024) generalizes ToT by modeling reasoning as a graph, allowing dynamic path selection, backtracking, and aggregation of information across multiple paths. Meanwhile, Sheleton-of-Thhugh (Ning et al., 2024) reduces latency by generating a skeleton outline of answers before completing the details in parallel. Recent research has also explored heuristic search techniques, such Monte Carlo Tree Search (MCTS) (Feng et al., 2023; Hao al., 2023) and A Search (Wang et al., 2024a), to identify optimal reasoning chains. Building on these advancements, our proposed PGTS approach enhances heuristic search by integrating learned policy to guide the search process, enabling more efficient exploration.
Beyond directly eliciting reasoning capabilities from LLMs, external components have been developed to verify reasoning steps (Cobbe et al., 2021: Lightman et al., 2023; Li et al., 2022; Wang et al., 2024b; Zhang et al., 2024) or refine generated reasoning through iterative revision (Qu al., 2024; Havrilla et al., 2024). Additionally, iterative bootstrapping techniques
13

---

Policy Guided Tree Search for Enhanced LLM Reasoning
(Zelikman et al., 2022; 2024; Singh et al., 2023; Yuan al., 2023) have been explored to progressively enhance LLM reasoning performance. These approaches are complementary to our method and can be integrated achieve further improvements.
Inference Time Scaling Optimizing inference-time compute has emerged as a critical area of research for improving the performance of LLMss. Unlike the traditional focus on scaling model parameters during training, inference-time scaling explores strategies to leverage additional compute during inference to enhance model outputs. One foundational approach in inference-ting time scaling best-of-N sampling, where multiple completions are generated for a given prompt, and the best response is selected using a secondary mechanism, such as a reward model (Lightman et al., 2023). Recent studies have sought to formalize the relationship between inference compute and performance, proposing scaling laws for inference-time optimization (Snell et al., 2024; Wu et al., 2024; Chen et al., 2024a; Brown et al., 2024). These works show that carefully allocating inference resources, such as by balancing the number of samples with selection quality, can yield significant gains in model performance. Our proposed PGTS approach offers an efficient alternative to traditional inference- -ime scaling methods. Rather than passively sampling multiple outputs, the policy actively and dynamically decides when and where to allocate additional inference resources. This targeted approach enables more strategic exploration, focusing compute on the most promising reasoning paths.
C. Evaluation Datasets
GSM8K GSM8K (Cobbe et al., 2021) is benchmark dataset designed to assess mathematical reasoning skills in language models. It comprises approximately 8,500 high-quality grade school math problems that require multi-step reasoning to arrive at the correct solution. We adhere to the original train-test split and evaluate models in a 4-shot setting, using in-context examples and a prompt template adapted from OpenCompass (Contributors, 2023), which provide representative coverage of various problem types.
During generation, each reasoning step is defined as single sentence. To manage computational complexity, we set the tree breadth limit to 4, meaning each parent node can have up to 4 child nodes, and the depth limit to 16, which suffices to reach terminal nodes for most questions in this dataset. Additionally, the PGTS policy is configured to allow a maximum of 64 reasoning steps, balancing efficiency with the ability to explore complex reasoning chains.
MATH500 MATH500 (Hendrycks et al., 2021; Lightman et al., 2023) is subset of the MATH dataset, designed to evaluate advanced mathematical problem-solving capabilities. It encompasses questions from high school and undergraduate-Ievel topics, including calculus, geometry, linear algebra, and number theory. We use the original training split to train our PGTS policy and evaluate the framework in a 4-shot setting with in-context examples sourced from OpenCompass.
Each reasoning step defined as a single sentence or a single line. To control computational complexity, we set the tree breadth limit to 4, the depth limit to 64, and the maximum reasoning steps for PGTS to 256, ensuring sufficient capacity for exploring complex mathematical reasoning chains.
AQUA AQUA (Ling et al., 2017) is a dataset comprising algebraic word problems and reasoning- based multiple-choice questions. Its diverse question formats and logical reasoning challenges make an effective benchmark for assessing models' general problem-solving abilities. We adhere to the original train-test split and evaluate our framework using a 10-shot setting, with in-context examples sourced from LLM-Reasoner (Hao al., 2024), enabling the model address problems involving numerical reasoning and reasoning by elimination.
Each reasoning step is defined as a single sentence. To manage computational complexity, we set the tree breadth limit to 4, the depth limit to 16, and the maximum reasoning steps for PGTS to 64, ensuring adequate exploration capacity for solving the dataset's challenges.
StrategyQA StrategyQA (Geva et al., 2021) is commonsense reasoning dataset designed to evaluate models' ability to answer binary (yes/no) questions by applying implicit multi-step reasoning strategies. The dataset features diverse and challenging questions where the intermediate reasoning steps are not explicitly provided. We randomly sample 1,000 examples for testing, with the remaining examples used to train the PGTS policy. Evaluation is conducted in a 5-shot setting, leveraging examples from OpenCompass to emphasize the model's capacity to infer and synthesize information across multiple reasoning steps.
14

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Each reasoning step is defined single sentence. To balance computational efficiency and exploration, we set the tree breadth limit to 4, the depth limit 16, and the maximum reasoning steps for PGTS to 64, ensuring sufficient capacity to handle the dataset's complexity.
PrOntoQA PrOntoQA (Saparov & He, 2022) is benchmark designed to evaluate logical reasoning and deductive inference. It consists of structured questions that require precise logical deductions based on predefined rules and premises. We adopt the same dataset splits used in RAP (Hao et al., 2023) and evaluate the model in 5-shot setting to assess its ability to perform deductive reasoning while adhering to strict logical constraints. This dataset serves as a valuable tool for testing structured reasoning in a controlled environment.
Each reasoning step is defined as a single sentence. To balance computational efficiency and thorough exploration, we set the tree breadth limit 4, the depth limit 16, and the maximum reasoning steps for PGTS to 64, ensuring sufficient capacity to tackle the dataset's 's logical challenges.
Blocksworld Blocksworld (Valmeekam et al., 2022) is benchmark for evaluating real-world planning tasks, requiring models reason over sequences of actions achieve specific goals in simulated block environment. The dataset assesses the ability to plan and execute multi-step strategies efficiently. We evaluate two versions that require 4 and 8 reasoning steps, using the same dataset splits and prompt template as RAP (Hao et al., 2023). Specifically, split-v1 is used for training and split-v2 for testing, with 4-shot setup. In the easy setup, the in-context examples include the same number of reasoning steps as the test problems.
Unlike RAP, which employs world model to predict the state after each action, we directly generate entire action trajectories without modeling intermediate states. Each reasoning step corresponds to one action described in single sentence. We set the breadth limit to 4 and the depth limit match the required number of reasoning steps for each version. The maximum reasoning steps for PGTS are capped at 32.
D. Baseline Methods
Chain-of-Thoughe CoT is a reasoning strategy that decomposes complex tasks into a sequence of intermediate steps, allowing models reason step-by-step towards a final answer. To enhance the robustness of CoT, we incorporate self- consistency (SC) (Wang et al., 2022), where multiple reasoning chains are generated for the same input, and the outcomes aggregated via majority voting. This voting scheme helps mitigate errors by leveraging diverse reasoning paths.
Monte Carlo Tree Search MCTS has been employed to improve LLM reasoning (Xie al., 2024a; Gao al., 2024), and our implementation closely follows RAP (Hao et al., 2023), with key modifications. Specifically, omit the use of world model for predicting state changes and rely solely on the likelihood of each generated reasoning step as the intermediate reward, instead of LLM--ased self- eeff-evaluation. Additionally, we introduce an oracle setting for MCTS, granting access to task rewards by comparing generated answers with ground truths during the search, providing an upper-bound performance benchmark.
MCTS inherently explores multiple reasoning chains by expanding nodes during the search process. We evaluate two result aggregation strategies: selecting the final answer from the chain with the highest trajectory reward and performing weighted voting over final answers based on their respective rewards. To ensure computational efficiency, we limit the tree's breadth to 4, consistent with our PGTS approach.
E. PGTS Approach
The PGTS policy employs a Graph Neural Network (GNN) architecture to dynamically model the reasoning tree as the search progresses. Specifically,, we utilize GPS layers (Rampäšek et al., 2022), which integrate local message passing with global attention mechanisms. This design effectively captures both fine- fine-grained structural relationships and broader contextual information within the reasoning process. The policy network consists of two GPS layers for feature aggregation, followed by linear layer applied to the current node's features to predict both actions and values. The action prediction is modeled as categorical distribution parameterized by the policy network, while the value prediction outputs a continuous scalar.
The input to the policy network is the reasoning tree constructed up to the current step. Each node the tree represents a
15

---

Policy Guided Tree Search for Enhanced LLM Reasoning
partial reasoning step, with its features initialized using the embedding of the corresponding reasoning text. Edge features encode the immediate reward R(s, a), which captures the contribution of each reasoning step the overall trajectory.
Training the policy leverages Proximal Policy Optimization (PPO) (Schulman et al., 2017), reinforcement learning algorithm designed to balance exploration and exploitation while maintaining stable updates. The training objective maximizes the expected cumulative reward across reasoning paths, with constraints to ensure updates remain within a stable region. Key PPO hyperparameters include:
PPO Clip Range: 0.2
Discount Factor (7): 0.99
GAE Lambda (A): 0.95
During training, the current policy to simulate the reasoning trajectories and collect the partial reasoning trees. These trajectories are collected and used to optimize the policy based on PPO objectives, iteratively improving the efficiency and accuracy of reasoning over time.
F. Alternative Policy Implementations
SLM The Small Language Model (SLM) implementation substitutes the graph- policy with a language-moder driven approach. We employ disti www.bberrbbbeeuccceeee trajectory. The input to the model is the serialized reasoning chain corresponding the current node. This serialized trajectory is passed through the model, which generates both the predicted next action and the value estimation using linear output layer. During training, the backbone language model remains frozen, and only the output layer parameters are optimized. The optimization follows the same PPO objectives as used PGTS.
lightweight transformer model, to process the reasoning
LLM Agent The LLM Agent represents an agentic approach where the same target LLM used for reasoning is also leveraged to predict the next exploration action. At each step, the reasoning trajectory provided as prompt to the LLM, along with instructions to suggest the next action. Please see Fig. 4 for the prompt template to select the exploration action. This approach aligns conceptually with collaborative multi-agent framework where the LLM functions both generator communicate with user and planner to reason about the optimal action. However, the LLM Agent lacks explicit structural modeling of the reasoning process and relies heavily prompt engineering guide its predictions. While this approach demonstrates reasonable performance, it is limited by the inherent inefficiencies of treating the LLM as both the reasoning environment and the policy, leading to increased computational costs and suboptimal exploration complex scenarios.
G. Examples
Figure 5, 6, 7, and 8 illustrate examples of the PGTS reasoning process applied to problems from GSM8K, MATH, PrOntoQA, and StrategyQA, respectively. The left side displays the reasoning tree, while the right side outlines the detailed reasoning steps.
16

---

Policy Guided Tree Search for Enhanced LLM Reasoning
You are an Al assistant designed to navigate through a reasoning graph to solve problems. At each node in the graph, you can take one of these actions:
1. Expand: Dive deeper from current node to explore its subnodes
2. Branch: Explore alternative parallel paths by investigating sibling nodes
3. Backtrack: Return to previous node to explore different possibilities
4. Terminate: End the reasoning process when you't 've reached a satisfactory answer
Important notes:
Sibling nodes are generated dynamically when you choose explore them
Your goal is find the most logical path the correct answer
Consider each step carefully before choosing action
Please choose only from the available actions list given below
Choose 'Branch' when current node seem incorrect or inefficient towards a solution
Choose 'Backtrack' to revise a previous node that may lead to incorrect or inefficient solution
When Backtrack' to a previous node in the current path, specify its node id the format 'Node Id: NODE.ID*
Choose "Terminate" only when you'le 're confident in the complete solution
# QUESTION:
Here is the problem to be solved:
{question]
CURRENT REASONING PATH:
Node IDs and reasoning steps:
{reasoning path}
# AVAILABLE ACTIONS:
Due to graph constraints in depth and breadth, not all actions are available to choose at each step. The following actions are currently available:
{available_actions)
Instructions:
1. Analyze where you in solving the problem
2. Review which actions are available to you
3. Decide your next action based on:
the current path promising?
Are there better approaches to try?
Have you reached a complete solution?
You can only choose one action from the available actions above
5. Format your final action choice <answer> YOUR CHOSEN ACTION< and optionally provide
your thinking process.
should be one of the available actions.
Figure 4. Prompt template to select the optimal exploration action for LLM Agent based policy.
17

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Question: Janet's ducks lay 16 day. She eats three for breakfast every momiin and bakes
her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh every day farmers'
Root
duck egg.
dollars market?VnLat's think step by stepinAnswer:
Expand
Root
Node Janet's ducks lay 16 day.
1
Expand
Janet eats 3 eggs for breakfast, she 16-3-13 eggs left.
Node
2
Expand
She bekes for muffins, she has 13-4=9 eggs left.
Node
3
Expand
She sells remainder the farmers' market, which eggs.
5
Node
Branch
Janet sells remaining eggs the farmers" market for $2 each.
Expand
makee day farmers' market.
Figure 5. An example of the PGTS reasoning process applied to problem from the GSM8K dataset. At node 4, the policy decide to branch explore alternative reasoning path.
18

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Convert the point (0,3) in rectangular coordinates to polar coordinates.
Root
Enter your answer the form (r, 8), where and0 I Please step and put your final answer within (boxed]]
Expand
Root
To convert the point (0,3) from rectangular coordinates to polar coordinates, need find
Node
angle B.
Expand
radius given by
4 y², Plugging
Node
2
get
VOp 3 V
Expand
3
5
The angle can be found using the formla However, since 0, we need to be cáreful.
Node
Backtrack step
6
We have that
Node
+y² -==+3=
Expand
We also have that
Nodie
Expand
Node
Since tan
2
have that
Expand
Node
the coordinates (344)
Figure 6. An example of the PGTS reasoning process applied problem from the MATH dataset. At node 3, the policy backtracks 1 step to explore alternative reasoning path.
19

---

Policy Guided
Tree Search for Enhanced LLM Reasoning
Q: Negative numbers negative, Prime numbers not composite. Every natural number negative. Every prime number natural number. Real numbers not imaginary. Each real number is number. integers.
Root
number. Every integer is real number. Every Mersenne prime is not composite. Mersenne True false: 131071 negative.
Mersenne
131071
A:
Expand
Root
Node
131071 is Mersenne prime.
Expand
Every Mersenne prime is prime number. 131071 prime number.
Node
2
Expand
Node
numbers not composite. So 131071
is not composite.
3
6
Expand
4
7
Node
Every prime number netural number. So 131071 a natural number.
Expand
5
8
Node
Every natural number is negative. 131071 not negative.
Backtrack steps
Node
number natural number. 131071 is a natural number.
Expand
Node
Every natural number negative, 131071 is not negative.
Expand
Node 8
The false.
Figure 7. An example of the PGTS reasoning process applied to problem from the PrOntoQA dataset. Al node the policy backtracks steps to explore an alternatiee reasoning path.
20

---

Policy Guided Tree Search for Enhanced LLM Reasoning
Question: is sunscreen unhelpful for the condition that killed Bob Martey?
Root
Answer:
Expand
Root
Bob Marley died from acral lentiginous melanoma,, type
Node
Expand
1
Sunscreen helps prevent skin cancer by protecting against
Node
2
Expand
Thus, sunscreen would be helpful preventing the condition that kiled Bob Marley.
Node 3
3
Expand
4
Node
the
Backtrack 2 steps
Sunscreen can help protect against skin cancer, Including melanoma.
Node
Expand
However, acral lentiginous melanoma is a type that the palms the hands, soles the feet, or under the nails.
Expand
Sunscreen may not be as effective in these due the thickness
Node 7
Expand
Thus, sunscreen may be less helpful for the specific condition that killed Bob Marley.
Node 8
Expand
Node
So the answer
Figure 8. An example of the PGTS reasoning process applied problem from the Strategy dataset. Although node 4 has reached answer. the policy chooses to explore alternative reasoning paths to refine and rectify the answer.
21


Reviews:
[{"review_id": "SWmhZBhTx0", "reviewer": "Reviewer_GYUV", "summary": "This paper introduces Policy-Guided Tree Search (PGTS), a novel framework that combines reinforcement learning with structured tree search to enhance reasoning capabilities in Large Language Models (LLMs). The key innovation is a learned policy that dynamically decides between expanding the current reasoning path, branching to alternative paths, backtracking to previous states, or terminating exploration.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful and constructive feedback. We address the key points raised below:\n\n## PGTS Performance on Complex Tasks like GPQA\n\nThe reviewer correctly notes that PGTS underperforms MCTS on GPQA (27.78% vs. 34.85%). This performance gap can be attributed to three primary factors:\n\n1. 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.\n2. 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.\n3. 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.\n\nWe 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:\n\n* Developing more sophisticated reward mechanisms tailored to complex reasoning tasks\n* Exploring synthetic data generation to improve training diversity and coverage\n\n## Training with More Examples\n\nThe 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.\n\nGPQA 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.\n\nTo address this challenge in future work, we plan to explore synthetic data generation techniques to expand training diversity.\n\n## Action Definitions Beyond Sentence Splitting\n\nThe 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. \n\nWe 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.\n\n## Clarifications on Algorithm 1\n\nThank you for noting the undefined symbols in Algorithm 1. We will revise the algorithm description to clearly define:\n\n- H: The entropy of the policy outputs, which is used as regularization to encourage exploration.\n- $\\omega$: The entropy regularization coefficient (set to 0.01 in our experiments)\n- $\\eta_{\\pi}$ and $\\eta_{V}$: Learning rates for the policy and value networks (both set to 3e-4)\n\nWe will include these definitions in the final version.\n\n## Ablation on Cost Function Settings\n\nWe 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. \n\nWe'd be happy to include a more detailed analysis of cost function sensitivity in the final version.\n\nWe appreciate the reviewer's thoughtful comments and will incorporate these suggestions in the final version.", "claims_and_evidence": "The paper's claims about PGTS's effectiveness are generally well-supported by empirical evidence across multiple reasoning tasks. The authors present comprehensive experimental results comparing PGTS against Chain-of-Thought (CoT) and Monte Carlo Tree Search (MCTS) baselines, showing consistent performance improvements while reducing token usage.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in this paper are generally appropriate and well-aligned with the reasoning enhancement problem the authors aim to solve. The evaluation strategy spans diverse reasoning tasks, including mathematical reasoning (GSM8K, MATH, AQUA), commonsense reasoning (StrategyQA), logical reasoning (PrOntoQA, GPQA), and planning (Blocksworld), which provides a comprehensive assessment of the framework's capabilities across different problem domains.", "theoretical_claims": "The paper does not contain formal mathematical proofs for theoretical claims that require rigorous verification. The primary theoretical contribution is the formulation of the reasoning process as a Tree Search MDP (TS-MDP), which is presented in Section 2.3.2.", "experimental_designs_or_analyses": "I have carefully examined the experimental designs and analyses presented in the paper and find them generally sound, though with some limitations.\n1. The paper limits tree breadth to 4 for computational reasons, which may artificially constrain the exploration space and affect comparative performance.\n2. The training process for the policy network could be more transparent about potential overfitting issues in Figure 3.\n3. There is limited analysis of how performance scales with the number of training examples beyond Figure 3.", "relation_to_broader_scientific_literature": "The key contributions of PGTS relate to several important research directions in the broader scientific literature.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "**Strengths**\n1. The paper creatively combines reinforcement learning with tree search to produce a dynamic, adaptive approach for LLM reasoning that avoids both the rigidity of predefined heuristics and the computational expense of exhaustive search.\n2. PGTS demonstrates substantial efficiency improvements over MCTS approaches, requiring significantly fewer tokens while achieving comparable or better performance.\n3. The experiments span diverse reasoning tasks, from mathematical problem-solving to logical deduction and planning, demonstrating the framework's versatility.\n\n**Weaknesses**\n1. There is no ablation analysis on the cost function C(a) settings, leaving unanswered questions about how sensitive the approach is to these hyperparameters across different tasks.\n2. Value modeling in complex reasoning processes remains challenging, and the paper provides limited evidence that two GPS layers can adequately predict both actions and values.\n3. The paper doesn't fully explore whether using higher-level actions (as shown in the LLM Agent approach in Figure 4) could be more effective than the sentence-level splitting used in the main experiments.", "other_comments_or_suggestions": "Algorithm 1 is missing explanations for several symbols:\n- H: Likely represents the horizon or maximum trajectory length, but should be explicitly defined\n- ω: Appears to be the entropy regularization coefficient in line 20, but lacks definition\n- ηπ and ηV: Presumably learning rates for policy and value networks, respectively, but require clarification", "questions_for_authors": "1. Table 1 shows that PGTS significantly underperforms MCTS on complex tasks like GPQA (e.g., 27.78% vs. 34.85% accuracy for LLaMA3.1-8B). Given that one major advantage of state-of-the-art LLMs like o1-class models is their performance on precisely these difficult reasoning tasks, does this indicate a fundamental limitation of PGTS in enhancing model performance on the most challenging reasoning problems? If so, what modifications might address this limitation?\n2. Have you explored training with substantially more examples, particularly for complex tasks like GPQA where performance lags behind MCTS?\n3. How does PGTS compare to MCTS in terms of sample efficiency when both are trained on identical data? The current comparison focuses on inference efficiency, but a direct comparison of how effectively each method utilizes the same training examples would provide important context. \n4. The paper uses sentence-level segmentation for reasoning steps, while Figure 4 shows the LLM Agent approach using higher-level structured actions. Have you investigated whether using more semantically meaningful action definitions beyond simple sentence splitting could improve performance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2ZULNmDzcS", "reviewer": "Reviewer_LMkW", "summary": "The paper introduces Policy-Guided Tree Search (PGTS), a novel framework designed to acquire the reasoning capabilities of LLMs by combining reinforcement learning with structured tree exploration. The key innovation of PGTS is a learned policy that dynamically decides between four actions: expanding, branching, backtracking, or terminating exploration. This approach eliminates the need for predefined heuristics or exhaustive search, making it more efficient and scalable compared to existing methods like Chain-of-Thought (CoT) prompting and Monte Carlo Tree Search (MCTS). The key empirical claim is that PGTS across various reasoning tasks, including mathematical reasoning, logical deduction, and planning benchmarks, shows significant improvements in accuracy while reducing computational costs.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful and constructive feedback. We address the key points raised below:\n\n## Fair Comparison of Computational Budget\n\nThe reviewer raises an excellent point about fair comparisons between methods.\n\nHowever, on most datasets, the token usage ratio is much more favorable, and PGTS often outperforms even CoT with self-consistency:\n\n* 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)\n* On StrategyQA: PGTS achieves 77.70% vs CoT (SC4)'s 77.10% while using only 1.58× tokens compared to CoT\n* On GPQA: PGTS achieves 18.69% vs CoT (SC4)'s 15.66% with modest token increase\n* On Blocksworld (4): PGTS achieves 27.63% vs CoT (SC4)'s 22.37% with efficient token usage\n\nMore 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.\n\nPGTS effectively occupies a middle ground in the accuracy-efficiency tradeoff space:\n\n1. CoT methods: Lowest token usage but limited exploration capability\n2. PGTS: Moderate token usage with adaptive exploration capability\n3. MCTS: Highest token usage with exhaustive exploration capability\n\nFor 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.\n\nIn 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.\n\n## Technical Details and Algorithm Clarity\n\nThe 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:\n- Policy network: Two GPS layers with 256 hidden units each (~1M parameters)\n- Input: Tree structure with node features (768d hidden states from the LLM's final layer) and edge features (1d immediate rewards)\n- Output: Categorical distribution over D+2 actions (expand, branch, D-1 backtrack depths, terminate)\n\nFor 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.\n\nWe 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.\n\n## Reward Signal for Termination\n\nFor 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.\n\n## Wall-time Efficiency\n\nLLM 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.\n\n## Fast Convergence and Broader Applicability\n\nPGTS 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.\n\nWe appreciate the reviewer's encouraging words about the research direction and will address these concerns in the revised version to improve clarity and rigor.", "claims_and_evidence": "##### Claim 1 - Policy-Guided Tree Search (PGTS)**:  \n  - The paper introduces a novel framework that combines **reinforcement learning (RL)** with **structured tree exploration** to enhance the reasoning capabilities of large language models (LLMs).  Its core innovation is a **learned policy** that dynamically decides between four actions: **expand**, **branch**, **backtrack**, or **terminate** exploration. This eliminates the need for **predefined heuristics** or **exhaustive search** strategies.  \n\nI think this claim is well supported in the experiments and literature review.\n##### Claim 2 - Superior Reasoning Accuracy**:  \n  - PGTS achieves **state-of-the-art performance** across multiple reasoning benchmarks, including **mathematical reasoning**, **logical deduction**, and **planning tasks**.  \n  - On the **MATH dataset**, PGTS achieves **41.00% accuracy** with LLaMA3.1-8B, compared to **34.40%** for Chain-of-Thought (CoT). When augmented with **self-consistency (SC)**, PGTS-SC8 achieves **52.20%**, outperforming CoT-SC8 (48.20%).  \n  - Similar improvements are reported on other datasets, such as **GSM8K**, **AQUA**, **StrategyQA**, and **Blocksworld**.  \n\nI am very concerned about the meaning of this claim. The authors are clearly aware of the necessity of taking the computation budget into account. This is, however not done in a systematic way. \n\nFor simplicity let me concentrated on the statement from the introduction: \"Using LLaMA3.1-8B, PGTS achieves a 41.00% accuracy on MATH, significantly improving upon CoT’s 34.40% while using only one-third of the tokens required by MCTS.\"\n\n\nThe full results extracted from Table 1 are\n\n| Method | SC1   | SC4   | SC8   |\n| ------ | ----- | ----- | ----- |\n| CoT    | 34.40 | 42.20 | 48.20 |\n| PGTS   | 41.00 | 47.00 | 52.20 |\n\nHowever, PGTS used $5\\times$ more tokens than CoT SC. It is thus more justified to compare PGTS (SC1) = $41\\%$ with Cot (SC4) = $42\\%$. In this light is it disputable if PGTS indeed better. \n\nThe authors should be precise about what comparisons they view as apple-to-apple, and why.", "methods_and_evaluation_criteria": "The proposed method is sound. The evaluation criteria and benchmark used are appropriate, except from the computational budget issue stated in the previous section.", "theoretical_claims": "Not applicable for this work.", "experimental_designs_or_analyses": "Experimental design is overall correct, there are diverse benchmarks. Baselines are valid (though I'd love to see at least one more). However there are various issues as well:\n- lack of proper apple to apple comparison\n- lack of statistical analysis, e.g., confidence intervals\n- lack of hyperparamter tuning (e.g., I am not completely sure if the temperature setting is optimal for the CoT baseline)", "relation_to_broader_scientific_literature": "The literature is properly discussed.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "First of all, I am very excited about the presented idea. I strongly encourage the authors to pursue this direction. \n\nUnfortunately, I think that the current version lack quality needed for a publication in a A* conference. Here is a list of issues\n\n1. The presentation is somewhat chaotic. While textual parts are too long in my opinion, on the other hand I am missing a lot of  technical details (e.g. the number of the policy network, what is exactly the input, including dimensionality etc.)\n\t1. I find it hard to grasp all details of the algorithm. Having pseudocode would be of help!\n\t2. I am very confused what reward signal is used for termination eq. (5)? \n\t3. Also, the authors could comment on wall-time efficiency. CoT+SC is very efficient in this respect, using modern frameworks like vLLM one can achieve high MFU, which translates to attractive wall-time (dollar value). Do we expect the same from the presented methods.\n2. The method use additional training. It should be explained how does it affect comparisons. \n3. I'd love to see more extensive analysis of the policy. Does using only two layers implies that it is 'local' (taking into account only a small ball of the tree?)\n\t1. Following that, is it possible to interpret the policy or even extract it?\n4. I am not quite sure how the depth constraint work, but it resembles somewhat beam-search. Is it a relevant comparison. If so, I'd consider including this as a baseline.", "other_comments_or_suggestions": "none", "questions_for_authors": "Can this method applied outside of LLMs? Can you state assumptions when it can be applied.\n- I find it fascinating that the convergence is so fast (Fig 3). Can you comment what is that so?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g6IBwh45Ou", "reviewer": "Reviewer_jFF8", "summary": "This paper presents a novel algorithm that integrates **tree search** and **reinforcement learning** which can be used alongside an **off-the-shelf LLM** to improve reasoning capabilities. Unlike conventional approaches that focus on **LLM fine-tuning**, the authors propose training a **Graph Neural Network (GNN)** to serve as a policy during inference, reducing the complexity of learning such a policy.  \n\nThe proposed method is **simple, scalable**, and supported by **comprehensive experiments** that effectively demonstrate its advantages.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive assessment of our work and the constructive feedback. We address the key points below:\n\n## Missing References\n\nWe 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:\n\n- **STAR** bootstraps reasoning by selecting and refining high-quality reasoning chains, focusing on improving reasoning through iterative training.\n- **REST-MCTS*** combines tree search with self-training, using process rewards to guide exploration.\n\nOur 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.\n\n## Action Importance\n\nRegarding $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.\n\nThe 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.\n\nExamples 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).\n\n## Policy Function Compared to PRM\n\nThe reviewer correctly identifies that our policy functions similarly to a Process Reward Model (PRM) at inference time. However, there are important distinctions:\n\n1. **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.\n\n2. **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.\n\nThis design allows PGTS to more efficiently guide the reasoning process while maintaining the flexibility to explore diverse paths as needed.\n\n## Blocksworld Training Set Sampling\n\nRegarding the Blocksworld training data, we use their split-v1 for training and split-v2 for testing. For clarification:\n\n1. We used the original problem generator from Valmeekam et al. with the same configuration parameters.\n2. Following RAP's methodology, we generated distinct training (split-v1) and testing (split-v2) sets.\n\nThis approach ensures a fair comparison with existing methods while maintaining the challenging nature of the planning task.These details are already included in appendix.\n\n## Comparison with Other Training-Based Approaches\nWe 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:\n\n- 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.\n- Model Preservation: PGTS keeps the LLM parameters unchanged, maintaining its general capabilities while learning to guide its reasoning process.\n\nOur 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.\n\nWe appreciate the reviewer's valuable feedback and believe addressing these points will further strengthen our paper.", "claims_and_evidence": "Yes, the claims are clear and the evidence is provided via experiments and ablations.", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria make sense and are a good list.", "theoretical_claims": "Not applicable, since there are none.", "experimental_designs_or_analyses": "Yes, I checked all the experiments and found them to be sound with details shared about training and implementation.", "relation_to_broader_scientific_literature": "This paper relates to LLM reasoning, planning, and inference time search in general.", "essential_references_not_discussed": "I believe that the authors could have discussed the following papers:\n1. Zelikman, Eric, et al. \"Star: Bootstrapping reasoning with reasoning.\" Advances in Neural Information Processing Systems 35 (2022): 15476-15488.\n2. Zhang, Dan, et al. \"Rest-mcts*: Llm self-training via process reward guided tree search.\" Advances in Neural Information Processing Systems 37 (2024): 64735-64772.\n\nThe above two papers share similar ideas and a comparison against the two would have been helpful.", "other_strengths_and_weaknesses": "**Strengths**  \n1. The paper is well-structured and clearly presented, with intuitive ideas.  \n2. The method is demonstrated across a diverse set of tasks, highlighting its broad applicability.  \n3. The approach achieves the most significant improvements in cases where the base LLM performs poorly, which is a promising outcome.  \n\n **Weaknesses**  \n1. The proposed method is **training-based**, yet the paper does not compare it against other training-based approaches. Such comparisons would strengthen the evaluation.  \n2. For tasks where the **LLM already performs well**, the performance gains are less pronounced.", "other_comments_or_suggestions": "None. Overall, this paper offers valuable insights and would make a strong addition to **ICML 2025**. I would be open to increasing my score if my questions and concerns are adequately addressed.", "questions_for_authors": "1. It seems that **c(a)** is biased toward **expansion and termination**, and the appendix does not provide examples of branching. Could the authors conduct an ablation study to assess the importance of each action?  \n\n2. To the best of my understanding, the policy functions similarly to a **Process Reward Model (PRM)** at inference time, where the **LLM is guided by the learned policy and prompted to take actions accordingly**. Is this interpretation correct?  \n\n3. **BlocksWorld**, as introduced by **Valmeekam et al.**, does not include a predefined training set. Could the authors clarify their **strategy for sampling the training set** in this work?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1nLioxcJd9", "reviewer": "Reviewer_5qMX", "summary": "The paper introduces Policy-Guided Tree Search (PGTS), a framework that enhances large language models (LLMs) for complex reasoning tasks by integrating reinforcement learning with structured tree search. Unlike traditional approaches such as Chain-of-Thought (CoT) or Monte Carlo Tree Search (MCTS), PGTS employs a learned policy to dynamically decide between expanding, branching, backtracking, or terminating reasoning paths, reducing reliance on predefined heuristics and computationally expensive exploration. The reasoning process is formalized as a Markov Decision Process (MDP), optimized using Proximal Policy Optimization (PPO), with rewards balancing reasoning accuracy and efficiency. PGTS utilizes a Graph Transformer-based policy network to process reasoning trees, enabling adaptive navigation of complex reasoning steps. Extensive experiments across mathematical reasoning (GSM8K, MATH, AQUA), logical deduction (PrOntoQA, GPQA), commonsense reasoning (StrategyQA), and planning (Blocksworld) demonstrate that PGTS improves accuracy while using fewer tokens than MCTS. By intelligently guiding inference-time exploration, PGTS emerges as a scalable, efficient alternative for structured reasoning in LLMs.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n## On Quantifying \"Overthinking\" in LLMs\n\nThe reviewer rightly suggests quantitative metrics on reasoning step counts to support our claim about reducing overthinking. Our analysis shows that:\n\n- For StrategyQA, PGTS uses on average 5.97 reasoning steps per problem vs. MCTS's 50.38 and CoT's 4.75 steps.\n- For GSM8K problems, PGTS uses on average 10.73 reasoning steps per problem vs. MCTS's 91.97 and CoT's 7.35 steps.\n- For MATH problems, PGTS uses 61.64 steps on average vs. MCTS's 220.24 steps.\n- For Auqa problems, PGTS uses 17.46 reasoning steps vs. MCTS's 64.20 steps.\n\nImportantly, PGTS's \"terminate\" action effectively stops the reasoning process once sufficient evidence has been generated, reducing steps compared to MCTS while improving accuracy.\n\nThese metrics confirm that PGTS achieves better accuracy while using fewer reasoning steps, directly addressing the overthinking problem.\n\n## On Generalization Without Dataset-Specific Tuning\n\nWe'd like to clarify our approach to hyperparameter tuning. While we did train separate policies for each benchmark, we deliberately used the same:\n- Network architecture (2-layer GPS)\n- Action costs (0.1, 0.2, 0.5, 0.0 for expand, branch, backtrack, terminate)\n- PPO hyperparameters\n- State/action representations\n\nThis 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.\n\n## On Additional Baseline Comparisons\n\nWe 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:\n\n- ToT requires task-specific design for proposal and evaluation prompts, which is not a standardized process and introduces variability that would complicate fair comparisons.\n- 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.\n\nPGTS addresses these limitations by learning a task-agnostic policy that doesn't require manual prompt engineering or code generation capabilities.\n\n## On Qualitative Evaluation\n\nWe 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.\n\n## On Evaluation on Open-ended Tasks\n\nWe 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.\n\nFor 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.\n\n## Conclusion\n\nWe 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.", "claims_and_evidence": "1. The claims about PGTS improves accuracy and more efficient on token usage are well-supported by experiments in section 4. \n2. The claim about PGTS learns an adaptive policy which can dynamically select actions that improves reasoning quality is proved by the ablation in section 4.1. Removing key compoments like edge features and global attention degraded the performance. \n3. For the claim `PGTS reduces the ‘overthinking’ problem in LLMs.`, it'd be better if you can provide more quantitative metric, e.g. reasoning step counts in PGTS vs CoT/MCTS. \n4. For `PGTS can generalize across diverse reasoning tasks without dataset-specific tuning.`, from what I understand, PGTS still required dataset-specific tuning, at some level, similar to RAP requires task-specific reward design. More experiments to show that PGTS maintains effectiveness without any dataset-specific hyperparameter adjustments would be supportive.", "methods_and_evaluation_criteria": "Yes, I believe the method is well-motivated, though it shares similar ideas to previous works like RAP, LLM-MDP, it provides some novelty. \nThe evaluations are pretty solid to support the main claims on performance and token-level efficiency, compared against strong baselines. I do have some concern on the scale of the novelty, but the work itself is self-contained.", "theoretical_claims": "This work doesn't explicitly show any theoretical claims. But here are some I'd say can be helpful:\n1. For the tree search MDP and action constraints in section 2.3.2 and 2.3.3, we want to make sure these constrains won't affect the model from selecting an optimal reasoning path. The constrained action space (i.e., limiting backtracking depth, restricting branching) ensures computational feasibility and prevents infinite loops. The binary constraint vector encoding action validity is a reasonable way to enforce these rules. **There are no major theoretical flaws**, but the paper does not formally prove optimality guarantees.\n2. The PGTS policy is implemented as a Graph Transformer (GPS) model, allowing it to process tree-structured reasoning paths effectively. The use of local message passing and global attention aligns with existing graph neural network (GNN) theory. The paper does not provide a theoretical proof or discussion of why GPS is the best choice for this problem compared to other graph-based architectures. It'd be better to have more details on this, like the discussion in Appendix E&F.", "experimental_designs_or_analyses": "1. The benchmarks selected are quite diverse, from math problems to planning tasks like Blocksworld. \n2. The baselines are recent and competitive. There are some other strong baselines can be included in the experiments, e.g.\n    - Tree-of-Thought (ToT) (Yao et al., 2024) which also optimizes tree search in LLM reasoning.\n    - Algorithm-of-Thought (AoT) (Sel et al., 2023), a more structured variant of CoT.\n    - *(The authors comment the challenges in comparing the model against the baselines above)\n3. Evaluation metrics on performance and token efficiency are quite sufficient for current tasks. It'd be better to report other metrics targeting the overthinking claim, e.g. chain length vs correctness. (chain length is further provided in the rebuttal, which shows the effectiveness of PGTS)", "relation_to_broader_scientific_literature": "1. PGTS builds upon prior approaches aimed at enhancing reasoning in LLMs by structuring problem-solving into intermediate steps.\n    - Relation to Chain-of-Thought (CoT) Prompting [Wei et al., 2022]\n      - PGTS extends CoT by dynamically exploring different reasoning paths, rather than committing to a single linear sequence.\n      - PGTS does not rely on predefined prompts and instead learns an optimal reasoning policy.\n    - Relation to Tree-of-Thought (ToT) Reasoning [Yao et al., 2024], RAP [Hao et al., 2023]\n      - ToT/RAP often relies on heuristic or manually-defined exploration strategies, whereas PGTS learns an adaptive policy using reinforcement learning.\n    - Relation to Algorithm-of-Thought (AoT) [Sel et al., 2023]\n      - AoT prescribes structured reasoning rules, while PGTS learns them dynamically from data.\n2. PGTS contributes to the emerging field of LLM inference-time optimization, improving efficiency during reasoning.\n    - Relation to Self-Consistency in LLM Reasoning [Wang et al., 2022]\n      - Self-consistency (SC) improves reasoning by generating multiple solutions and selecting the most common answer. PGTS benefits from SC, but unlike SC, it **actively selects reasoning paths** during search rather than **aggregating** pre-generated answers.\n    - Relation to Skeleton-of-Thought (SoT) [Ning et al., 2024]\n      - SoT reduces inference time by first generating an outline, then filling in details in parallel. PGTS is different in that it does not perform parallel reasoning but instead optimizes search exploration, which helps to optimize reasoning structure rather than token generation.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "### Strengths:\n  - Combining Policy Learning with Tree Search:\n    - The key innovation of PGTS is integrating a learned policy with tree search for LLM reasoning.\n    - While Monte Carlo Tree Search (MCTS) has been explored for reasoning (e.g., RAP, ToT), PGTS removes the need for predefined/manually annotated heuristics by learning a policy to dynamically guide search.\n    - This reduces human-engineered reward functions, making PGTS more adaptable across tasks.\n  - Backtracking and Termination in LLM Reasoning:\n    - PGTS introduces structured backtracking to recover from incorrect paths, which is not standard in prior CoT-based methods.\n    - The terminate action ensures LLMs do not overthink, addressing a key issue in existing search-based methods.\n### Weaknesses:\n  - PGTS is mainly evaluated on well-defined and structured benchmarks. \n    - Evaluate PGTS on more commonsense reasoning tasks (e.g., ARC, BigBench) or real-world multi-hop QA datasets would be beneficial. \n  - Missing comparison with some baselines (e.g. ToT) as mentioned\n  - Qualitative evaluation on overthinking is missing.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Can PGTS generalize to open-ended reasoning tasks since it still requires dataset-specific tuning? It seems to be pretty solid on structural/well-defined tasks, but if tuning is always necessary, maybe PGTS can't scale well to unseen reasoning tasks. (clarified by the authors in the rebuttal)\n2. Do you have more detailed analysis on the reasoning chains? The accuracy of PGTS is sounded, but conducting human evaluations on the steps in the reasoning chain can further prove if the improvement does come from better reasoning quality over gaming the evaluation metrics.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper received four high-quality reviews, with overall scores of 4, 4, 2, and 4.

All reviewers recognize the innovative nature of the approach and its promising empirical results. Even the reviewer who provided a lower rating acknowledged the solid contribution of the work, though they expressed some concerns regarding clarity in presentation and the additional computational costs introduced during the training phase.

The authors' rebuttal adequately addressed most of the concerns raised. Notably, the critical reviewer affirmed that the presented results are indeed publishable following the authors' clarifications.

After evaluating the reviews and authors' responses, the Area Chair agrees that the paper clearly meets the acceptance criteria. However, it is suggested that further clarification and refinement of presentation be included in the final manuscript.
Decision: Accept (poster)