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

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

#### 5.3.3 Comparison of Tree Search and Sequential Revisions

As we mentioned in Section 5.3.1 , tree search typically samples multiple candidate solutions simultaneously, which allows some tree search algorithms, such as BoN, to be accelerated using parallel strategies. At the same time, the candidate solutions in tree search are relatively independent of each other, and the search process can encourage exploration during sampling, such as in MCTS. This broadens the coverage of candidate solutions. On the other hand, sequential revisions generally continue to search around a previously sampled solution, allowing for, incremental improvements, which may lead to better solutions in some cases. For example, Snell et al. (2024) show that sequential revisions with internal guidance outperform majority voting (a tree search strategy with internal guidance). However, since sequential revisions rely on previous attempts, the computational cost increases with the number of revisions made.

#### 5.3.4 Combination of Tree Search and Sequential Revisions

Tree search and sequential revisions can be used together. In tree search, using solution-level search nodes can be seen as a combination of tree search and sequential revisions ( Zhang et al., 2024c ; Zhang et al., 2024d ) . Additionally, Snell et al. (2024) combines BoN with sequential revisions by first randomly sampling N candidate solutions, then applying sequential revisions to these N solutions, and finally using a verifier to select the best solution from all of them. This combined approach achieves performance that surpasses BoN. Such results may demonstrate the potential of combining these two search strategies.

### 5.4 Speculation About the Search of o1.

In this section, we discuss the guiding signals and search strategies used during both the training-time search and inference-time search in o1.

##### Train-time search.

During training, o1 is more likely to employ tree search techniques, such as BoN or tree search algorithms, and primarily relies on external guidance. This is because the model needs to gradually enhance its reasoning capabilities during training, and tree can sample a large number of candidate solutions in parallel, efficiently providing the model with abundant high-quality training data. Furthermore, since real-time interaction is not required during training, various external environments can be accessed to validate the sampled solutions, such as executing code or verifying the accuracy of mathematical computations. This external guidance helps more accurately direct the model’s search process.

##### Test-time search.

For test-time search, o1 is more likely to use sequential revisions, combining internal guidance to continuously refine and correct its search through reflection. Through the example in the o1 blog ( OpenAI, 2024a ) , we can observe that o1’s reasoning style is closer to sequential revisions. Besides, using Tree Search for long reasoning processes can result in significant overhead. And during inference, it is difficult to rely on real-world environments for guidance, and ( Gao et al., 2023 ; Stroebl et al., 2024 ) point out that performing extensive searches based on proxy feedback, like reward models, can lead to overoptimization problem. While increasing the computations during inference, it can actually degrade performance. However, observations from o1’s blog suggest that the model’s performance continues to improve as the computational effort during inference increases. This leads us to believe that o1 primarily uses internal guidance during the inference stage. The computations during inference is mainly reflected in the length of the inference chain.

### 5.5 Scaling Law of Search

Prior to OpenAI’s release of o1, several studies explored inference-time scaling laws. Brown et al. (2024) studied BoN, dividing BoN into two stages: sampling solutions and selecting the best one via a verifier. They found that increasing the number of samples improved "coverage" (pass@1 accuracy) following a power law. Small models could achieve near 100% pass@1 accuracy on MATH with sufficient scaling. However, they noted a gap between best-of-n accuracy and pass@1, attributing it to the verifier’s limitations in identifying correct solutions.

Snell et al. (2024) analyzed the scaling of tree search and sequential revisions, showing both approaches scaled well with more computational resources. They found that, for a given computational budget, increasing model size was more effective for complex tasks, while sampling more tokens sufficed for simpler tasks. Lightman et al. (2024) examined reward signals in scaling and observed that outcome-reward-based best-of-n search plateaued as samples increased, while PRM-based best-of-n search avoided this issue.

While these studies showed performance gains with increased search computation, Gao et al. (2023) found the inverse scaling law where scaling best-of-n search could degrade performance due to distribution shift. The reward model, trained on original policy data, struggles to generalize to new policy. Stroebl et al. (2024) noted similar issues and suggested that other search algorithms, like MCTS, may also face inverse scaling challenges.

### 5.6 Challenges of Search for Reproducing o1

##### How to overcome Inverse Scaling

One approach is to reduce test-time search, as the inverse scaling phenomenon occurs primarily during large-scale search. However, this limits the scale of search. Alternatively, improving the reward model’s generalization to handle unseen states is another solution. Inspired by LLM development, this can be achieved by increasing the model’s size and training data, as validated in Gao et al. (2023) .

##### How to avoid over-thinking on simple tasks?

Not all problems require complex reasoning or search. For straightforward queries like “1+1=?”, engaging in elaborate analysis wastes computational resources and potentially introduces errors. Forcing reasoning on such problems wastes resources and causes delays. To address this, we could constrain the length of the chain of thought (CoT) using reward shaping with a length penalty. This reshaped reward balances minimizing superfluous search and solving problems effectively.

##### How to trade off tree search and sequential revisions?

Search scales across two dimensions: tree search and sequential revision. Combining both improves performance ( Zhao et al., 2024a ) , but with a fixed computational budget, the optimal allocation of resources remains unclear. This challenge is similar to balancing model size and data size under a fixed budget. Empirical scaling laws can provide guidance for resource allocation.

##### How to improve efficiency of Search?

A key challenge in scaling search is efficiency, as LLMs’ auto-regressive generation is limited by memory read-write speeds, restricting GPU utilization ( Dao et al., 2022 ) . Additionally, some tree search algorithms, like MCTS, lack inherent parallelism. Improving efficiency requires both engineering and algorithmic solutions. For example, engineering optimizations involves implementing KV-cache sharing, while algorithmic improvements can perform KV-cache compression ( Xiao et al., 2024 ; Zeng et al., 2024b ) and speculative sampling ( Yi et al., 2024 ) .

## 6 Learning

We have introduced policy initialization, which learns from human-expert data 1 1 1 policy initialization can also learn from a stronger model, which we ignores here for simplicity. , why do we still need reinforcement learning? What makes reinforcement learning essential is that the training data of reinforcement learning is unlimited which comes from the interaction with an environment. In contrast, the human-expert data is limited and expensive. Furthermore, reinforcement learning has the potential to achieve superhuman performance, since it learns from trial and error instead of human-expert data. While human-expert data captures human behavior and knowledge, reinforcement learning can lead to the discovery of strategies that humans may not be capable of. AlphaGo ( Silver et al., 2016 ) which utilized reinforcement learning, was able to defeat world-class human players in the game of Go by discovering novel strategies that were previously unknown to experts, for example the famout “37 move” of AlphaGo ( Silver et al., 2016 )

The reinforcement learning typically samples the trajectory using a policy and improves the policy based on the received rewards. In the context of o1, we hypothesize that the reinforcement learning process generates trajectories through a search algorithm rather than relying solely on sampling. One advantage of search methods is their ability to explore superior states or solutions compared to random sampling. For example, beam search prioritizes actions with the highest expected action values. Thus, search techniques can provide higher-quality training data than simple sampling. Under this assumption, the reinforcement learning of o1 could involve an iterative process of search and learning. In each iteration, the learning phase utilizes the search-generated outputs as training data to enhance the policy, while the improved policy is then applied to the search process in the next iteration. A prominent example of this search-and-learn iteration is AlphaGo Zero ( Silver et al., 2017 ) , which uses trajectory data obtained via Monte Carlo Tree Search (MCTS) for policy learning.

The train-time search is different from the test-time search. The test-time search outputs a solution with the maximum rewards or confidence in all candidate solutions. But at the training, all candidate solutions generated from search may be utilized by learning. We denote the set of state-action pairs output from search as D search D_{\text{search}} , and the set of state-action pairs in the optimal solutions from search as D expert D_{\text{expert}} . Therefore, D expert D_{\text{expert}} is a subset of D search D_{\text{search}} . We visualize the difference between D search D_{\text{search}} and D expert D_{\text{expert}} in Figure 8 .

### 6.1 Learning Methods

Given D search D_{\text{search}} , the policy can be improved using either policy gradient methods ( Sutton et al., 1999 ) or behavior cloning. This section provides a detailed overview of several key policy gradient methods. Among these, Proximal Policy Optimization (PPO) ( Sutton et al., 1999 ) and Direct Policy Optimization (DPO) ( Rafailov et al., 2023 ) are the most widely used reinforcement learning techniques for LLMs. It is also common to preform behavior cloning or supervised learning on the searched data ( Zelikman et al., 2022 ; Touvron et al., 2023 ) . The behavior cloning here differs from that in policy initialization discussed in Section 3 in that it learns the behavior derived from the search process, rather than from human expert.

#### 6.1.1 Policy Gradient

The policy gradient parameterizes the policy with θ \theta , which can be defined as π θ ​ ( a | s ) \pi_{\theta}(a|s) . θ \theta is updated with the following gradient:

∇ θ J ​ ( θ ) = 𝔼 π θ ​ [ Q ⁡ ( s t , a t ) ​ ∇ θ ​ log ⁡ π θ ​ ( a t | s t ) ] \nabla_{\theta}J(\theta)=\mathbb{E}_{\pi_{\theta}}\left[Q(s_{t},a_{t})\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})\right] (3) Q-value Q ⁡ ( s t , a t ) Q(s_{t},a_{t}) is often approximated with Monte Carlo sampling ( Sutton et al., 1999 ) or Temporal-Difference estimation ( Konda & Tsitsiklis, 1999 ) .

##### REINFORCE

REINFORCE algorithm ( Sutton et al., 1999 ) uses Monte Carlo to approximate action value: ∇ θ J ​ ( θ ) = 1 | D search | ​ ∑ ( s t , a t ) ∈ D search [ G t ​ ∇ θ ​ log ⁡ π θ ​ ( a t | s t ) ] , \nabla_{\theta}J(\theta)=\frac{1}{|D_{\text{search}}|}\sum_{(s_{t},a_{t})\in D_{\text{search}}}\left[G_{t}\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})\right], (4) where the return G t G_{t} is the discounted cumulated rewards, D search D_{\text{search}} denotes all state-action pairs in the searched data.

##### Actor-Critic

The return G t G_{t} may bring high variance for the gradient estimation. Therefore, Actor Critic ( Konda & Tsitsiklis, 1999 ) replace G t G_{t} with the advantage function or Temporal-difference error, which defined A ⁡ ( s t , a t ) = R t + 1 + γ ​ V π θ ​ ( s t + 1 ) − V ⁡ ( s t ) A(s_{t},a_{t})=R_{t+1}+\gamma V_{\pi_{\theta}}(s_{t+1})-V(s_{t}) .

The advantage function in Actor-Critic algorithm involves the estimation of value function V ⁡ ( S ) V(S) , therefore it learns a value model to approximate the value function. Another benefit of using Temporal-difference error A ⁡ ( s t , a t ) A(s_{t},a_{t}) instead of G t G_{t} is that the estimation of G t G_{t} may require more computation for sampling. For example, some trajectories in beam-search are pruned before reaching the terminal state. We need to make additional sampling to estimate G t G_{t} for these intermediate states.

##### PPO

In policy gradient methods, the policy is updated using gradient ascent. However, when the policy update is too large, it can cause drastic changes in the policy, leading to poor performance or even divergence. TRPO ( Schulman et al., 2015 ) and PPO ( Schulman et al., 2017 ) address this problem with KL divergence constraint and clipping, which ensuring that policy updates are not too large.

Another benefit of TRPO and PPO is that it improves data utilization. It is more data-efficient to update the parameters θ \theta multiple times on the same example, like the multi-epoch training in machine learning. Repeatably learning on the same data is called Replaying buffer in reinforcement learning. Replaying buffer improves the data utilization, but making the learning off-policy, where the behavior policy ( π search \pi_{\text{search}} ) that generates D search D_{\text{search}} is not the same as the target policy ( π θ \pi_{\theta} ) that is to be optimized. TRPO ( Schulman et al., 2015 ) and PPO ( Schulman et al., 2017 ) alleviates this problem with KL divergence constraint and clipping.

PPO is a variant of Actor-Critic algorithm, which has been widely used for reinforcement learning for large language model, especially the RLHF ( Ouyang et al., 2022 ; Bai et al., 2022a ) . But the training of PPO typically involves 4 models for function approximation of π θ \pi_{\theta} , π search \pi_{\text{search}} , reward, value, which is expensive. GRPO ( Shao et al., 2024 ) alleviate this issue by estimating value function with Monte Carlo estimation instead of function approximation. Remax ( Li et al., 2024c ) is also proposed to simplify the training of PPO, which use the return of greedy decoding as the baseline of reduce the variance of reinforce algorithm.

It is not trivial to reimplement the performance of PPO, Huang et al. (2022) highlight 13 core implementation details of PPO, for example, using two distinct networks to predict policy and value respectively is better than sharing their parameters. Zheng et al. (2023c) ; Wang et al. (2024a) also empirically validate the influence of some implementation details on the performance of RLHF, for example, they found that it is helpful to normalize and clip reward.

##### DPO