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

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

DPO eliminate the cost of modeling value model and reference policy, meanwhile having low variance in reward/value estimation. The idea of DPO is to reparameterize reward function with the optimal policy π ∗ \pi^{*} , and transform the optimization of Bradley-Terry reward modeling to the optimization of the policy. The loss function of DPO is: ℒ DPO ​ ( π θ , π ref ) = − 𝔼 ( x , y w , y l ) ∼ 𝒟 search ​ [ log ⁡ σ ⁡ ( β ​ log ⁡ π θ ​ ( y w | x ) π ref ​ ( y w | x ) − β ​ log ⁡ π θ ​ ( y l | x ) π ref ​ ( y l | x ) ) ] \mathcal{L}_{\text{DPO}}(\pi_{\theta};\pi_{\text{ref}})=-\mathbb{E}_{(x,y_{w},y_{l})\sim\mathcal{D_{\text{search}}}}\left[\log\sigma\left(\beta\log\frac{\pi_{\theta}(y_{w}|x)}{\pi_{\text{ref}}(y_{w}|x)}-\beta\log\frac{\pi_{\theta}(y_{l}|x)}{\pi_{\text{ref}}(y_{l}|x)}\right)\right] (5)

where x x denotes the question, y w y_{w} denotes the positive solution and y l y_{l} denotes the negative solution. The parametrization of DPO can derive a kind of reward shaping: f ⁡ ( r , π ref , β ) ​ ( x , y ) = r ⁡ ( x , y ) − β ​ log ​ ∑ y π ref ​ ( y | x ) ​ exp ⁡ ( − r ⁡ ( x , y ) β ) f(r,\pi_{\text{ref}},\beta)(x,y)=r(x,y)-\beta\log\sum_{y}\pi_{\text{ref}}(y|x)\exp\left(-\frac{r(x,y)}{\beta}\right) (6) Here, the term β ​ log ​ ∑ y π ref ​ ( y | x ) ​ exp ⁡ ( − r ⁡ ( x , y ) β ) \beta\log\sum_{y}\pi_{\text{ref}}(y|x)\exp\left(-\frac{r(x,y)}{\beta}\right) serves as the baseline for r ⁡ ( x , y ) r(x,y) , which explains why DPO does not encounter the high variance issue associated with the REINFORCE algorithm. Estimating this baseline is computationally expensive due to the large action space, but DPO eliminates this estimation by leveraging the Bradley-Terry model.

DPO has been widely used as the alternative for PPO in RLHF ( Dubey et al., 2024 ) . Even though the search does not provide the preference data directly, we can still construct preference pairs from D search D_{\text{search}} . LLama-3 ( Dubey et al., 2024 ) uses the positive and negative examples from rejected sampling as the preference data for DPO. MCTS-DPO ( Xie et al., 2024 ) and SVPO ( Chen et al., 2024b ) collect a number of step-level preference pairs from state-action pairs from MCTS, and use the preference data for DPO.

#### 6.1.2 Behavior Cloning

Behavior cloning, or supervised learning, employs the expert policy as the target and seeks to reduce the discrepancy between the policy π θ \pi_{\theta} and the expert policy π expert \pi_{\text{expert}} . While π expert \pi_{\text{expert}} may not represent a true expert policy, it is generally more effective than π θ \pi_{\theta} . Furthermore, π expert \pi_{\text{expert}} improves with increased computational effort, indicating that, with sufficient resources, π expert \pi_{\text{expert}} could approach the performance of a true expert policy. For LLM, the loss function of behavior cloning is cross-entropy loss. min θ − 𝔼 ( s , a ) ∼ π expert ​ [ log ⁡ π θ ​ ( a | s ) ] \text{min}_{\theta}-\mathbb{E}_{(s,a)\sim\pi_{\text{expert}}}\left[\log\pi_{\theta}(a|s)\right]

D search D_{\text{search}} contains all state-action pairs from search, including the negative solution, therefore it can not be taken as the expert data. In contrast, the state-action pairs in D expert D_{\text{expert}} receive the highest rewards compared with other pairs. Therefore, D expert D_{\text{expert}} can be taken as the expert data. The expert data from search may not really have a high quality, but it is better than the data sampled from the current policy. And as we scale the compuation of search, the expert data generated from search becomes more optimal. For example, Gao et al. (2023) found scaling best-of-n search enlarges the distance between policy of the best-of-n solution and the current policy.

min θ − 1 | D expert | ​ ∑ ( s , a ) ∈ D expert [ log ⁡ π θ ​ ( a | s ) ] , \text{min}_{\theta}-\frac{1}{|D_{\text{expert}}|}\sum_{(s,a)\in D_{\text{expert}}}\left[\log\pi_{\theta}(a|s)\right], (7)

The approach of iterating search and behavior cloning is called Expert Iteration in Anthony et al. (2017) . AlphaGo Zero ( Silver et al., 2017 ) is a well-known example of expert iteration, which uses the policy of MCTS for behavior cloning.

Although Expert iteration uses behavior cloning which can be seen as the supervised learning, it is different from policy initialization introduced in Section 3 , where the expert data comes from human expert or a stronger model. The expert data of expert initialization comes from the interaction with the environment, i.e., search.

A well-known example of expert iteration for large language model is STaR ( Zelikman et al., 2022 ) that uses reject sampling as the search algorithm to filter the examples generated by LLM with incorrect answers and finetune LLM on the remaining examples. LLama-2 ( Touvron et al., 2023 ) utilizes the method of STaR for RLHF and found it achieve the comparable performance with PPO. V-Star ( Hosseini et al., 2024 ) follows Star to use reject sampling to filter high-quality examples, but they train a verifier using DPO to provide reward for reject sampling rather than relying on the ground-truth answer. Their verifier also works at the test-time, by applying on best-of-n search. Wan et al. (2024) ; Tian et al. (2024a) ; Chen et al. (2024a) ; Zhang et al. (2024b) follows AlphaGo Zero to use MCTS as the search algorithm and finetune LLM on the solutions with the highest or the top-k highest reward in the tree. The main difference of their work between STaR is that they used different search algorithm.

#### 6.1.3 Speculation about the learning of o1

We make some comparison on the PPO, DPO and behavior cloning to see which one is more likely to be used for o1. The differences between the three methods in terms of memory cost and data utilization are summarized in Table 3 .

##### Memory Cost

PPO requires storing reward function, value function and reference policy in memory, which is expensive. While DPO eliminates the reward model and value model, which is simpler and more memory-efficient than PPO. But DPO is based on Bradley-Terry model and requires the preference data. For some environment, where the ground truth is available, it would be better to learn a policy with PPO using the reward signal, instead of learning from preference data. Behavior cloning even does not need reference policy, thus is the most memory-effient among three learning methods.

##### Data Utilization

The difference among the 3 learning methods is not only on the learning algorithm, but also on the training data. PPO and DPO use all state-action pairs from search ( D search D_{\text{search}} ), even those with negative rewards. While behavior cloning adopts the subset of state-action pairs from search with high rewards ( D expert D_{\text{expert}} ). Therefore, PPO and DPO have better data utilization than behavior cloning, since the negative action or solution can also provide useful signal for improving policy.

Havrilla et al. (2024) empirically compare the performance between PPO ( Schulman et al., 2017 ) and behavior cloning on GSM8K and MATH, and they find the behavior cloning (expert iteration) consistently outperforms policy gradient. It requires more empirical comparison on these methods.

It is likely that the learning of o1 results from the combination of multiple learning methods. In this framework, we hypothesize that the learning process for o1 begins with a warm-up phase using behavior cloning, followed by a transition to PPO or DPO once the improvements from behavior cloning plateau. This approach is grounded in the idea that behavior cloning is more efficient than PPO or DPO, thereby accelerating the warm-up phase. However, behavior cloning is limited in that it only learns from the highest-reward solutions and disregards negative solutions. Consequently, further optimization may require the use of PPO or DPO, which provide better data utilization. This pipeline is consistent with the post-training strategies employed in LLama2 ( Touvron et al., 2023 ) and LLama3 ( Dubey et al., 2024 ) .

### 6.2 Scaling Law of Reinforcement Learning

The relationship between the loss, computational cost, model parameters, and data size during the pre-training phase of large models follows a power-law scaling law. But does reinforcement learning (RL) in large language models also exhibit a scaling law? The blog of OpenAI ( OpenAI, 2024a ) shows that there is a log-linear scaling law between reasoning performance and train-time compute. However, aside from this, there has been little research on the scaling laws related to reinforcement learning.

OpenAI has also investigated the scaling laws in reinforcement learning ( Hilton et al., 2023 ) , their experiments focused on traditional reinforcement learning tasks, such as Dota 2, and did not involve LLMs. They found that the performance of reinforcement learning is improved with the increase of the number of model parameters and interactions with environment, which follows power law. They were able to derive the optimal number of model parameters and interactions with environment under a fixed computational budget.

Tuyls et al. (2024) studied the scaling laws of imitation learning in traditional reinforcement learning tasks, such as Atari. They also observed a power-law relationship between the imitation-learning loss and scale factors, including model size, data size, and computation budget. However, the data used for imitation learning came from expert annotations, rather than being generated through interactions with environment.

Some studies on reinforcement learning with large models have experimentally explored the relationship between LLM performance and the number of iterations ( Zelikman et al., 2022 ; Xie et al., 2024 ) . However, these studies only conducted small-scale validations, showing that scaling the number of iterations is effective, but they did not establish scaling laws that could be used to predict results for larger-scale experiments. In order to conduct large-scale reinforcement learning like o1, it is crucial to investigate the scaling laws of reinforcement learning on LLMs.

### 6.3 Challenges of Learning for Reproducing o1

##### How to improve Training Efficiency?

The primary bottleneck in training efficiency arises from the train-time search process, as the time required for LLM generation on the same batch exceeds that of training, with search being particularly slow. For instance, in the open-source project MCTS-DPO ( Xie et al., 2024 ) , the majority of the training time is consumed by the MCTS search, leading to training times of up to one week on an A800 GPU using the MATH dataset. There are two potential strategies for accelerating training: 1. Improving the search algorithm and implementation, which has been discussed in the search section and will not be reiterated here. 2. Extending learning beyond the data generated by online search to include data from previous search iterations. While reusing data from previous iterations may introduce issues related to off-policy learning, it increases data utilization and consequently reduces the search scale.

##### How to learn a Strong Question Generator.

We have introduced how the solutions generated by search can be used for learning. But the data for learning includes not only the solution but also question. As the policy of LLM improves, the challenging questions become simple. Therefore, it may be important to update the question or initial states. For example, the new questions can be more challenging or explore the new domains. Wizzrd-LM Xu et al. (2024) proposes to prompt a strong LLM like chatgpt to generate a more challenging question or another question that is related to the given question. But the generation of new question is independent from the performance of the policy to be optimized. The problem of updating questions is also related to Curriculum Learning ( Wang et al., 2022a ) , especially automatic curriculum learning ( Kumar et al., 2010 ) . Generating questions conditioned on the policy for LLM is challenging, since the generated questions may not be suitable or even unsolvable. Learning on these questions are not helpful for improving policy of LLM.

##### How to Narrow the Distribution Shift in off-policy learning?

The solutions generated by search is typically better than the data sampled from the current policy. The result is that the data produced by search can be considered as originating from a better policy. Using search-generated data for policy gradient training constitutes off-policy learning. To mitigate the distribution shift of policy in off-policy learning, one straightforward approach is to limit the scale of the search. The smaller the search scale, the less pronounced the distribution-shift issue becomes. Alternatively, we can implement search with sampling from the current policy, then there is no distribution shift problem.

An alternative approach is to incorporate off-policy learning methods. While techniques such as importance sampling and KL divergence constraints, employed in TRPO ( Schulman et al., 2015 ) and PPO ( Schulman et al., 2017 ) , are effective. However, they require knowledge of the policy probability associated with the search data, denoted as π search \pi_{\text{search}} , which is unavailable.

Another approach to make use of behavior cloning to turn the off-policy learning to on-policy learning. We can first apply behavior cloning on searched data D search D_{\text{search}} to narrow the gap between π θ \pi_{\theta} and π search \pi_{\text{search}} . After completing the behavior cloning, policy gradient training can then be conducted on the searched data. The pipeline of policy gradient following behavior cloning occurs at each iteration, which is different the speculation introduced in 6.1.3 , where we perform behavior cloning at early iterations and PPO at the remaining iterations. Furthermore, this approach can be supplement for our speculation: We perform behavior cloning at the warm-up stage, and then combine behavior cloning and PPO at each iteration.

## 7 Open-source o1 Project

Although o1 has not published a technical report, the academic community has made available several open-source implementations of o1. All of these implementations can be viewed as components or special cases of the o1 framework introduced in this paper. In Table 4 , we summarize the methods adopted by these open-source projects in the areas of policy initialization, reward design, search, and learning.

There are also some o1-like models from industry, for example, k0-math, skywork-o1 ( o1 Team, 2024 ) , Deepseek-R1 ( Shao et al., 2024 ) , QwQ ( Team, 2024 ) and InternThinker ( Cai et al., 2024 ) . We do not discuss about these models, since they have not open sourced their techniques.

### g1

g1 ( bklieger, 2024 ) may be the earliest project that attempts to reimplement o1, they take the approach of prompt engineering. They prompt LLM to self-reflect and propose multiple solutions to clone the behavior of o1.

### Thinking Claude

Thinking Claude ( richards, 2024 ) works similarly to g1, which prompts LLM with more complex and fine-grained actions, like problem analysis and progress tracking. Both g1 and Thinking Claude can reshape the behavior of LLM to be like that of o1, but they have not validated their prompts on reasoning benchmarks.

### Open-o1

Open-o1 ( Open-Source-O1, 2024 ) provides a SFT dataset, where each response contains long cot. But it is not clear where does these data come from. We suspect that these data are coming from human experts or a strong LLM. Open-o1 found that training LLama-3-8B and Qwen-7b on their dataset can not only shape the style of model response to be like o1, but also improve the model performance on reasoning benchmarks.

### o1 Journey

The o1 journey ( Qin et al., 2024 ; Huang et al., 2024b ) is outlined in two technical reports.