Scaling of Search and Learning: A Roadmap to Reproduce o1 from Reinforcement Learning Perspective — part 5/8

Cleaned-text code points [62257, 78216). Bibliography omitted here; source.md includes it.

This section introduces tree search algorithms like Best-of-N (BoN), beam search, and MCTS. BoN generates multiple independent candidate solutions but lacks dynamic adjustment of the model’s probability distribution, leading to inefficiencies, such as over-sampling high-probability options. BoN can be seen as a special case of tree search with depth-1 nodes. In contrast, other tree search strategies dynamically adjust at each step, balancing exploration and exploitation with heuristics, and can use lookahead search, backtracking, and pruning to improve efficiency and reduce sampling costs. An important question in using tree search algorithms on LLMs is defining the granularity of the tree node. As shown in Figure 7 , common granularities of tree nodes include: token level, step level, and solution level. The token level represents the finest granularity, while the solution level represents the coarsest granularity. Generally, the smaller the granularity of the search tree nodes, the deeper the search tree becomes.

##### Best-of-n Sampling

Best-of-N sampling (BoN) is a simple yet effective search method, which can be seen as a solution-level tree search. It first generates multiple candidate solutions from the model, then selects the best solution via a reward model while discarding the rest ( Cobbe et al., 2021 ) . When BoN uses oracle rewards (e.g., comparing with ground truth), its accuracy improves with more samples, as coverage increases ( Brown et al., 2024 ) . However, in most cases, oracle rewards are unavailable, making that learning a reward model becomes the primary bottleneck. Additionally, increasing the number of samples raises computational costs, further limiting BoN’s scalability.

To reduce BoN’s cost, Sun et al. (2024a) propose speculative rejection, leveraging the correlation between partial and complete sequence scores. By scoring partial sequences and discarding low-scoring solutions early, it reduces computation while achieving BoN-like performance. An alternative to multiple sampling is fine-tuning the policy model to mimic BoN’s distribution. Gui et al. (2024) theoretically prove that BoN is optimal under specific constraints and fine-tune the model using best-of-n and worst-of-n data with preference-based IPO loss and SFT loss. Similarly, Amini et al. (2024) introduce variational BoN (vBoN), which uses PPO to minimize the gap between the policy model and BoN’s distribution, achieving superior performance without repeated sampling. Sessa et al. (2024) propose BOND, which fine-tunes the model to approximate BoN’s distribution by minimizing Jeffreys divergence (a combination of forward and backward KL divergences). Through iterative distillation, BOND reduces inference to a single sample while maintaining BoN-level performance.

##### Beam Search

Beam search is a classic tree search algorithm that traditionally expands and prunes branches based on partial sequence probabilities. Beam search is typically conducted at the token level. Recent advancements in large language models have led to modifications of beam search by integrating additional guidance signals. TreeBoN ( Qiu et al., 2024 ) iteratively expands branches and prunes low-quality responses using token-level rewards from DPO ( Rafailov et al., 2024 ) , maintaining high-quality sequences while reducing computation in a manner similar to beam search, on step level. Similarly, Yu et al. (2024a) train a value model from outcome supervision to score partial sequences, replacing token probabilities. Xie et al. (2023) improve stochastic beam search ( Kool et al., 2019 ) , balancing exploitation and exploration by using a policy model for self-evaluation for sampling instead of token probabilities. Ma et al. (2023) propose a process reward model to guide greedy search, introducing backtracking when child nodes have negative rewards. Snell et al. (2024) implement lookahead search via k-step rollouts for evaluating partial sequences, which can be seen as a specia case of MCTS. Reward-guided beam search, compared to token-probability-based approaches, better aligns with downstream task performance. By incorporating value models, rollouts, or backtracking, it also balances exploitation and exploration, enhancing sampling coverage.

##### Monte Carlo Tree Search

Incontrast to BoN and traditional beam search, Monte Carlo Tree Search (MCTS) is a lookahead search algorithm used for making optimal decisions in large search spaces ( Browne et al., 2012 ) , which can choose actions based on their expected returns. Traditional MCTS use Monte Carlo estimations to estimate state values in a search tree. As simulations increase, the search tree grows, values become more accurate, and the action-selection policy improves by prioritizing higher-value children. MCTS is suitable for large-scale search spaces. Compared to traditional algorithms, it efficiently explores complex search trees using estimated action values (expected rewards). Moreover, MCTS balances exploration and exploitation, allowing the search to not only focus on the most promising areas but also explore unknown regions, thereby avoiding getting stuck in local optimal solutions. As a result, it has shown strong performance on complex reasoning tasks with large search spaces, such as Go ( Silver et al., 2016 ) .

In MCTS, the nodes of the tree represent states, and the edges represent actions. MCTS constructs a search tree through multiple MCTS simulations to estimate the value (expected reward) of candidate actions in the current state. A typical MCTS simulation consists of four stages:

• Selection : Starting from the root, the algorithm selects edges (actions) recursively based on the action values. To encourage exploration, MCTS adds an exploration term to the action values, which promotes actions that have been visited less frequently. This term is typically calculated using a variant of PUCT ( Rosin, 2011 ) . As a result, MCTS is able to explore more actions and make more accurate estimates of their values. After selecting an action, a new state is reached, which corresponds to a node in the tree. The selection process is repeated continuously until a leaf node is reached.

• Expansion : If the leaf node reached by the selection process is not a terminal node, such as not being the <eos> token in the text, MCTS will expand this node. It selects the possible actions for the leaf node’s state and the resulting states after executing those actions as the child nodes of the current node, using the probabilities of these actions as the initial action values. Typically, MCTS expands the same number of actions each time.

• Evaluation : For the node reached by the selection process, MCTS evaluates the value of the node’s state to refine the current action values. Traditional MCTS uses a more efficient policy for Monte Carlo estimation (e.g., using a smaller model), a process known as rollout, and the policy used for estimation is called the rollout policy ( Sutton & Barto, 1998 ) . Another approach is to use an additional value model to predict the value of the current state without requiring extra sampling. In practice, rollout and the value model can be used together. The advantage of this state value estimation is that MCTS, compared to traditional lookahead search algorithms, can more efficiently estimate the expected reward of the current state in the future without traversing the entire search tree, making it more suitable for problems with large search spaces.

• Backpropagation : After evaluating the value of the node’s state, MCTS performs backpropagation to update the values of the actions along the path as well as the visit counts of the actions. Through backpropagation, MCTS refines its estimates of the action values, making the estimates more accurate.

The MCTS algorithm iteratively executes four stages until reaching a predefined simulation limit, selecting the most visited action from the root node.

MCTS algorithms based on LLMs mainly differ in action granularity and reward selection. Wan et al. (2024) analyze the impact of token-level and step-level actions on the search space. Step-level actions simplify the search process by reducing tree width but face challenges with large action spaces due to sentence diversity. In contrast, token-level actions narrow the action space but increase tree depth. Some approaches also adopt solution-level actions, where nodes represent entire solutions, and actions modify these solutions.

• Token-level actions: Liu et al. (2024c) use a token-level action space, where the LLM predicts the next token as the MCTS action, with the PPO value model evaluating states. Choi et al. (2023) employ token-level MCTS to ground generation in reference knowledge, reducing hallucinations. Zhang et al. (2023b) use token-level MCTS to guide Transformers in generating better programs, evaluating node values via rollouts. Accurately assessing token values is critical, but sparse reward signals pose challenges in deriving token-level value or reward models. Additionally, the deep search tree leads to efficiency issues, making the design of an efficient inference framework essential for token-level MCTS.

• Step-level actions: For multi-step reasoning tasks, step-level actions are a natural choice. Hao et al. (2023) propose RAP, an MCTS-based planning algorithm, which defines actions based on task-specific steps, such as posing sub-questions in math problems or moving blocks in Blocksworld ( Valmeekam et al., 2023 ) . RAP uses an LLM to predict future states instead of observing state transitions. Chen et al. (2024a) treat partial solutions as states and reasoning steps as actions for mathematical problems. Putta et al. (2024) ; Zhou et al. (2024) use browser interaction commands as actions in WebShop, a simulated e-commerce platform. Unlike task-specific steps, Qi et al. (2024) design five human-like reasoning actions for MCTS. To achieve o1-style reasoning, we may need to manually design actions like reflection and error correction, which should be incorporated during policy initialization.

• Solution-level actions: Solution-level MCTS treats the complete solution as a node state and modifications as actions. Zhang et al. (2024c) propose MCTSr, where nodes represent different answer versions, actions involve self-refinement, and rewards come from self-evaluation. Starting with a naive model-generated answer, MCTSr combines self-refinement with MCTS to achieve GPT-4-level performance on math problems. Building on MCTSr, Zhang et al. (2024d) use a pairwise preference reward model to evaluate whether a modified answer improves, estimating the new node’s value. Yan et al. (2024) use self-improvement as actions and design a consistency-based method to ensure factual correctness.

##### Other Tree Search Algorithms.

In addition to MCTS and Beam Search, traditional graph search algorithms like DFS, BFS, and A ∗ A^{*} are also used. The Tree-of-Thought (ToT) framework ( Yao et al., 2023a ) models reasoning as a tree, where nodes represent reasoning steps and branches represent continuations. By expanding and evaluating multiple reasoning solutions, ToT explores a broader solution space, introducing mechanisms like search, reflection, and backtracking to reconsider steps and explore alternatives. ToT employs DFS and BFS for its search process. Koh et al. (2024) propose an A ∗ A^{*} -inspired best-first search, using a multimodal LLM to evaluate nodes and a heuristic to select or backtrack. Lehnert et al. (2024) use A ∗ A^{*} to collect execution traces, training an encoder-decoder model to imitate these traces.

#### 5.3.2 Sequential Revisions

Compared to tree search, sequential revisions primarily conducts the search by iteratively refining the previous answer. The key feature of sequential revisions is that it generates improved answers based on reflections on previous answers or changes in the environment. Sequential revisions requires the model to have basic abilities for self-reflection and error correction, which can be introduced during the poicy initialization through SFT or Prompting, as discussed in Section 3 .

Sequential revisions can directly rely on the internal guidance such as self-evaluation. Madaan et al. (2023) propose to SELF-REFINE. It first generates an initial answer using a LLM, and then utilize the same LLM to iteratively provide feedback for its output and refine its output based on the self feedback. The iteration process stops when the maximum number of iterations is reached or when the model determines that it can stop. Snell et al. (2024) also uses a model with revision capabilities to perform revisions during inference, demonstrating that the accuracy of the answer improves as the number of revisions increases. Other work rely on observation or feedback from external environment. Both Yao et al. (2023b) and Shinn et al. (2023) are based on the model reflecting on or improving the next action based on external feedback observed after taking an action. Chen et al. (2024c) and Gou et al. (2024) use feedback based on the execution of generated code to allow the model to modify the generated code or debug.

There is still debate over whether Sequential Revisions is truly effective. For example, Huang et al. (2024a) argue that large models cannot self-correct properly without external feedback. However, the opposing viewpoint suggests that due to the existence of the DG-gap ( Leike, 2022 ) , large models may have a stronger ability to discern and improve the answers they have already generated, allowing for further refinement. Chen et al. (2024d) conduct empirical studies across a range of tasks and found that Sequential Revisions only achieves performance improvements over simpler methods like BoN when the accuracy of the discriminator (guidance) ≥ \geq 90%.

Paper Search Guidance Search Strategies Internal Guidance External Guidance Tree Search Sequential Revisions Math Verifier ( Cobbe et al., 2021 ) ✗ Env Feedback Best-of-N ✗ Self-consistency ( Wang et al., 2023a ) Model Uncertainty ✗ Best-of-N ✗ Speculative Rejection ( Sun et al., 2024a ) ✗ Env Feedback Best-of-N ✗ BoN ( Gui et al., 2024 ) ,vBoN ( Amini et al., 2024 ) ,BOND ( Sessa et al., 2024 ) ✗ Env Feedback Best-of-N † ✗ TreeBoN ( Qiu et al., 2024 ) ✗ Env Feedback Beam Search ✗ OVM ( Yu et al., 2024a ) Value Function Value Function Beam Search ✗ Xie et al. (2023) Self-evaluation ✗ Beam Search ✗ Ma et al. (2023) ✗ Env Feedback Beam Search ‡ ✗ Snell et al. (2024) ✗ Env Feedback Beam Search ✗ TS-LLM ( Wan et al., 2024 ) ✗ Env Feedback MCTS ✗ PPO-MCTS ( Liu et al., 2024c ) Value Function Value Function MCTS ✗ KCTS ( Choi et al., 2023 ) ✗ Env Feedback MCTS ✗ Zhang et al. (2023b) ✗ Env Feedback MCTS ✗ RAP ( Hao et al., 2023 ) Self-evaluation Heuristic Rules MCTS ✗ AlphaMath ( Chen et al., 2024a ) Value Function Value Function MCTS, Beam Search ✗ AgentQ ( Putta et al., 2024 ) Self-evaluation ✗ MCTS ✗ LATS ( Zhou et al., 2024 ) Self-evaluation, Model Uncertainty ✗ MCTS ✗ rStar ( Qi et al., 2024 ) Model Uncertainty ✗ MCTS ✗ MCTSr ( Zhang et al., 2024c ) Self-evaluation ✗ MCTS ✓ LLaMA-Berry ( Zhang et al., 2024d ) ✗ Env Feedback MCTS ✗ Mirror ( Yan et al., 2024 ) ✗ Heuristic Rules MCTS ✓ ToT ( Yao et al., 2023a ) ✗ Env Feedback, Heuristic Rules DFS,BFS ✗ Koh et al. (2024) ✗ Env Feedback A ∗ A^{*} ✗ Beyond A ∗ A^{*} ( Lehnert et al., 2024 ) ‡ ✗ Heuristic Rules A ∗ A^{*} ✗ Self-refine ( Madaan et al., 2023 ) Self-evaluation ✗ ✗ ✓ Snell et al. (2024) Self-evaluation ✗ ✗ ✓ ReACT ( Yao et al., 2023b ) ✗ Env Feedback ✗ ✓ Reflexion ( Shinn et al., 2023 ) ✗ Env Feedback ✗ ✓ Self-debug ( Chen et al., 2024c ) ✗ Env Feedback ✗ ✓ Critic ( Gou et al., 2024 ) ✗ Env Feedback ✗ ✓

As a summary, we list the research work in Table 2 , along with the specific categories of guiding signals and search strategies used, for the reader’s reference.