Learning to Better Search with Language Models via Guided Reinforced Self-Training

Thank you for your helpful feedback, particularly regarding the choice of backbone models. We would like to address your questions below.

---

Q1. The choice of the backbone models

Thank 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.

To 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.

---

Q2. More benchmarks

Regarding 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.

---

Q3. Definition of nodes in multi-hop reasoning or coding

We 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.

PrOntoQA

Context:

• Each cat is a carnivore.
• Every carnivore is not herbivorous.
• Carnivores are mammals.
• All mammals are warm-blooded.
• Mammals are vertebrates.
• Every vertebrate is an animal.
• Animals are multicellular.

Question: Fae is a cat. True or false: Fae is not herbivorous

Solution:

1. Fae is a cat.
2. Cats are carnivores.
3. Fae is a carnivore.
4. Every carnivore is not herbivorous.
5. Fae is not herbivorous

In 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 → 2 → 3 → 4 → 5).

MuSiQue

Question: When did the people who captured Malakoff come to the region where Philipsburg is located?

Solution:

1. What is Philipsburg capital of? Saint Martin
2. Saint Martin is located on what terrain feature? Caribbean
3. Who captured Malakoff? French
4. When did the French come to the Caribbean? 1625

Coding 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.

---

Again, 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.

---

References

[1] Kanishk Gandhi et al., Stream of Search (SoS): Learning to Search in Language, COLM 2024 [2] Abulhair Saparov and He He, Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought, ICLR 2023 [3] Harsh Trivedi et al., MuSiQue: Multihop Questions via Single-hop Question Composition, arXiv 2021
[4] Jacob Austin et al., Program Synthesis with Large Language Models, arXiv 2021

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.

---

Q1. Reliance on ground-truth solutions

We 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.

1. Can outcome-level rewards alone improve an LLM's in-context search capabilities?
2. If outcome rewards are insufficient, what additional supervision is needed?

Our 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].

While 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.

In 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:

1. Token efficiency
   • 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).
   • In typical reasoning tasks, the depth is much greater and each step can be multiple sentences, so the entire search trace grows rapidly.
2. Branching factor
   • In Countdown, the branching factor is fixed and relatively small.
   • In natural language reasoning tasks, the branching factor is unbounded.
3. Existence of oracle
   • 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.
   • In general reasoning tasks, these verification and generation steps can be noisy.

---

Q2. Use of Llama-3.3-70B-Instruct

We 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.

---

Q3. Clarity of the experiment section

Thank you for the suggestion. In the revised manuscript, we will add separate sections for the Countdown experiments and the math-reasoning experiments.

---

Q4. Line 20-21

Thank you for pointing it out. We will correct the sentence to: “auto-regressive training on massive web-scale data.”

---

Q5. Figure 2

You 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.

---

Q6. Figures 4-5-6-7

Yes, these figures report the performance on the Countdown task. We will make this explicit in each caption.

---

We again thank you for providing constructive feedback, which truly enhances the quality of our paper. We hope that our response adequately addresses your questions.

---

References

[1] Yang Yue et al., Does Reinforcement Learning Really Incentivize Reasoning, arXiv 2025
[2] Andre He et al., Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening, arXiv 2025
[3] Alon Albalak et al., Big-Math: A Large-Scale, High-Quality Math Dataset for Reinforcement Learning in Language Models, arXiv 2025

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.

---

Q1. Extension to general reasoning tasks

Yes. 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].

• The root node is the original prompt.
• Each node expansion appends a single reasoning step.

Thus, a single root-to-leaf path in the search tree corresponds to one complete CoT.

However, there might be some practical issues to apply this framework general-purpose reasoning tasks.

1. Token efficiency
   • 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).
   • In typical reasoning tasks, the depth is much greater and each step can be multiple sentences, so the entire search trace grows rapidly.
2. Branching factor
   • In Countdown, the branching factor is fixed and relatively small.
   • In natural language reasoning tasks, the branching factor is unbounded.
3. Existence of oracle
   • 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.
   • In general reasoning tasks, these verification and generation steps can be noisy.

At 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.

We will add a separate "Limitations and Future Work" section to the Appendix in the revised manuscript.

---

Q2. Contribution limited to data generation scheme

Thank 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

1. Superior standalone performance (Section 5.2, Figure 3): The subgoal-augmented traces generated by our method consistently outperform the previous RFT baseline (ReST).
2. 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.

These 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 → RL pipeline.

We 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.

To 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.

While 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.

We will highlight the synergy with RL as our primary contribution in both the Introduction and Results sections of the revised manuscript.

---

Q3. Clarity of Section 3

We 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.

1. Expansion – creating child nodes
   • Countdown: choose two remaining numbers and an operator; the child node is the updated set of numbers after the operation.
   • Mathematical reasoning: generate one paragraph of reasoning; the child node is the concatenation of all previously generated paragraphs with this new paragraph.
2. Selection – moving to the next child node for exploration
3. Verification – checking a leaf node
   • Countdown: when only one number remains, verify that it matches the target.
   • Mathematical reasoning: when the solver produces an answer, verify that it is correct.

Thus, a search trace is a sequential text record of expansion, selection, and verification lines in the exact order they occur during the search.

In 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.

Returning 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.

In 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.

---

Q4. Other fine-tuning methods

Our 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.

---

Once 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.

---

References

[1] Violet Xiang et al., Towards System 2 Reasoning in LLMs: Learning How to Think With Meta Chain-of-Thought, arXiv 2025
[2] Yang Yue et al., Does Reinforcement Learning Really Incentivize Reasoning, arXiv 2025
[3] Andre He et al., Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening, arXiv 2025

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.

---

Q1. Reliance on exact optimal traces

We 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.

1. Can outcome-level rewards alone improve an LLM's in-context search capabilities?
2. If outcome rewards are insufficient, what additional supervision is needed?

Our 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.

We 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.

---

Q2. Actual latency and GPU-hours

First, 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.

Regarding PPO, its computational cost is much smaller than a single iteration of ReST or Guided-ReST.

• For data generation, PPO only uses 25K training samples, while ReST or Guided-ReST uses 200K training samples.
• 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.
  • 1 forward for computing old_log_prob
  • 1 forward for computing ref_log_prob
  • 1 forward for computing value
  • 2 forward and backward for updating actor and critic * 4 PPO epochs

Therefore, PPO introduces only a marginal overhead to the overall compute budget.

For 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.

---

Q3. More analysis on accuracy and trajectories

Could 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.

ReST 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.

Regarding trajectory length, the table below reports the quantile statistics of the search traces for 10,000 test samples with unseen targets.

Guided-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.

Thank you for the suggestion and we will incorporate these analyses into the main paper.

---

Q4. Link between trace quality and accuracy

Thank 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:

1. 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.
2. 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.

We 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.

We also report the performance of Guided-ReST under each strategy with 95% confidence intervals.

These results highlight that high-quality, efficient search traces are crucial for optimal performance.

---

Q5. Confidence intervals

We 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.

---

Q6. Large model on different task

As 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.

Regarding 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.

---

Once 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.

---

References

[1] Yang Yue et al., Does Reinforcement Learning Really Incentivize Reasoning, arXiv 2025
[2] Andre He et al., Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening, arXiv 2025