Learning to Better Search with Language Models via Guided Reinforced Self-Training
Authors: Seungyong Moon
Bumsoo Park
Hyun Oh Song
Venue: NeurIPS 2025 poster ()
Keywords: learning to search
self-improvement
reinforcement learning
Primary area: deep_learning
TL;DR: We propose a novel self-improvement algorithm to teach language models to perform effective search.

Abstract: While language models have shown remarkable performance across diverse tasks, they still encounter challenges in complex reasoning scenarios. Recent research suggests that language models trained on linearized search traces toward solutions, rather than solely on the final solutions, exhibit improved generalization, despite the search traces being potentially noisy or suboptimal. However, relying on such imperfect traces can result in inefficient use of test-time compute. To address this, we propose guided reinforced self-training (Guided-ReST), a fine-tuning algorithm designed to improve the model’s capability for effective search during inference. The key insight behind Guided-ReST is that optimal solutions can serve as valuable step-by-step landmarks to guide the model’s search process. Based on this insight, we introduce a novel data generation method that seamlessly incorporates optimal solutions into the model’s search procedure, enabling the generation of high-quality search traces. By fine-tuning the model on these search traces, we effectively distill improved search strategies into the model. Our method significantly enhances the search capabilities of language models on arithmetic reasoning and code self-repair tasks, including Countdown, CodeContests, and CodeForces. We release the source code at https://github.com/snu-mllab/guided-rest.

Learning to Better Search with Language Models via Guided Reinforced Self-Training
Seungyong Moon', Bumsso Park?, Hyun Oh Song!
Seoul National University, *KRAFTON (symoon11, hyyno1,,mynnhhmmm1111e snu. ac kr bumsoo, . park96kkrafton. com
Abstract
While language models have shown remarkable performance across diverse tasks, they still encounter challenges in complex reasoning scenarios. Recent research suggests that language models trained on linearized search traces toward solutions, rather than solely on the final solutions, exhibit improved generalization, despite the search traces being potentially noisy or suboptimal. However, relying such imperfect traces can result in inefficient use of test-time compute. To address this, we propose guided reinforced self-training (Guided-ResS) designed to improve model's s capability for effective search during inference. The key insight behind Guided-ReST" is that optimal solutions can serve as valuable step-by-step landmarks to guide the model's ' search process. Based on this insight, we introduce a novel data generation method that seamlessly incorporates optimal solutions into the model's search procedure, enabling the generation of high-quality search traces. By fine-tuning the model on these search traces, we effectively distill improved search strategies into the model. Our method significantly enhances the search capabilities of language models on arithmetic reasoning and code self relf-repair tasks, including Countdown, CodeContests, and CodeForces. We release the source code at battps://github wtw.ttt...oooooe.eeoooootttttttt snu- mwwab/ www.mm/gudeddrreeeeeeee guited-rest
tine-tuning algorithm
Introduction
Transformer-bant language models have achieved remarkable success, demonstrating human-level performance across a wide range of natural language tasks, including conversation, code generation, [S] 26). Their impressive performance is primarily attributed auto-regressive training on massive, internet-sourced corpora. However, language models still encounter challenges in tasks that require complex planning and reasoning [I8] 33
and mathematical problem-solving [ 31 23 12
To address
these challenges, recent approaches have leveraged prompt-based strategies, enabling self-correction and planning via external symbolic search algorithms [34 27 38). While these methods have shown success in specific contexts, their reliance on explicit prompting templates and manually engincered search structures restricts their flexibility and scalability.
More recent studies indicates that training language models to internalize the search process, rather than solely outputting the final solution, significantly improves their reasoning capabilities and allows performance to scale effectively with increased test-time compute 4 7 An instructive example of test-time scaling is stream of search (SoS) [4, which trains model to imitate search traces that encompass the complete decision- decision-making process finding optimal solutions through trial and error, including exploration and backtracking upon failure. This work demonstrates that models trained on search traces, despite being noisy and even suboptimal, exhibit superior generalization performance compared to those trained solely to imitate optimal solution via behavior cloning (BC) [22). However,
*Corresponding author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Reinforced self-training
Guided reinforced self-training
Stage 1
Stage 1
Stage 2
Stage N
Explore only
Language model
Explore
Explore
Language model
Jump
Inefficient test time-time
Jump
Explore
Efficient test time compute
Noisy search traces without any guidance
compute
Optimal search traces with guided exploration
O: Correct nodes
O: Incorrect nodes
Figure 1: Overview of guided reinforced self-training, illustrating how search traces are generated through the progressive integration of optimal solutions during self-generation. The numbers indicate the order in which nodes are explored.
excessively noisy search traces can result in inefficient utilization of test-time compute, as the model may repeatedly explore irrelevant paths or perform redundant backtracking.
In this work, we examine how the sub-optimality search traces affects test-time compute efficiency and propose novel self-training algorithm, guided reinforced self-training (Guided-ReST), designed to improve search efficiency during inference. Our key insight is that while optimal solutions alone insufficient for direct imitation, they can effectively serve as valuable landmarks to guide the search process. Inspired by jump-start reinforcement learning (JSRL) (32], we introduce new data generation approach that progressively integrates optimal solutions into the model's search process. generating high-quality search traces, as illustrated in Figure [ Subsequently, we fine-tune the model on these traces, thereby distilling more effective search strategies into the model. Finally, we further enhance performance by applying reinforcement learning (RL) fine-tuning.
We evaluate method on the Countdown benchmark 4 challenging arithmetic reasoning task that requires complex search. Our method achieves over a 10% accuracy improvement compared to baseline fine-tuning algorithms. Furthermore, utilizes test-time compute more efficiently, reaching comparable performance while using less than half the number of tokens required by other baselines. Finally, we demonstrate the applicability of our method to more realistic code self-repair task the CodeContests and CodeForces benchmarks [1 21
2 Preliminaries
2.1 Stream of search
In this paper, consider the problem of training language model my for tasks that require complex searching capabilities. Suppose that we have a training dataset 'D of question-solution pairs (g, S), where each solution consists of step-by-step reasoning S (81, ..... 8T). The most straightforward
approach to train the model on this dataset is behavior cloning (BC) [22] which optimizes the model to generate the optimal solution using supervised learning:
max () [log (S |4]
However, prior studies have demonstrated that this approach struggles to generalize to unseen test examples 37
39].
To address this, the steam of search (SoS) method introduces novel training approach that emphasizes the search process itself rather than just the final solution 4 Specifically, SoS formulates the problem as tree search, exploring potential solutions through trial and error with operations until a successful solution is discovered. Each node in the tree represents a partial solution, and each edge represents single reasoning step. The method represents primitive tree-search operations in language, including node generation, exploration, verification, and backtracking. For each question, it employs symbolic search algorithms such as depth-first search (DFS) and breadth-first search (BFS) to generate a search trace Z = (Zi,..., 2T") where each 2, denotes an individual search operation expressed in language.

---

The model is then trained to predict this search trace via supervised learning:
max (2 Dlog n ())
Note that in this framework, the optimal solution can also be viewed search trace, namely path the search tree. Additional details are provided in Appendix A.
2.2 Fine-tuning with self-generated data
While SoS demonstrates improved generalization by training language model on search traces rather than optimal solutions, these traces are often noisy and suboptimal, resulting imperfect alignment with the target task. To address this, SoS further fine-tunes the model using reinforcement learning techniques. One approach is reinforced self- self-training (ReST) [42 6 (88), which fine-tunes the model on successful examples drawn from its self ge-eeeeeeredd responses via supervised learning:
max D, Z-mo(-lo))) [1 R(Z())>> log To 9)]
where R denotes the terminal reward function that evaluates the quality of the responses and 7 denotes the threshold parameter.
Another approach is reinforcement learning (RL) fine-tuning 30 17 which fine-tunes the model to directly maximize the reward while constraining the KL divergence from the reference model Tef initialized from the pre-trained model:
max (. [R(Z] 9) - ( (g) || Teee(: g))],
where 3 > 0 denotes the coefficient controlling the strength of the KI. penalty term. This objective is typically optimized with proximal policy optimization (PPO) or group-relative policy optimization (GRPO) 25 26
2.3
Countdown benchmark
We use Countdown as the primary benchmark, following previous studies on searching with language models (38 Each problem consists of input numbers and target number, and the objective is to combine the inputs using the four basic arithmetic operations to reach the target. For example, given target number 26 and input numbers (84, 14, 15). valid solution is 26 84 x (14+ 15). Each arithmetic operation in this solution can be considered a single reasoning step. This task involves high branching factor up to O each node of the search tree, where & the number of remaining inputs. To ensure a tractable difficulty level, we set the number of initial inputs to 4, consistent with the setup in Gandhi al. [ Additional details are provided in Appendix B.
Methods
3.1
Motivation
SoS demonstrates that a language model trained on noisy, suboptimal search traces generalizes better than one trained only clean, optimal solutions [4]. However, optimal solutions still offer valuable guidance during search. By providing the model with partial optimal solutions as hints, we can steer the model toward reduced and more manageable search space, enabling more efficient exploration and significantly increasing the likelihood success under a given token budget. Although optimal solutions are available only during training. the resulting trajectories provide high-quality supervision for fine-tuning.
This approach is closely aligned with the principle of jump-start reinforcement learning (JSRL) 32 which introduces a two-policy framework comprising a guide policy and an exploration policy. The guide policy, informed by prior knowledge demonstrations, initiates the search process by steering the model toward promising regions of the solution space. The exploration policy then takes over and continues the search to reach the optimal solution, which significantly reduces the cost of exploration.
To investigate whether SoS can also benefit from this approach, we conduct a preliminary experiment. Given the optimal solution S = (S1,...,87) ST), we define a partial solution (s1,...,5t) 5t) consisting of
3

---

Step 1. Select random child node from the last subgoal node.
Step 2. Replace the selected node to the next subgoal node and truncate the search trace up to it.
Step 3. Continue the search from the modified node.
Inputs: 84,2,14,15 Tanget: Solution: 84 1114-15)=26
84,2,14,15
84,2,14,15
84,2,14,15
B4/14-6
B4/2-42
15-14-1
84/14-6
14+15-22
29
84/14-6
14+15-29
2,15,6
14,45,42
84,2,1
2,15,6
84,2,29
2,15,6
84,2,29
Language model
15-2=13
15-14=1
A2/14:3
64/2=42
15-2-13
15-2-13
84/2-12
29-2-27
6,13
15,3
1,42
6,13
6,33
84,27
12/1
42
12-1-411
15:3+5
15+3-18
42:29-13
42
41
18
O: Correct nodes
O: Incorrect nodes
O: Selected node
Figure 2: Overview of the subgoal augmentation algorithm. The numbers indicate the order in which nodes are explored. The red box highlights the modifications through node replacement.
Table 1: Accuracy and cross-entropy loss of search traces from partial solutions of different lengths.
Length
2
3
Accuracy (%) (t)
55.96
85.28
95.27
100.00
Cross-entropy loss (L)
0.0481
0.1245
0.2412
0.2539
the first reasoning steps. We utilize this partial solution as an initial prompt to condition the model in generating response. Table shows the accuracy of search traces initialized from partial solutions of different lengths, evaluated over 10K training examples. The model successfully continues the search and discovers solutions when starting from partial solutions, even though it has not encountered these specific partial solutions during training. Moreover, increasing the amount of guidance by providing longer partial solutions greatly improves the accuracy of the resulting traces. This motivates to leverage such high-quality self celf-generated traces for fine-tuning, enabling the model to effectively internalize the successful search strategies.
However, naively distilling partial solutions as hints might negatively impact the model's performance. Table shhows the cross-entropy loss of successful search traces initialized from partial solutions of different lengths, demonstrating that traces from longer partial solutions results in higher loss values. This happens because longer partial solutions follow distribution that differs significantly from the model's 's search behavior. Consequently, fine-tuning directly on these traces might cause substantial parameter updates, potentially degrading the model's search capabilities. Therefore, it is important to explore methods for effectively integrating optimal solutions into self-generated traces, ensuring that the resulting traces maintain both high likelihood and solution quality.
3.2 Guided reinforced self-training
In this subsection, we introduce novel fine-tuning algorithm, called guided reinforced self-training (Guided-Res) ReST,, which seamlessly incorporates optimal solutions into the model's search process for data generation and effectively distills improved search behaviors into the model. The main idea is to leverage unsuccessful search attempts as context for each intermediate reasoning step within the optimal solution, mimicking the model's inherent search process and resulting in trajectories with high likelihood under the model. Before delving into approach, we first introduce some notations. We define subgoal node as any node along the optimal solution in the search tree. Thus, the optimal solution consisting of T reasoning steps contains T subgoal nodes.
Now, consider an unsuccessful search trace that has reached (1 1) subgoal node but fails to navigate its subtree, and thus does not reach the 1-th subgoal node. To transform it into successful search trace, our method first randomly selects an unsuccessfully explored child node of the (t 1) subgoal node. Next, it replaces this child node with the correct k-th subgoal node and truncate the

---

Algorithm AUGMENTŠUBGOAL: Generating search traces via subgoal augmentation
Require: model no, question Ω. search trace Z, optimal solution S, subgoal index t 1: if Z has already explored the t-th subgoal node of S then
2: return Z
3: end if
4: Randomly select explored child node from the 1)-th subgoal node of S
5: Replace the selected child node with the 1-th subgoal of
6: Truncate the search trace up the modified node
7: Continue the search: Z + Te(- q, 2)
8: return Z
Algorithm 2 Guided reinforced self-training
Require: model Tref, question optimal solution
1: Initialize model parameters: T <- Fref
2: for i= 1,
MAXITER do Generate search trace: Z Te(- (g)
4: for t = ...,...,1 , do 5: if successful then
6:
break
7: end
8: Add the t-th subgoal to the search trace: Z + AUGMENTSUBGOAL S,t) end for
9:
10: Fine-tune the model via supervised learning: <- TRAIN Trof 2)
end for
search trace up the modified node, the subsequent search history becomes inconsistent with the updated subgoal. After the modification, the model continues the search from this augmented subgoal node. This algorithm, called subgoal augmentation, is summarized in Figure 2 and Algorithm [ We iteratively run this algorithm until all subgoal nodes have been incorporated into the search trace. An example of the resulting search trace provided in Appendix C. Finally, the model fine-tuned on the successful search traces via supervised learning over multiple iterations. Algorithm summarizes the overall training procedure.
key advantage of Guided-ReS" is its ability to explicitly teach the model where to backtrack during unsuccessful searches. Unlike ReST, which treats failed search traces uninformative, or RL, which relies solely on terminal reward signals, it leverages optimal subgoal information to provide corrective feedback failure points and guide the model continue the search from more promising points. By systematically injecting the subgoal information to search traces, we expose the model to high-quality local corrections grounded in successful reasoning steps. This enables the model not only generate correct answers but also to internalize structured reasoning strategies and the ability to recover from failure.
3.3 Operation-level RL fine-tuning
Building upon the fine-tuned model using Guided-ReST, further improve its performance via RL fine-tuning using PPO. In standard PPO fine-tuning, the importance ratio is computed in a token-level Markov Decision Process (MDP). However, since our goal is to optimize the model's search strategy rather than token-level generation, we reformulate PPO in an operation-level MDP. Specifically, each action aę defined as a sequence of tokens (ah.1..... a.,L) )) that represents a single tree operation.
Consequently, each state is defined as the concatenation of the question and all previously executed operations. In this MDP formulation, the log importance ratio between the learner policy no and the sampling policy Wo.s is defined as
log Th
Σ (log To (anee (h,1:t-1,85) log To.ta (an,e (h.1:e-1,8hh) 8h))
5

---

Algorithm 3 Guided reinforced self-training (episode-level)
Require: model reef question q. optimal solution S, maximum turn limit T
1: Initialize model parameters: Tp f- Fzef
2: for i = 1, ... MAXITER do
3: Generate a multi-turn episode: Z t- To(- g)
4: for / = 1, 1, ...,,,, T do
5: if is successful then
6:
break
7:
end if
8:
simplified subgoal augmentation Truncate the episode up to the f-th tur:: Z + (S1, .2....,,i. St) Append user feedback containing the optimal solution < (S1, Ez...., Et, E141)
9:
11:
Continue the episode: Z + To(- 1.. 2) end for
13: Fine-ture tune the model via supervised learning: To < TRAIN Tof 2.
14:
end for
Finally, train the learner policy xo and the value function V with the standard PPO objective:
max B(sn:an)~Tes [min (r% (0) An, clip (rx (0), c, + +c) An)).
min (oon.s))Ntoo.l
(v)) - R2.)
where >> O is the clipping coefficient. We compute the advantage Ad and return R. using the Monte Carlo method instead of GAE 24 and remove the KI. penalty term, following the recent practice of Yu et al. 40
4 Application to episode-level search: code self-repair
So far, we have established our method on Countdown, challenging yet formally defined task with finite search space. In this section, we extend it to more realistic domains involving longer reasoning chains and more complex responses. Unlike Countdown, where the entire step-wise search trace fits within the context window, the increased output length and complexity make fine-grained tree search impractical. Therefore, we consider episode-level search, which generates complete solution attempts and iteratively refines them through multiple rounds of revision (20] 29 10). Specifically, given question y, a search trace Z is represented as multi-turn episode:
Z = E Sr),
Ey,
where S, denotes complete set of reasoning steps and the resulting solution generated by the model and E, denotes user feedback containing verification results for the previous model's 's response and a revision instruction. The episode ends when the model's response passes verification or the maximum turn limit Taax is reached.
In this setting, we apply our method to the code self-repair task, where the model first generates an initial code solution and then iteratively refines it based on verification results from a public set. The optimal solution is relatively easy to obtain because correctness can be automatically verified via code execution, making this task an effective testbed for evaluating our method. However, directly applying our method to this task still presents challenges, since subgoals are difficult to define within cither the optimal solution or the search trace.
To address this issue, we employ simplified, episode-level variant of Guided that progressively guides the model through user feedback augmented with the optimal solution hint. At each round t, the algorithm first truncates the current unsuccessful episode up to the t-th turn, then augments the subsequent user feedback E141 with the optimal solution, and finally regenerates the remaining turns to complete the episode. The process repeats until the round reaches the maximum turn limit. As the round progresses, the number turns augmented with the optimal-solution hint gradually increases, yielding progressively guided search traces. Finally, the model is fine-tuned on the successful search traces via supervised learning over multiple iterations. Algorithm 3 summarizes the overall training procedure. Additional details are provided in Appendix D.

---

Experiments
5.1
Setup
5.1.1 Countdown
We use Llama-3.2- IB-Instruct the base model [5] with maximum response length of 4K tokens For SoS, we generate search traces using heuristic-guided DFS and BFS over 500K training examples and fine-tune the base model for two epochs. For GuSd we generate a single search trace for each problem using 200K training examples and fine-tune the model for two epochs, repeating this procedure for three iterations. For PPO, we fine-tune the model on 200K training examples for two epochs. We adopt an outcome reward function that assigns for success and 0 otherwise. Additional details are provided in Appendix E.
For evaluation, we follow the protocol of Gandhi et al. [ We construct 10K test examples for each of two settings: (1) seen targets, where the target numbers overlap with those in the training data but are paired with different input numbers, and (2) unseen targets, where the target numbers are entirely disjoint from the training data. We measure accuracy while varying the token budget.
We compare our method against BC, SoS, and the two fine-tuning methods introduced in Gandhi et al. 4:: (1) ReST and (2) PPO. Although the original SoS paper employs iterative advantage- induced policy alignment (APA) as the base RL algorithm [43, we find that PPO outperforms iterative APA. Unless otherwise specified, we conduct all RL fine-tuning experiments using operation-level PPO.
5.1.2 Code self-repair
We use Qwen2.5-7B 2BB.5-77BIInttree as the base model [36. with a maximum response length of 16K tokens. We build the training data using Primelntellect's SYNTHETIC-1 [16), which was curated from APPS, CodeContests, and TACO [8] [3). We select problems that have at least one ground-truth solution, resulting in 16K problems. For Suided-RRST.. We generate eight episodes per problem, each with a maximum of four turns, and fine-tune the model for two epochs, repeating this procedure for three iterations. We compute the loss only on the last model response, following the practice of Snell et al. 29 Additional details provided in Appendix E.
For evaluation, we consider two code benchmarks: (1) CodeContests and (2) CodeForces [4 211. CodeContests consists of 165 problems and represents near-distribution evaluation, data was partially curated from the same source. On the other hand, CodeForces consists of 408 more recent problems, providing more challenging cuttof-diitributtio benchmark. We measure accuracy by extracting code from the model's 's response and executing on private test cases.
training
We compare our method against the base model and ReST. We do not perform RL fine-tuning due to limited computational resources. Instead, we evaluate performance using pass@ accuracy, which reliable proxy for RL fine-tuning [41]
5.2 Countdown results
Figure 3 shows the performance of each method across different token budgets on Countdown, where responses are sampled using greedy decoding. Our method significantly outperforms SoS and the baseline fine-tuning methods. At the maximum token budget, our method achieves 87% accuracy on both seen and unseen targets, outperforming the second-best PPO by more than 10%. It is important to note that the improved performance unseen targets indicates that our method does not merely memorize the optimal solutions provided as guidance. Instead, it internalizes a more effective search strategy, leading to stronger generalization. Moreover, our method achieves comparable accuracy to the second-best PPO while consuming only about 50% of its tokens, demonstrating a more efficient utilization of test-time compute. BC achieves less than 40% accuracy on both seen and unseen targets, indicating substantially weaker generalization performance compared to tearch-based models.
We choose the Llama-3 model for Countdown due to its efficient tokenization of integers.
7

---

SoS
RcST
PPO
Ours
ReST PPO
PPO
100
100
80
80
Accuracy (%)
Accuracy (%)
60
60
40
20
20
1K
2K
3K
4K
1K
2K
3K
4K
Token budgets
Token budgets
Figure 3: Accuracy of method and baselines on Countdown. Solid lines represent seen targets, and dotted lines represent unseen targets.
Figure 4: Accuracy of PPO initialized with ReST on Countdown. PPO does not benefit from ReST in the high-budget regime.
Table Pass@ accuracy ReST and Guided-ReSt on Countdown (32 samples per problem).
Seen target
Unseen target
Method
2
4
8
16
32
1
2
4
8
16
32
SoS
55.3
63.8
69.6
73.3
75.5
77.1
53.2
62.3
68.7
73.0
75.6
77.5
ReST
62.3
67.7
71.3
73.6
75.3
76.6
60.8
66.7
70.8
73.5
75.5
77.0
Guided-ReST
62.7
75.3
84.6
90.8
94.7
96.8
61.0
74.1
83.8
90.3
94.3
96.8
Impact of Guided-Resh on RL fine-tuning
Our method leverages both Muided-ReeST and PPO for fine-tuning. This naturally raises the question of whether PPO can achieve similar synergy with standard ReST, which does not leverage optimal- solution guidance during data generation. Figure/4| the accuracy of PPO fine-tuning initialized with ReST. We observe that PPO does not benefit from ReST, but achieves significant improvements when combined with Guided-Rest ReS.
To examine why Guided-ReST exhibits stronger synergy with PPO, we evaluate the pass accuracy. Table 2 presents the pass @ accuracy of ReST and Guided-ReST on Countdown, where responses are sampled temperature 1.0. ReST outperforms SoS at pass@1 pass @ 32. Recent work suggests that RL mainly shifts probability mass from already-likely correct responses toward top-ranked ones [41]. Consequently, since ReST and SoS exhibit similar coverage of correct candidates, PPO initialized with ReST yields performance comparable to that initialized with SoS. contrast, Guided-Red -elieed-eeeTTaee achieves a similar pass accuracy to ReST but delivers much larger accuracy gains increases. This broader coverage provides PPO with more correct candidates to amplify, leading to significant performance improvements.
shows almost no difference
5.4 Analysis of operation-level MDP
We investigate the effectiveness of the operation-level MDP formulation for RL fine-tuning introduced in Section 3.3 Figure shows the performance of our method under the standard token-level MDP. We find that fine-tuning the model under the operation-level MDP not only improves performance from 83% 87% at the maximum token budget but also substantially enhances test-time compute efficiency. Compared to the token-level MDP, our approach achieves approximately 2 higher token efficiency in the low-budget regime and about 1.5x higher efficiency the high-budget regime. This result highlights the importance of defining MDP aligned with the optimization objective.

---

Ours (token-level)
Ours
PPO (dense)
PPO
100
100
80
80
Accuracy (%)
Accuracy (%)
60
60
40
20
20
1K
2K
3K
4K
1K
2K
3K
4K
Token budgets
Token budgets
Figure 5: Accuracy of our method when trained under the token-level MDP on Countdown.
Figure 6: Accuracy of PPO trained with dense reward on Countdown.
Table 3: @ accuracy of ReST and Guided-Reet Te on code self-repair (128 samples per problem).
CodeContests
CodeForces
Method
8
16
32
16
32
Base
4.5
7.5
11.3
15.7
20.6
25.8
5.5
8.5
12.5
17.2
22.4
27.7
ReST
9.4
13.1
17.0
21.3
25.9
30.4
9.7
14.2
19.3
24.8
30.5
35.9
Guided-ReST
10.5
14.8
19.4
24.1
28.9
33.9
9.7
14.5
20.2
26.2
32.0
37.6
5.5
RL fine-tuning with dense reward
Our method leverages subgoal guidance derived from optimal solutions during data generation. As alternative, one can directly perform RL. fine-tuning using subgoal-based rewards. Given an optimal solution S
(S,,..., ST) and search trace Z (ZI, 2T"). ), we define the subgoal reward
Raubyoal (Zi:e 9)) = {o (0.25 there exists s; such that 2p reaches St
otherwise,
We fine-tune the SoS model using PPO with dense reward that combines the subgoal and outcome reward functions! To prevent reward hacking, the subgoal reward is applied only once per subgoal.
Figure 6 presents the performance of PPO with the dense reward. The dense reward has little effect on overall performance and even causes slight degradation in the high- high-budget regime. These results underscores the importance of Guided-ReST in conditioning the model before the RL phase.
5.6 Code self-repair results
As observed in the Countdown experiment in Section 5.3 Guided-ReST achieves similar pass @ accuracy to ReST but yields higher pass k as & increases, which translates into greater improvements during PPO fine-tuning To examine whether this effect generalizes beyond Countdown, conducted the same analysis on the code self-repair task.
Table 3 presents the pass@k k accuracy of ReST and Mhided-ReST" on CodeContests and CodeForces. Guided-ReST consistently achieves higher accuracy than ReST, with the gap between them widening as & increases. Note that ReST also outperforms the base model, as the base model is general-purpose rather than task-specific. The overall improvement remains moderate compared to Countdown, likely due to fewer training examples, incomplete subgoal augmentation, and weaker search capabilities.
Thhe total reward ranges from 0 1.5, each Countdown problem contains non-trivial subgoals.

---

6
Related work
Searching with language models Yang al. [37 first introduce the concept of training sequence model to imitate expert MCTS policy search traces, primarily to enhance imitation learning. More recent studies have shifted focus toward directly improving search capabilities by training language models on these traces and fine-tuning them with self- geneeeateeedd data [1] In contrast to these prior studies, which apply generic self-improvemmnt algorithms not tailored for search [42 G our work analyzes the relationship between the quality of self- self-generated data and the efficiency of test-time compute, and introduces a self-improvement algorithm specifically optimized for search efficiency.
Several studies have also explored sequential revision and episode-leve search [35 20 29] but these do not address improvements in test-time compute efficiency.
Fine-tuning self-generate data Anthony et al. [2] first introduce the idea of iteratively distilling self-generated data into neural networks, generating high-quality trajectories using an expert MCTS policy and imitating them to improve neural network- based policy. Recent studies have extend this idea to fine-tune tune language models on self- gefereeee data for various tasks, including mathematical reasoning, question answering, and machine translation [19] 42 6. although these approaches not specifically tailored for search. The work most closely related to ours is that of Chang et al. [3], which employs stronger guide model alongside weaker base model to generate data and fine-tune the base model using PPO. Their method, however, is limited to a single round of interaction between the guide and base models. Our approach performs multiple rounds of guide-base interactions, enabling richer and more iterative supervision.
7 Conclusion
We introduce Guided-ResT, self-training algorithm designed to improve the search efficiency of language models. By incorporating optimal solution as guidance for generating high-quality search traces, our method enables models to internalize more effective search strategies. Experiments on the Countdown and code self-repair tasks demonstrate that our method not only enhances accuracy but also scales robustly with increased test-time compute, outperforming existing fine-tuning approaches.
Our work is not without limitations. It assumes access to optimal solutions, which may not always be available in practice. Nevertheless, this assumption can be relaxed by leveraging solutions generated by a sufficiently capable model. Once such model is available, a promising direction is to develop collaborative framework in which stronger model provides guidance when the base model fails to make progress, and this interaction is subsequently distilled into the base model through additional training. This teacher-studee student framework facilitates the transfer of effective search strategies from the teacher to the student model.
Another limitation lies in the assumption of an oracle in Countdown that precisely identify the first incorrect step and backtrack to that point. Although such an oracle provides a clear and reliable signal for targeted correction, this assumption becomes unrealistic when replaced with language model, particularly standard instruction-folloiing models that still struggle to recognize reasoning errors or perform localized backtracking. A promising direction for future work extend our approach to reasoning models that already exhibit emerging backtracking and self-correction capabilities [7].
Acknowledgements
This work was supported in part by Institute of Information & Communications Technology Planning & Evaluation (ITP) grant funded by the Korea government (MSIT) (No. R82020-11200882. (SW STAR LAB) Development of deployable learning intelligence via self-sustainable and trustworthy machine learning. No. RSS-002-11220488, Development of Training and Inference Methods for Goal-Oriented Artificial Intelligence Agents, No. RS-2021 10021-I0 1--220211121130 112111334, Artificial Intelligence Graduate School Program (Scoul National University)). and by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. supported by grant from KRAFTON AI. This material is based upon work supported by the Air Force Office of Scientific Research under award number FA2386-22 1-----2-----00 1-404. Hyun Oh Song is the corresponding author.
This research was
10

---

References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[2] Thomas Anthony, Zheng Tian, and David Barber. Thinking fast and slow with deep learning and tree search. In NeurlPS, 2017.
[3] Jonathan D Chang, Kiante Brantley, Rajkumar Ramamurthy, Dipendra Misra, and Wen Sun. Learning to generate better than your Ilm. arXiv preprint arXiv:2306. 1816, 2023.
[4] Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, and Noah Goodman. Stream of search (sos): Learning to search in language. arXiv preprint arXiv:2404.03683 2024.
[5] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
[6] Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, et al. Reinforced self-training (rest) for language modeling. arXiv preprint arXiv:2308. 089988 2023.
[7] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12948. 2025.
[8] Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, et al. Measuring coding challenge competence with apps. arXiv preprint arXiv:2/055 09938, 2021.
[9] Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Camey, al. Openai ol system card. arXiv preprint arXiv:2412.16726, 2024.
[10] Aviral Kumar, Vincent Zhuang, Rishabh Agarwal, Yi Su, John D Co-Reyes, Avi Singh, Kate Baumli, Shariq Igbal, Colton Bishop, Rebecca Roelofs, al. Training language models to self-correct via reinforcement learning. In ICLR, 2025.
[11] Lucas Lehnert, Sainbayar Sukhbaatar, Paul Movay, Michael Rabbat, and Yuandong Tian. Beyond a*: Better planning with transformers via search dynamics bootstrapping. arXiv preprint 2rXiv:2402.14000 2024.
[12] Raymond Li, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, LI Jia, Jenny Chim, Qian Liu, et Starcoder: may the source be with you! Transactions Machine Learning Research, 2023.
[13] Rongao Li, Jie Fu, Be- Zhang, Tao Huang, Zhihong Sun, Chen Lyu, Guang Liu, Zhi Jin, and Ge Li. Taco: Topics in algorithmic code generation dataset. arXiv preprint 2rXiv:2312.14852 2023.
[14] Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition- code generation with alphacode. Science, 2022.
[15] Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman. llya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint arXiv:2305.20005 2023.
[16] Justus Mattern, Sami Jaghouar, Manveer Basra, Jannik Straube, Matthew Di Ferrante, Fe- lix Gabriel, Jack Min Ong, Vincent Weisser, and Johannes Hagemann. Synthetic-1: Two million collaboratively generated reasoning traces from deepseek-rl, 2025. URI. httpps
//www, www.peeeeenttelec.....mmmm...... w:w.lo///ynthheiio--ieeee r-eeeas

---

[17] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida. Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray. et al. Training language models to follow instructions with human feedback. In NeurlPS, 2022.
[18] Vishal Pallagani, Bharath Muppasani, Keerthiram Murugesan, Francesca Rossi, Biplav Srivas- tava, Lior Horesh, Francesco Fabiano, and Andrea Loreggia. Understanding the capabilities of large language models for automated planning. arXiv preprint arXiv:2305. 16151, 2023.
[19] Stanislas Polu, Jesse Michael Han, Kunhao Zheng, Mantas Baksys, Igor Babuschkin, and llya Sutskever. Formal mathematics statement curriculum learning. arXiv preprint arX1v:2202.01344, 2022.
[20] Yuxiao Qu, Tianjun Zhang, Naman Garg, and Aviral Kumar. Recursive introspection: Teaching language model agents how to self-improve. In NeurIPS, 2024.
[21] Shanghaoran Quan, Jiaxi Yang, Bowen Yu, Bo Zheng, Dayiheng Liu, An Yang. Xuancheng Ren, Bofei Gao, Yibo Miao, Yunlong Feng, et al. Codeclo: Benchmarking competition-leeel code generation of Ilms with comaaacomearable elo ratings. arXiv preprint arXiv:2501.01257, 2025.
[22] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In AISTATS, 2011.
[23] Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, al. Code Ilama: Open foundation models for code. arXiv preprint arXiv:2308. 12950, 2023.
[24] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. In ICLR, 2016.
[25] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1.XXiv:707.0060 2017.
[26] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Mingchuan Zhang, YK Li, Yu Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03380 2024.
[27] Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. NeurIPS, 2023.
[28] Avi Singh, John D Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Xavier Garcia, Peter J Liu, James Harrison, Jachoon Lee, Kelvin Xu, et al. Beyond human data: Scaling self-training for problem-solving with language models. arXiv preprint arXi::331.006885 2023.
[29] Charlie Victor Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling tist-time compute optimally can be more effective than scaling Ilm parameters. In ICLR, 2025.
[30] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss. Alec Radford, Dario Amodei, and Paul Christiano. Learning summarize with human feedback. In NeurIPS, 2020.
[31] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Amn Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 2023.
[32] Ikechukwu Uchendu, Ted Xiao, Yao Lu, Banghua Zhu, Mengyuan Yan, Joséphine Simon, Matthew Bennice, Chuyuan Fu, Cong Ma, Jiantao Jiao, et Jump-start reinforcement learning. In ICML, 2023.
[33] Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models a critical investigation. In NewrIPS, 2023.
12

---

[34] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe Le, Ed H Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consisteney improves chain of thought reasoning in language models. In ICIR, 2023.
[35] Sean Welleck, Ximing Lu, Peter West, Faeze Brahman, Tianxiao Shen, Daniel Khashabi, and Yejin Choi. Generating sequences by learning to self-correct. In ICLR, 2023.
[36] An Yang. Baosong Yang. Beichen Zhang. Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2.5 technical report. arXiv preprint arXiv:2412.15115, 2024.
[37] Mengjiao Sherry Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Chain of thought imitation with procedure cloning. In NeurIPS, 2022.
[38] Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. In NeurlPS, 2023.
[39] Tian Ye. Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu Physics of language models: Part 2.2, how to learn from mistakes grade-school math problems. arXiv preprint arXiv:2408.16223, 2024.
[40] Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source Ilm reinforcement learning system at scale. arXiv preprint arXiv:2503. 14476, 2025.
[41] Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model? arXiv preprini 2rXiv:2504. 13837, 2025.
[42] Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. In NeurlPS, 2022.
[43] Banghua Zhu, Hiteshi Sharma, Felipe Vieira Frujeri, Shi Dong, Chenguang Zhu, Michael I Jordan, and Jiantao Jiao. Fine-tuning language models with advantage-induced policy alignment. arXiv preprint 22.222230. 2023.
13

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We clearly explain our paper's contribution in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We clearly explain the limitations in the conclusion.
14

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: Our paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all the information (e.g., hyperparameters) in the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
15

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. w
For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide a link the source code in the abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ 7/nips. /ttppp.cc/pbllc/guudeeee http://oeeubiiisssoooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
16

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We specify all details in the experiment section and supplementary material.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We run experiments with multiple random seeds (32 128) and report average results for those with high variance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the computation resources in the supplementary materials.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
17

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:://awaaasccc//uuu///ttcccca
Answer: [Yes]
Justification:
We have carefully reviewed the NeurIPS Code Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss potential societal impacts in the Supplementary materials.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: Our work poses no such risks, as it focuses on arithmetic and code self-repair tasks.
Guidelines:
18

---

The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We provide the license of the models we used in the supplementary material.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We provide a link to the source code the abstract.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
19

---

Justification:
We does not involve crowdsoureing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: We use LLMs only for writing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy httess//Deerrrppcccc///fferrreep for what should or should not be described.
20


Reviews:
[{"review_id": "tsxW9UFrvl", "reviewer": "Reviewer_xv2X", "summary": "The paper builds on Stream-of-Search (SoS) and introduces Guided-ReST, a fine-tuning framework that improves the (inference) compute-efficiency of language-model search. Guided-ReST treats each optimal solution as a sequence of landmark sub-goals: during self-generation it replaces a failing child node in the search tree with the next correct sub-goal, then lets the model resume its own exploration, creating high-likelihood, high-quality traces. These augmented traces are distilled back into the model through several supervised iterations, followed by RL that further refines the policy while shortening the credit-assignment horizon.  Experiments rely on optimal solutions with subgoals and provide mixed results on the arithmetic Countdown benchmark and is also adapted with operation-level MDP to MATH-500 and AMC23.", "questions": "- Could you provide approximate training GPU-hours/FLOPs and average inference latency (ms per token or ms per problem) for Guided-ReST versus baselines? This would let readers verify that the method is compute-efficient in wall-clock terms, not just by token budget.\n- Could you please provide 95% confidence intervals to establish that the gain is statistically reliable?\n- Do you have quantitative data—e.g., distributions of log-likelihood, median trace length, or success-per-1k tokens—showing that sub-goal augmentation yields measurably higher-quality traces and that this correlates with accuracy gains?\n- I’m not sure how manageable it would be but can you run Guided-ReST on at least one contemporary 7–13B open model on a different task that is not saturated easily by such large models to establish that the observed gains (with confidence intervals) are not limited to the 250M-parameter regime? This would solidify the validity of the main claim.", "limitations": "Training compute overhead and overall latency/FLOPs is an aspect that has not been discussed. Scalability beyond small-scale models is questionable. Reliance on optimal solutions is acknowledged by the authors.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The authors clearly articulate the problem, noisy search traces waste test-time compute, and the solution, Guided-ReST, that injects optimal-solution “landmarks” during self-training. The experiments are clear and show measurable gains in Countdown and mixed results with math benchmarks.\n- The paper improves test-time compute efficiency by reducing the unnecessary wasteful exploration of mistakes. The method is a training algorithm that encourages more economic behaviour at inference time.\n- Although the results are not as compelling on the math benchmarks, the initial effort and implementation of the idea in such benchmarks is commendable.\n- Generally well-written and clearly articulated. The paper situates itself nicely within prior work, offering a concise recap of SoS, JSRL, ReST, etc. that clarifies the baseline mechanics and helps the reader contextualize the improvements. Additionally, the paper follows a nice cycle of introducing a claim and then validating it with simple and clear-to-understand results throughout the paper.\n- Table 1 is particularly insightful and provides a clear empirical motivation: it reveals that while partial-solution hints boost success rates, they simultaneously lower likelihood—highlighting a distribution-shift problem that explains why naïve distillation fails and thereby setting the stage for the more principled Guided-ReST approach.\n- operation-level MDP for improved credit assignment due to long search traces.\n- Figure 3 shows the model maintains performance on unseen target numbers and continues to improve as the token budget increases, indicating it has learned transferable search strategies rather than merely memorizing the injected landmarks.\n- Ablation in Figure 4 demonstrates that Guided-ReST pre-alignment is essential for PPO to realize its performance gains. Similarly, Figure 7 shows naive subgoal rewards help only marginally, strengthening the case that data-centric Guided-ReST is the main lever.\n- Upfront acknowledgement of limitation, i.e., dependence on optimal solutions, is noteworthy.\n\n\n\nWeaknesses:\n\n- As acknowledged by the authors, the method—and much of the empirical lift—relies on having exact optimal traces for every training example. The authors themselves concede this as a limitation and propose a stronger “teacher” as a future remedy, but without such oracles, the pipeline is hard to replicate in domains where gold CoT traces are scarce or noisy.\n- While accuracy continues to rise up to 4k tokens, there is no accounting of *actual* latency or compute cost. Compute-efficiency claims are based on token budgets, not wall-clock or FLOPs. PPO fine-tuning and operation-level value networks introduce extra parameters and training epochs, so the headline “efficient search” narrative remains unquantified in resource terms.\nIn short, the pretraining part seems comparable to SoS, but there are extra sub-goal-augmented traces (for 3 iterations) + an operation-level value network, which adds more parameters and FLOPs. Therefore, an account of training FLOPs/GPU-hours would be helpful to substantiate the overall compute efficiency (as opposed to only considering the test-time efficiency).\n- Related to the above point, only top-1 accuracy is reported; no look at trajectory length, failure modes, or calibration of the value network. Without this, we can’t tell if the model simply finds more solutions or just spams longer traces that eventually hit a solution. Median / percentile trace lengths could shed light on whether the model spams many tokens at larger budgets compared to the baseline or not.\n- Not a reason to reject, however, benchmarking with stronger baselines, e.g., 7-13B parameter open models, would be helpful to show that the gains are not due to originally under-performing models. I understand the GPT-2 experiments follow SoS, and the gains are measurable in that regime; however, the marginal improvement in math benchmarks with larger models calls for scrutiny over the actual gain of the method on some other task in which gains can be demonstrated when large models are used.\n- Generality outside arithmetic search is only modestly demonstrated. On MATH-500 and AMC23, the lift over ReST is small and absolute accuracy remains low (≤60 % and ≤32 %, respectively). Moreover, the adaptation drops the elegant subgoal-replacement machinery in favor of a simpler “mask incorrect solution” heuristic, suggesting the method’s sophistication may not carry over cleanly to richer natural-language reasoning. Either a narrower scope claim (effective for tasks with verifiable sub-steps) or a stronger transfer experiment is needed for the main claims to be acceptable.\n- It would be helpful if the link between trace quality and accuracy could be probed. The authors argue that higher-likelihood traces improve learning, but beyond Table 1, this “trace noise” is not quantified or correlated with performance (this also is not a major reason for rejection).\n- Only average results reported—no error bands (despite marking yes in checklist question 7). A more accurate representation of the results and their variability is required to substantiate the claims, especially in Countdown, as the gains in math benchmarks are quite marginal.", "was_revised": true, "final_justification": "The author's rebuttals addressed most of my concerns. I have therefore increased my rating.", "rebuttal": "Thank you for your helpful feedback. We appreciate your positive remarks on the clarity of our manuscript and your suggestions regarding compute efficiency, pass@k accuracy, and other evaluation metric, all of which help us to strengthen our paper. Our responses to your specific questions are provided below.\n\n---\n\n**Q1. Reliance on exact optimal traces**\n\nWe acknowledge that the reliance on exact optimal traces makes our method not readily applicable to general reasoning tasks. Nevertheless, we believe the paper offers useful insight into two fundamental questions.\n1. Can outcome-level rewards alone improve an LLM's in-context search capabilities?\n2. If outcome rewards are insufficient, what additional supervision is needed?\n\nOur results indicate that neither ReST nor PPO on their own discover new search behavior; the model acquires novel search capabilities only when step-by-step guidance is provided through supervised learning. This is evident in the pass@k metrics shown below (please refer to Q3), where ReST improves pass@1 but lags behind the base model from pass@8 onward, whereas Guided-ReST continues to improve across all k.\n\nWe plan to explore how this strong assumption can be relaxed (e.g. using long CoT models)–and how far that relaxation can be pushed–in future work.\n\n---\n\n**Q2. Actual latency and GPU-hours**\n\nFirst, we compared the compute efficiency of ReST and Guided-ReST and report the GPU hours required for data generation and training below. Guided-ReST adds 7 RTX 3090 GPU hours at the data-generation stage, which amounts to less than a 9% increase in total compute–and even less once the much higher throughput of the A100 used for training is considered.\n\n| Method | Data generation | Training |\n| --- | --- | --- |\n| ReST | 20 hours (RTX 3090) | 55 hours (A100) |\n| Guided-ReST | 27 hours (RTX 3090) | 55 hours (A100) |\n\nRegarding PPO, its computational cost is much smaller than a single iteration of ReST or Guided-ReST.\n- For data generation, PPO only uses 25K training samples, while ReST or Guided-ReST uses 200K training samples.\n- For training, PPO performs 11 forward and 8 backward passes for each sample, while ReST and Guided-ReST perform 10 forward and 10 backward passes for each sample as they run for 10 epochs.\n   - 1 forward for computing old_log_prob\n   - 1 forward for computing ref_log_prob\n   - 1 forward for computing value\n   - 2 forward and backward for updating actor and critic * 4 PPO epochs\n\nTherefore, PPO introduces only a marginal overhead to the overall compute budget. \n\nFor inference, our method does not invoke the trained critic. Therefore, it introduces no additional computational overhead and its latency scales only with the number of tokens generated.\n\n---\n\n**Q3. More analysis on accuracy and trajectories**\n\nCould you please confirm whether the \"top-1 accuracy\" you mention corresponds to pass@1 accuracy? We repeated the experiments on the Countdown task using the more recent Llama-3.2-1B-Instruct model and computed pass@k accuracy on a subset of 1,000 test samples with unseen targets. The results are summarized below.\n\n| Method | Pass@1 | Pass@2 | Pass@4 | Pass@8 | Pass@16 | Pass@32 | Pass@64 |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| SoS | 0.551 | 0.633 | 0.686 | 0.736 | 0.774 | 0.805 | 0.841 |\n| ReST | 0.601 | 0.652 | 0.700 | 0.732 | 0.763 | 0.788 | 0.809 |\n| Guided-ReST | 0.630 | 0.731 | 0.835 | **0.896** | **0.944** | **0.965** | **0.975** |\n| Guided-ReST+PPO* | **0.734** | **0.793** | **0.839** | 0.883 | 0.903 | 0.925 | 0.941 |\n(* denotes the method presented as Ours in the paper.)\n\nReST outperforms SoS at pass@1 yet falls behind from pass@8, which is consistent with recent observation that RL mainly boosts the probability of already-likely correct solutions without discovering new ones [1, 2]. In contrast, Guided-ReST improves accuracy at every k. This explains why ReST+PPO yields limited gains whereas Guided-ReST+PPO remains effective. We also conducted an analysis of how many problems that DFS and BFS leave unsolved can be solved by our methods (Figure 4 in the supplementary material), which further highlights and clarifies the strengths of our approach.\n\nRegarding trajectory length, the table below reports the quantile statistics of the search traces for 10,000 test samples with unseen targets.\n\n| Method | 0.0 | 0.1 | 0.2 | 0.3 | 0.4 | 0.5 | 0.6 | 0.7 | 0.8 | 0.9 | 1.0 |  \n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| SoS | 197 | 200 | 467 | 664 | 1503 | 2123 | 3874 | 4096 | 4096 | 4096 | 4096 |\n| ReST | 197 | 259 | 518 | 777 | 1570 | 2400 | 3465 | 4096 | 4096 | 4096 | 4096 |\n| Guided-ReST | 197 | **197** | 394 | 653 | 1392 | 1928 | 2998 | 4096 | 4096 | 4096 | 4096 |\n| Guided-ReST+PPO* | 197 | **197** | **332** | **529** | **1127** | **1813** | **2619** | **3582** | 4096 | 4096 | 4096 |   \n(* denotes the method presented as Ours in the paper.)\n\nGuided-ReST and Guided-ReST+PPO consistently produce shorter traces than the baselines, whereas ReST sometimes generates longer traces than SoS. This confirms Guided-ReST does not simply generate overly longer traces that eventually stumble into a solution; instead, it conducts a more efficient and effective search.\n\nThank you for the suggestion and we will incorporate these analyses into the main paper.\n\n---\n\n**Q4. Link between trace quality and accuracy**\n\nThank you for pointing this out. We conducted experiments to control the quality (likelihood) of the subgoal-augmented search traces. In Algorithm 1, line 4 lets us control how much of the search trace is pruned. Concretely, we can modify line 4 in two ways:\n1. Aggressive pruning: replace the current node with the first child visited (e.g., node 2 in Figure 2). This removes the entire exploration history of that subtree, so the resulting traces are expected to be shorter but lower in quality.\n2. Conservative pruning: replace the current node with the last child node visited (e.g., node 9 in Figure 2). This preserves most of the previously explored trace, so the resulting traces are expected to be longer (inefficient) but higher in quality.\n\nWe report the average likelihood and token length of subgoal-augmented search traces generated by each strategy below; as expected, the more aggressive the pruning, the shorter and higher-likelihood the resulting traces. Note that optimal paths are typically around 200 tokens long.\n\n| Strategy | Likelihood | Token length |\n| --- | --- | --- |\n| Aggressive | 0.170 | 759.4 |\n| Original* | 0.102 | 1449.2 | \n| Conservative | 0.078 | 2038.5 |\n(* denotes the method presented as Ours in the paper.)\n\nWe also report the performance of Guided-ReST under each strategy with 95% confidence intervals.\n\n| Strategy | Accuracy (seen) | Accuracy (unseen) | \n| --- | --- | --- |\n| Aggressive | 0.6677 &plusmn; 0.0014  | 0.6551 &plusmn; 0.0025 |\n| Original* | **0.6900 &plusmn; 0.0076** | **0.6721 &plusmn; 0.0029** | \n| Conservative | 0.6708 &plusmn; 0.0041 | 0.6464 &plusmn; 0.0101 |\n(* denotes the method presented as Ours in the paper.)\n\nThese results highlight that high-quality, efficient search traces are crucial for optimal performance.\n\n---\n\n**Q5. Confidence intervals**\n\nWe ran each experiment on the Countdown task with three independent seeds and report the corresponding 95% confidence intervals below, which demonstrate very low variance across seeds.\n\n| Method | Accuracy (seen) | Accuracy (unseen) | \n| --- | --- | --- |\n| SoS | 0.5747 | 0.5342 | \n| ReST | 0.6282 &plusmn; 0.0019 |  0.5989 &plusmn; 0.0084 | \n| Guided-ReST | 0.6900 &plusmn; 0.0076 | 0.6721 &plusmn; 0.0029 | \n| Guided-ReST+PPO* | **0.7477 &plusmn; 0.0053** | **0.7330 &plusmn; 0.0073** |\n(* denotes the method presented as Ours in the paper.)\n\n---\n\n**Q6. Large model on different task**\n\nAs noted in Q3, we re-ran the Countdown experiments using the more recent, larger Llama-3.2-1B-Instruct model and found that its results were consistent with GPT-2, as shown in the table below.\n\n| Method | Accuracy (seen) | Accuracy (unseen) |\n| --- | --- | --- |\n| SoS (GPT-2) | 0.5747 | 0.5342 |\n| SoS (Llama) | 0.5718 | 0.5391 |\n| ReST (GPT-2) | 0.6279 | 0.6003 |\n| ReST (Llama) | 0.6330 | 0.6033 |\n| Guided-ReST (GPT-2) | 0.6823 | 0.6698 |\n| Guided-ReST (Llama) | 0.6791 | 0.6659 |\n| Guided-ReST+PPO* (GPT-2) | 0.7507 | 0.7361 |\n| Guided-ReST+PPO* (Llama) | 0.7561 | 0.7420 |\n(* denotes the method presented as Ours in the paper.)\n\nRegarding a new task, We are exploring an additional experiment that applies our method to a code-editing setting. While we may not have enough time to complete this study within the rebuttal period, we will run it regardless to strengthen our paper. If the results are ready before the rebuttal period ends, we will share them with you.\n\n---\n\nOnce again, we really appreciate your insightful questions, which greatly help us to improve our paper. We hope that our response above addresses all of your questions.\n\n---\n\n**References**\n\n[1] Yang Yue et al., Does Reinforcement Learning Really Incentivize Reasoning, arXiv 2025 \\\n[2] Andre He et al., Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening, arXiv 2025"}, {"review_id": "qlSSpAjklX", "reviewer": "Reviewer_5NRh", "summary": "The paper introduces a new fine-tuning method called Guided-ReST to improve the search abilities of language models. The method trains the model to \"correct\" its own search path/tree, trained on a synthetic dataset generated using optimal solutions. The paper shows this method ourperforms baseline fine-tuning methods on Countdown and math datasets.", "questions": "Please focus on addressing aforementioned weaknesses. I will consider raising my score if the following questions are well-addressed: \n1. How does the method extend to general reasoning tasks? What are the limitations? Discussion points may include: \n  a) Does search tree / search path make a difference? \n  b) The paper uses a stronger model as a substitue for the optimal solution in mathematical reasoning. Is the performance then bottlenecked by the strong model? \n2. Can you clarify if the main contributions are limited to a data generation scheme during training? \n\nI will most likely not lower my score as long as the following questions are addressed: \n1. Can you propose some revision to Section 3 to improve readability? Specifically, the main confusion may result from how $Z$ is framed in Algorithms 1 and 2. My understanding is that the revised search path is appended to the text during training to obtain a long CoT, but this isn't clearly stated. \n2. Can you discuss some other fine-tuning methods, and why they are not included as baselines?", "limitations": "Yes, though I would appreciate a more in-depth discussion on limitations and how the methodologies can be applied to complex CoT models, perhaps in the appendix.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \nThe method is well-motivated, the introduction is intuitive, and the conclusions are clear. The experiment results are well-presented. The idea of guided search, a form of refinement, is quite interesting. \n\nWeaknesses: \n1. The paper mainly illustrates the method on Countdown, where the reasoning problem is easily represented as a search problem. While the authors also demonstrate the method's effectiveness on math, relevant paragraphs are scarce, and discussions on how the method extends to general reasoning problems seem to be lacking. \n2. The main contribution seems to be limited to a data generation scheme in SFT, where an optimal search path is utilized to help the model more accurately pinpoint the current search bottleneck. \n3. Some parts of section 3 is confusing to me. For one, it is not fully clear what model inference looks like. According to the appendices, the search trace has crucial steps \"Moving to Node #...\" for Countdown and \"There may be an error in your solution\" prompt for math, characterizing the LLM's self-adjustment. I feel like these should also be discussed briefly in the main paper.", "was_revised": true, "final_justification": "The authors have addressed my concerns through the rebuttal, and now I believe this work is a good contribution to the field. Despite this, I raised the score to borderline accept since I think the contributions are not incredibly significant, but it's a good work.", "rebuttal": "Thank you for your constructive feedback. It clearly highlighted which parts needed more explanation and guided us on how to improve them. We would like to address your questions below.\n\n---\n\n**Q1. Extension to general reasoning tasks**\n\nYes. Any task that can be solved by step-by-step chain-of-thought (CoT) reasoning can, in principle, be cast as a search problem, often called meta-CoT [1].\n- The root node is the original prompt.\n- Each node expansion appends a single reasoning step.\n\nThus, a single root-to-leaf path in the search tree corresponds to one complete CoT.\n\nHowever, there might be some practical issues to apply this framework general-purpose reasoning tasks.\n1. Token efficiency\n   - Countdown involves only three hops, and each CoT step is a single arithmetic operation, so the entire search trace remains short and easy to log (still it consumes 4096 tokens to achieve strong performance).\n   - In typical reasoning tasks, the depth is much greater and each step can be multiple sentences, so the entire search trace grows rapidly.\n2. Branching factor\n   - In Countdown, the branching factor is fixed and relatively small.\n   - In natural language reasoning tasks, the branching factor is unbounded.\n3. Existence of oracle\n   - In Countdown, we assume an oracle that pinpoints the first incorrect step in the search track, backtracks to it, and generate the correct next reasoning step.\n   - In general reasoning tasks, these verification and generation steps can be noisy.\n\nAt a high level, our algorithm alternates between an exploratory student and an oracle teacher that provide step-level corrections. This interplay produces search traces with far more diverse backtracking scenarios than the teacher would encounter in self-play. Training on this richer distribution of errors and fixes broadens the student’s experience and can ultimately enable it to surpass the teacher.\n\nWe will add a separate \"Limitations and Future Work\" section to the Appendix in the revised manuscript. \n\n---\n\n**Q2. Contribution limited to data generation scheme**\n\nThank you for raising this point. While our approach introduces a new data-generation scheme for rejection-sampling fine-tuning (RFT), its contribution goes beyond data generation alone. Specifically, we show that\n1. Superior standalone performance (Section 5.2, Figure 3): The subgoal-augmented traces generated by our method consistently outperform the previous RFT baseline (ReST).\n2. Strong synergy with RL fine-tuning (Section 5.3, Figure 4): When followed by PPO, our method yields large additional gains, whereas ReST shows almost no benefit from the same RL stage.\n\nThese results show that the value of our work lies not simply in introducing a new data-generation procedure, but in how it enables a much more effective RFT &rarr; RL pipeline.\n\nWe analyzed why our method shows a strong synergy with RL fine-tuning. Recent studies suggest that RL-based self-improvement mostly amplifies the probability of solutions that are already likely, rather than discovering entirely new ones; in other words, it improves pass@1 accuracy only if the base model already has a high pass@k accuracy for larger k [2, 3]. This observation leads to the hypothesis that our method provides a much stronger initial model with higher pass@k, giving RL more high-quality candidates to amplify.\n\nTo test this hypothesis, We repeated the experiments on the Countdown task using the more recent Llama-3.2-1B-Instruct model and computed pass@k accuracy on a subset of 1,000 test samples. The results are summarized below.\n\n| Method | Pass@1 | Pass@2 | Pass@4 | Pass@8 | Pass@16 | Pass@32 | Pass@64 |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| SoS | 0.551 | 0.633 | 0.686 | 0.736 | 0.774 | 0.805 | 0.841 |\n| ReST | 0.601 | 0.652 | 0.700 | 0.732 | 0.763 | 0.788 | 0.809 |\n| Guided-ReST | 0.630 | 0.731 | 0.835 | **0.896** | **0.944** | **0.965** | **0.975** |\n| Guided-ReST+PPO* | **0.734** | **0.793** | **0.839** | 0.883 | 0.903 | 0.925 | 0.941 |\n(* denotes the method presented as Ours in the paper.)\n\nWhile ReST outperforms SoS at pass@1 accuracy, it falls behind SoS from pass@8 accuracy onward. Therefore, adding an RL stage after ReST may be less efficient than applying PPO directly to the SoS baseline. In contrast, our proposed data generation scheme significantly improves pass@k accuracy at every k. This implies that additional RL fine-tuning can be applied effectively.\n\nWe will highlight the synergy with RL as our primary contribution in both the Introduction and Results sections of the revised manuscript.\n\n---\n\n**Q3. Clarity of Section 3**\n\nWe recognize that Section 2 does not fully explain how a search trace Z is expressed in text, which may have caused confusion when reading Section 3. A search trace is the textual log of every step the solver–symbolic or LLM–takes while traversing the search tree. Each trace consists of three operation types.\n\n1. Expansion – creating child nodes\n   - Countdown: choose two remaining numbers and an operator; the child node is the updated set of numbers after the operation.\n   - Mathematical reasoning: generate one paragraph of reasoning; the child node is the concatenation of all previously generated paragraphs with this new paragraph.\n2. Selection – moving to the next child node for exploration\n3. Verification – checking a leaf node\n   - Countdown: when only one number remains, verify that it matches the target.\n   - Mathematical reasoning: when the solver produces an answer, verify that it is correct.\n\nThus, a search trace is a sequential text record of expansion, selection, and verification lines in the exact order they occur during the search.\n\nIn the Countdown task, logging a search trace is token-light–each operation fits into just two short lines. In mathematical reasoning, however, every node contains the full chain of reasoning generated so far, so expanding multiple children at once would duplicate that text for every child and quickly inflate the prompt. To keep the search trace compact, we therefore enforce the search to expand only one child at a time and move to it immediately. In this regime the combined expansion + selection step is recorded by appending just a single paragraph to the existing trace. In Section 4, we refer to this setting as “episode-level search”–as opposed to step-wise search. The following table represents how each operation is represented in text for each benchmark.\n\n| Operation | Countdown | Mathematical reasoning |\n| --- | --- | --- | \n| Expansion | \"Exploring Operation...\\nGenerated Node...\" | A single reasoning step separated by the double newline |\n| Selection | \"Moving to Node...\\nCurrent State...\" | None |\n| Verification | \"a,b unequal: No Solution\" or \"a,a equal: Goal Reached\" | \"There may be an error in your solution...\" |\n\nReturning to Section 3, Algorithm 1 details how we modify and prune an LLM-generated search trace using known sub-goal information (see also Figure 3 in the supplementary material) and then prompt the model to regenerate the trace. Algorithm 2 describes how we fine-tune the model on the data produced by Algorithm 1.\n\nIn the revision, we will add a precise definition of a search trace and give a detailed explanation of how it is represented in text in Section 2.\n\n---\n\n**Q4. Other fine-tuning methods**\n\nOur paper does not include DPO-style preference fine-tuning. Although we could have constructed preference pairs by ranking search traces on success and length, we concluded that, given a well-defined scalar reward function, directly optimizing it with PPO would be both more straightforward and more effective.\n\n---\n\nOnce again, we really appreciate your constructive feedback, which greatly help us to improve our paper. We hope that our response above addresses all of your questions.\n\n---\n\n**References**\n\n[1] Violet Xiang et al., Towards System 2 Reasoning in LLMs: Learning How to Think With Meta Chain-of-Thought, arXiv 2025 \\\n[2] Yang Yue et al., Does Reinforcement Learning Really Incentivize Reasoning, arXiv 2025 \\\n[3] Andre He et al., Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening, arXiv 2025"}, {"review_id": "Teyzeu0Yda", "reviewer": "Reviewer_yPG3", "summary": "This paper proposes a fine-tuning algorithm for LLMs. Instead of using complete thought traces (as done in the Stream of Search paper), the approach distills effective search traces, used later for fine-tuning. These optimal search solutions are generated using guided subgoal augmentation that replaces random unsuccessful search traces with new searches.\nThe results show how this technique improves LLMs search capabilities for arithmetic and mathematical reasoning tasks.", "questions": "- In Figure 2, the guided subgoal augmentation algorithm is shown. Why is the right path (the one that has node 9 and then 10) removed at step 2? The selected node (at step 1) is 4, and correctly, all the history coming from that is deleted. I do not understand why the path with nodes 9 and 10 is also removed, since it stems from node 1, not node 4.\n- What is the reasoning behind choosing  Llama-3.3 70B-Instruct for generating the step-by-step solutions (Section 5.1.2)?\n- Are the results in Figures 4-5-6-7 representing the performances on Countdown? This is not outlined either in the caption or when referencing those figures.", "limitations": "While the limitations are outlined in the Conclusion section, they could be more detailed. The proposed solution is not sufficient to address the shortcoming previously explained. The criteria to distinguish 'standard' and 'long' CoT models should be outlined earlier in the paper.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths: \n- The paper communicates the main ideas clearly. \n- The references are plentiful and proper, and the authors provide clear attribution when borrowing ideas from other papers. \n- The claims are supported by the results.\n\nWeaknesses: \n- The approach for generating training data relies on having ground truth solutions, which may not be available for the most challenging reasoning tasks;\n- Experiments on the BigMath dataset (Section 5.1.2) rely on Llama 3.3 70B-generated solutions as substitutes for missing optimal solutions, but there's no verification that these generated solutions are actually optimal, which undermines the quality of the training data.\n- The paper is mostly clearly written, but the Experiments section could be restructured. I also suggest having a separate section for results to improve readability.\n- Lines 20-21: \"Their impressive performance is primarily attributed to auto-regressive training on high-quality, internet-scale data\". While internet-scale data is correct, claiming that the data used for training is high-quality is debatable, considering that for closed-source models we do not know exactly what data was used for training.", "was_revised": true, "final_justification": "The authors clearly answered questions and argued the significance of this work. Various improvements have been made, including the addition of computational overhead metrics, metric variance measurements, and experiments using llama models. This raises the quality of this work. Due to this, I raise my score to Accept.", "rebuttal": "Thank you for the valuable feedback, which has greatly helped us improve the clarity of our paper. We would like to address your questions below.\n\n---\n\n**Q1. Reliance on ground-truth solutions**\n\nWe acknowledge that relying on exact optimal traces limits the immediate applicability of our method to general reasoning tasks. For that reason, we use the term \"search\" rather than \"reasoning\" in the title and throughout the main text, reserving \"reasoning\" only for fixed phrases such as \"reasoning step,\" \"arithmetic reasoning,\" and \"mathematical reasoning,\" where it is standard and aids clarity. Although our method rests on a strong assumption, we believe it nevertheless offers valuable insights into two fundamental questions.\n\n1. Can outcome-level rewards alone improve an LLM's in-context search capabilities?\n2. If outcome rewards are insufficient, what additional supervision is needed?\n\nOur results indicate that neither ReST nor PPO alone uncovers new search capabilities. The model develops such capabilities only when step-by-step guidance is added through supervised learning. To substantiate this claim, we report pass@k accuracy below–a metric widely used in recent work to test whether a model truly acquires novel reasoning abilities [1, 2].\n\n| Method | Pass@1 | Pass@2 | Pass@4 | Pass@8 | Pass@16 | Pass@32 | Pass@64 |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| SoS | 0.551 | 0.633 | 0.686 | 0.736 | 0.774 | 0.805 | 0.841 |\n| ReST | 0.601 | 0.652 | 0.700 | 0.732 | 0.763 | 0.788 | 0.809 |\n| Guided-ReST | 0.630 | 0.731 | 0.835 | **0.896** | **0.944** | **0.965** | **0.975** |\n| Guided-ReST+PPO* | **0.734** | **0.793** | **0.839** | 0.883 | 0.903 | 0.925 | 0.941 |\n(* denotes the method presented as Ours in the paper.)\n\nWhile ReST surpasses SoS at pass@1, its accuracy slips below SoS from pass@8 onward, suggesting that ReST mainly boosts the probability of solutions already favored by the model rather than uncovering new ones. By contrast, Guided-ReST improves accuracy at every k, indicating that it enables the model to develop genuinely novel search capabilities. We will state our contributions explicitly in the revised manuscript.\n\nIn addition, we will add a dedicated “Limitations” section. The main limitations we have identified when applying our method to general-purpose reasoning tasks are as follows:\n1. Token efficiency\n   - Countdown involves only three hops, and each CoT step is a single arithmetic operation, so the entire search trace remains short and easy to log (still it consumes 4096 tokens to achieve strong performance).\n   - In typical reasoning tasks, the depth is much greater and each step can be multiple sentences, so the entire search trace grows rapidly.\n2. Branching factor\n   - In Countdown, the branching factor is fixed and relatively small.\n   - In natural language reasoning tasks, the branching factor is unbounded.\n3. Existence of oracle\n   - In Countdown, we assume an oracle that pinpoints the first incorrect step in the search track, backtracks to it, and generate the correct next reasoning step.\n   - In general reasoning tasks, these verification and generation steps can be noisy.\n\n---\n\n**Q2. Use of Llama-3.3-70B-Instruct**\n\nWe followed the same protocol as the Big-Math paper [3]–generating solutions with Llama-3 models and validating them using the final answers–but opted for Llama-3.3-70B-Instruct instead of Llama-3.1-405B-Instruct. The 70B model fits on a single node with 8 GPUs without quantization, yet reportedly matches (and on the MATH benchmark even surpasses) the performance of the 405B model. In this work, we measure solution quality only by correctness, without evaluating clarity, succinctness, or other factors. Therefore, we acknowledge that, Llama-3.3-70B-Instruct may not always produce the highest-quality answers on dimensions beyond correctness, and can sometimes produce the correct final answer even though the reasoning it provides is wrong.\n\n---\n\n**Q3. Clarity of the experiment section**\n\nThank you for the suggestion. In the revised manuscript, we will add separate sections for the Countdown experiments and the math-reasoning experiments.\n\n---\n\n**Q4. Line 20-21**\n\nThank you for pointing it out. We will correct the sentence to: “auto-regressive training on massive web-scale data.”\n\n---\n\n**Q5. Figure 2**\n\nYou are correct that, in an abstract search tree, modifying node 4 does not change the branch that contains node 9 and 10. However, in our LLM-driven search, the model explores nodes sequentially and every new expansion is conditioned on the full context of all nodes visited so far. Concretely, when the model reached node 9, its prompt already included the original contents of node 4. If we later revise node 4, then the prompt becomes stale. Therefore, we discard all nodes visited after the modified node–even those on different logical branches–because they were generated from an outdated context. We will clarify this in the revised manuscript.\n\n---\n\n**Q6. Figures 4-5-6-7**\n\nYes, these figures report the performance on the Countdown task. We will make this explicit in each caption.\n\n---\n\nWe again thank you for providing constructive feedback, which truly enhances the quality of our paper. We hope that our response adequately addresses your questions.\n\n---\n\n**References**\n\n[1] Yang Yue et al., Does Reinforcement Learning Really Incentivize Reasoning, arXiv 2025 \\\n[2] Andre He et al., Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening, arXiv 2025 \\\n[3] Alon Albalak et al., Big-Math: A Large-Scale, High-Quality Math Dataset for Reinforcement Learning in Language Models, arXiv 2025"}, {"review_id": "JDIWyKAQhi", "reviewer": "Reviewer_sgBD", "summary": "This work proposes Guided-ReST that integrates optimal solutions into the model’s search procedure. It leverages unsuccessful search attempts as context for each intermediate reasoning step. Then, Guided-ReST conducts finetuning on these automatically generated search traces. The experiments are evaluated on Countdown, MATH-500, and AMC23, where Guided-ReST shows improvements over baselines in terms of the same token budget.", "questions": "1. How do we define nodes for other reasoning tasks like multi-hop reasoning or coding?", "limitations": "Yes.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength: \n\n1. Incorporating the partial solution as an initial prompt to guide the model to generate more diverse and effective reasoning trajectories is a natural idea and effective in the data preparation stage.\n\n2.  Guided-ReST shows improved performance on the Countdown task with a good coverage of baselines, such as BC, SoS, ReST, and PPO\n\nWeakness:\n\n1. My major concern for this work is about the choice of the backbone models and the benchmarks. First, there is an inconsistent choice of base models without enough justification used in the Countdown task - GPT2-250M and mathematical reasoning tasks Llama-3.1-8B-it, respectively. Since the method is mainly developed for the Countdown task, the 250M model may not be evident enough to reflect the generalization of the model to other scales of the base model. \n\n2. The method itself, as acknowledged by the author in the limitation section, requires optimal step-by-step solutions, and the reasoning steps need to be easily decomposed into steps. However, the evaluation on Countdown and two mathematical reasoning tasks is not convincing to demonstrate the effectiveness of the method to other reasoning tasks. I suggest that authors include more base models and benchmarks.", "was_revised": true, "final_justification": "I thank the rebuttal from the authors. I appreciate the results from llama and observe its comparable performance with the GPT-2 model. Though the author provides justifications on how the proposed method generalize to other task domains, the evaluation from the original paper is still only conducted in a limited domain with a particular focus on the Countdown tasks. Given the factors above, I have increased my final score to 3.", "rebuttal": "Thank you for your helpful feedback, particularly regarding the choice of backbone models. We would like to address your questions below.\n\n---\n\n**Q1. The choice of the backbone models**\n\nThank you for raising this point. Following the protocol in the stream-of-search (SoS) paper [1], our Countdown experiments used a GPT-2 model with 250M parameters trained from scratch. While this gave us a clean baseline–a fully controlled model trained from scratch–we recognize that modern LLMs differ greatly from GPT-2 in scale, architecture, and tokenization. Consequently, we acknowledge your concern that findings based on a GPT-2 model trained from scratch may not readily generalize to modern LLMs.\n\nTo check how well our approach generalizes, we repeated the Countdown experiments with a more recent LLM. Because our paper already employs Llama-3.1-8B-Instruct for the mathematical reasoning experiments, we chose Llama-3.2-1B-Instruct–a model from the same Llama-3 family–for the Countdown task. We conducted the experiments with the same hyper-parameters as in the original setup, but reduced the number of epochs to 2. In addition, we reformatted the dataset into an instruction-chat style. The results are summarized below, showing that our method scales to larger, instruction-tuned backbones.\n\n| Method | Accuracy (seen) | Accuracy (unseen) |\n| --- | --- | --- |\n| SoS (GPT-2) | 0.5747 | 0.5342 |\n| SoS (Llama) | 0.5718 | 0.5391 |\n| ReST (GPT-2) | 0.6279 | 0.6003 |\n| ReST (Llama) | 0.6330 | 0.6033 |\n| Guided-ReST (GPT-2) | 0.6823 | 0.6698 |\n| Guided-ReST (Llama) | 0.6791 | 0.6659 |\n| Guided-ReST+PPO* (GPT-2) | 0.7507 | 0.7361 |\n| Guided-ReST+PPO* (Llama) | 0.7561 | 0.7420 |\n(* denotes the method presented as Ours in the paper.)\n\n---\n\n**Q2. More benchmarks**\n\nRegarding a new benchmark, We are exploring an additional experiment that applies our method to a code-editing setting. While we may not have enough time to complete this study within the rebuttal period, we will run it regardless to strengthen our paper. If the results are ready before the rebuttal period ends, we will share them with you.\n\n---\n\n**Q3. Definition of nodes in multi-hop reasoning or coding**\n\nWe provide a concrete example of nodes in multi-hop reasoning tasks, using problems from PrOntoQA [2] and MuSiQue [3], which are well-known benchmarks for multi-hop logical reasoning and question answering.\n\n**PrOntoQA**\n\nContext: \n- Each cat is a carnivore. \n- Every carnivore is not herbivorous.\n- Carnivores are mammals.\n- All mammals are warm-blooded.\n- Mammals are vertebrates. \n- Every vertebrate is an animal.\n- Animals are multicellular. \n\nQuestion: Fae is a cat. True or false: Fae is not herbivorous\n\nSolution: \n1. Fae is a cat.\n2. Cats are carnivores. \n3. Fae is a carnivore.\n4. Every carnivore is not herbivorous.\n5. Fae is not herbivorous\n\nIn this case, each reasoning step forms a node in the search trace. These nodes are connected sequentially, with each step building on the previous one, forming a sequential reasoning graph (1 &rarr; 2 &rarr; 3 &rarr; 4 &rarr; 5).\n\n**MuSiQue**\n\nQuestion: When did the people who captured Malakoff come to the region where Philipsburg is located?\n\nSolution:\n1. What is Philipsburg capital of? Saint Martin\n2. Saint Martin is located on what terrain feature? Caribbean\n3. Who captured Malakoff? French\n4. When did the French come to the Caribbean? 1625\n\nCoding problems like MBPP [4] can also be framed as search problems, where the goal is to find an optimal solution (code) that satisfies the given prompt. In this context, search traces would represent the process of iterating over potential solutions by generating a new code snippet or modifying an existing one.\n\n---\n\nAgain, we sincerely appreciate your constructive feedback, which has been very helpful in improving our paper. We hope that our response above addresses all of your questions.\n\n---\n\n**References**\n\n[1] Kanishk Gandhi et al., Stream of Search (SoS): Learning to Search in Language, COLM 2024\n[2] Abulhair Saparov and He He, Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought, ICLR 2023\n[3] Harsh Trivedi et al., MuSiQue: Multihop Questions via Single-hop Question Composition, arXiv 2021 \\\n[4] Jacob Austin et al., Program Synthesis with Large Language Models, arXiv 2021"}]

Author rebuttal:

Meta-review: This paper proposed a new fine-tuning algorithm to improve the search capability in inference time of LLMs. The topic is well motivated and the paper well written and easy to follow. There are some limitations of the proposed method: it depends on having access to optimal solutions, and the main benefits are primarily demonstrated on the Countdown tasks. Despite this, most of the reviewers found this paper and the preliminary results interesting. We recommend acceptance of this paper, but encourage the authors to add more empirical investigations w/o direct access to optimal solutions and other math / reasoning tasks in the final revision.
Decision: Accept (poster)