Scaling of Search and Learning: A Roadmap to Reproduce o1 from Reinforcement Learning Perspective — part 4/8
Cleaned-text code points [46861, 62257). Bibliography omitted here; source.md includes it.
An intuitive way to build a reward signal for a general task is through an ensemble rewards in specfic domains. Quan (2024) trains the reward model as a form of MoE ( Jiang et al., 2024a ; Zeng et al., 2024a ) , where each expert is trained on a different task to provide the corresponding reward signal, and the outputs are aggregated finally. Nguyen et al. (2024) frames this problem as a multi-armed bandit question, learning to select the most suitable reward model. In short, the crucial question is for these methods is how to combine the reward signals effectively.
#### 4.5.2 One for all: World Model
The world model can not only provide the reward signal but also predict the next state ( Ha & Schmidhuber, 2018 ; Bar et al., 2024 ) . Some works claim that a video generator is a world model since the video generator can predict the images in the future time steps ( Liu et al., 2024b ; Bar et al., 2024 ; Wang et al., 2024d ; Brooks et al., 2024 ) . Dawid & LeCun (2023) also proposes a framework where world model do not need to predict the next state, instead they predict the representation of the next state, which is easier and more efficient than predicting a image. This is inline with Muzero ( Schrittwieser et al., 2020 ) , which also let the environment simulator to predict the representation of the next state in the game of go. The current works of world model focus on modeling of the next state prediction, but we believe modeling the reward signal is also critical and challenging for an agent to accomplish real-environment tasks.
## 5 Search
For LLMs, performing random sampling during generation has become a mainstream method to improve output quality, with techniques such as nucleus sampling being prominent examples ( Holtzman et al., 2020 ) . Furthermore, many studies ( Kulal et al., 2019 ; Chen et al., 2021 ) have observed that the pass@k metric improves consistently as the number of model samples increases. Brown et al. (2024) show that even small models can outperform large models leveraging search. This suggests that language models have the potential to explore the correct solution by more sampling during inference, which requires consuming more inference-time computations. Search refers to the process of finding the correct solution through multiple attempts or strategic exploration based on certain guidance, such as rewards or heuristic rules. Well-known inference strategies like self-consistency ( Wang et al., 2023a ) and Best-of-N (BoN) ( Cobbe et al., 2021 ) both can be seen as search methods. For model like o1, which are designed to solve complex reasoning tasks, search may play an important role in both training and inference processes. In this section, we first analyze the possible role of search in o1, and then introduce the details of various promising search methods which have the potential to be used in the construction of o1-like models.
### 5.1 The role of Search in o1
Search often relies on guiding signals thus can be viewed as a process of strategy iteration, we call it search policy. Compared to naive sampling, search policies are usually more likely to find better solutions. Solutions generated by search policies can either be directly used as the final output or incorporated into training algorithms, like Expert Iteration ( Anthony et al., 2017 ) , to iteratively improve the policy. We believe that search plays a crucial role in both the training and inference processes of o1 . We will refer to the search during these two stages as training-time search and test-time search , respectively.
In the training phase, the trial-and-error process in online reinforcement learning can also be seen as a search process ( Sutton & Barto, 1998 ) , where the agent performs naive sampling based on its own policy and learns solution that yield high rewards. However, since o1 involves longer inference lengths and includes human-like reasoning behaviors, the search space becomes large, and naive sampling may become inefficient. Therefore, some advanced search strategies are needed to more efficiently explore better solutions and use them as training data to update the policy model. This process can be iteratively conducted during training. In the inference phase, o1 demonstrates that increasing compute during inference with more time spent thinking can continuously improve model performance ( OpenAI, 2024a ) . We argue that o1’s way to think more can be seen as a kind of search, using more inference-time computation to find better answers.
The two key aspects of search are the guiding signals for the search and the search strategies to get candidate solutions. Search strategies are used to obtain candidate solutions or actions, while guiding signals are used to make selections. We first discuss the guiding signals for the search process, dividing them into internal and external guidance. Inspired by Snell et al. (2024) , we categorize the search strategies into tree search and sequential revisions. It is worth noting that these two classification dimensions are orthogonal, for example, tree search methods can utilize either internal or external guiding signals. We present a schematic diagram of these categories in Figure 6 .
### 5.2 Search Guidance
In the following sections, we will discuss methods for providing guiding signals to the search process. Search based on internal guidance does not rely on real-world feedback from the external environment or proxy models, but instead uses certain states or evaluation capabilities of the model itself to guide the search process. Classic text generation decoding algorithms, such as greedy decoding and beam search, typically use the probability of tokens or sequences as internal guidance for the search process. External guidance is generally independent of specific policies and relies solely on environment- or task-related signals to guide the search process.
#### 5.2.1 Internal Guidance
##### Model Uncertainty
Model uncertainty ( Malinin & Gales, 2018 ) is a common form of internal guidance. Many studies leverage it to select high-quality responses from candidates, with self-consistency ( Wang et al., 2023a ) being a notable example. Self-consistency uses majority voting or weighted sums to choose the answer with the lowest uncertainty. Universal self-consistency ( Chen et al., 2023 ) extends this approach to free-form responses, enabling large language models to select the most consistent answer without task-specific constraints, thus broadening its applicability. Building on the idea that different sentences can express the same meaning, Kuhn et al. (2023) shift the focus to semantic rather than syntactic uncertainty. They use NLI models ( Williams et al., 2018 ) for semantic clustering, treating two sentences as equivalent if they exhibit bi-directional entailment. The entropy of the semantic distribution is then used to measure uncertainty. This semantic entropy has also been applied to detect hallucinations ( Farquhar et al., 2024 ) , showing strong performance. Model uncertainty is often derived through unsupervised methods, making it easy to be obtained. However, its reliability depends heavily on the model’s calibration ( Guo et al., 2017 ) . A well-calibrated model is essential for effectively utilizing model uncertainty.
##### Self-evaluation
Model uncertainty is a useful guide but does not directly reflect response accuracy. For example, high-uncertainty answers can still be correct. To improve model performance, self-evaluation aims to let models assess their own outputs, leveraging the assumption that evaluation is easier than generation, known as the generator-discriminator gap (DG-gap) ( Leike, 2022 ; Davis et al., 2024 ) . A key application of the DG-gap is reinforcement learning from human feedback (RLHF) ( Stiennon et al., 2020 ) . Compared to directly annotating the correct model output, researchers have found that human annotators’ preferences for different candidate outputs can serve as effective and scalable supervisions. This approach helps align the model’s behavior with human intentions and values. For instruction-following language models, LLM-as-a-Judge ( Gu et al., 2024 ) is an efficient self-evaluation method, achieving high agreement with human evaluators on MT-Bench ( Zheng et al., 2023b ) . Using this approach, Yuan et al. (2024) demonstrate that models can improve self-rewarding through iterative DPO training. Similarly, Wu et al. (2024a) introduce LLM-as-a-Meta-Judge to enhance self-evaluation. The main challenge in task-specific self-evaluation is determining whether a DG-gap exists and how to fully exploit it for greater accuracy. It is important to note that some studies ( Huang et al., 2024a ) suggest that models cannot accurately evaluate answers without feedback. To address this, we may need to scale up the model size, train it for self-evaluation, specify more detailed evaluation criteria.
#### 5.2.2 External Guidance
##### Environmental Feedback
Using environmental feedback corresponding to downstream tasks is one of the most commonly used forms of external guidance and is typically directly related to the evaluation metrics of downstream tasks. Reward is a typical form of environmental feedback and is often used to guide the search process, directly corresponding to the final performance. We have discussed details of reward design in Section 4.2 . During the search process, we may need multi-scale rewards for guidance, such as outcome rewards and process rewards. Using rewards as guiding signals often requires constructing the environment or utilizing proxy feedback, which can improve search effectiveness but may also introduce additional computational overhead. Moreover, when using proxy feedback such as a reward model for search, if the sampled solutions deviate significantly from the distribution of the reward model’s training data, the search may actually lead to a decrease in performance. Besides, environmental feedback that is indirectly related to the final performance can also be used to guide the search process, such as code compilation outcomes, correctness of intermediate steps in solving mathematical problems, unit tests, and more. This type of feedback often serves as an alternative to feedback on the final task outcomes.
##### Heuristic Rules
Many search algorithms use heuristics as guiding signals alongside environmental feedback. Traditional informed search methods, such as greedy search and A ∗ A^{*} search, require the use of heuristic rules developed based on the specific information of the task to guide the search process ( Russell & Norvig, 2020 ) . Additionally, studies for enhancing LLMs’ reasoning abilities ( Hao et al., 2023 ; Koh et al., 2024 ) also apply task-specific heuristics to guide search. When we want to reduce the cost of environmental feedback or when environmental feedback is unavailable, using heuristic rules is a good alternative.
#### 5.2.3 Comparison of Internal and External Guidance
This section compares internal and external guidance. Internal guidance relies solely on the model, avoiding the need for external environments or ground truth, and is generally task-agnostic. Thus, internal guidance is highly transferable and useful when specific evaluation criteria for downstream tasks are unavailable . External guidance relies on specific downstream task information, such as rewards from interactive environments or ground truth, making it more aligned with model performance and better at guiding search strategies. However, it introduces construction costs and computational overhead. During inference, the ground truth is often unavailable, and interacting with the environment or simulator is costly. Additionally, external guidance from surrogate models (like a fixed reward model) can face out-of-distribution (OOD) issues ( 4.4 ). Therefore, external guidance during inference requires careful consideration.
#### 5.2.4 Combination of Internal and External Guidance
Internal and external guidance can be combined to direct the search process, with typical approaches integrating the model’s own uncertainty and proxy feedback from reward models. For example, Wang et al. (2024c) and Snell et al. (2024) combine self-consistency with feedback from a process reward model to select the highest-quality responses.
The value function is another type of signal that combines both internal and external guidance. In reinforcement learning, the value function estimates expected cumulative rewards from a state (V function) or state-action pair (Q function), guiding agents to select actions that maximize long-term rewards. It typically relies on environmental reward signals and uses a separate neural network ( Yu et al., 2024a ; Liu et al., 2024c ; Wang et al., 2024c ) . The calculation formula for the value function is as follows (refer to Sutton & Barto (1998) ): v π ( s ) ≐ 𝔼 π [ G t ∣ S t = s ] = 𝔼 π [ ∑ k = 0 ∞ γ k R t + k + 1 ∣ S t = s ] , for all s ∈ 𝒮 , v_{\pi}(s)\doteq\mathbb{E}_{\pi}[G_{t}\mid S_{t}=s]=\mathbb{E}_{\pi}\left[\sum_{k=0}^{\infty}\gamma^{k}R_{t+k+1}\mid S_{t}=s\right],\,\text{for all }s\in\mathcal{S}, (2)
In this equation, v π ( s ) v_{\pi}(s) denotes the value of state s s under policy π \pi , while 𝔼 π \mathbb{E}_{\pi} represents the expectation taken over all possible trajectories following the policy π \pi . The term G t G_{t} is the return, which is defined as the cumulative sum of future rewards discounted by the factor γ \gamma , where γ ∈ [ 0 , 1 ] \gamma\in[0,1] determines the importance of future rewards relative to immediate rewards. Specifically, R t + k + 1 R_{t+k+1} is the reward received at time step t + k + 1 t+k+1 , conditioned on the state at time t t being s s . Finally, 𝒮 \mathcal{S} refers to the set of all possible states in the environment.
The value function plays a fundamental role in reinforcement learning, as it quantifies the long-term expected reward for starting in a specific state and taking actions according to the given policy. By evaluating the value of states, the value function enables agents to compare and make informed decisions about which states or actions are more favorable under a particular policy.
Rafailov et al. (2024) show that DPO is an inverse Q-learning algorithm, where DPO model logits act as the Q function, enabling search processes guided by the value function. They also demonstrate performance gains from applying beam search to the DPO model. Additionally, some work ( Silver et al., 2017 ; Wan et al., 2024 ; Chen et al., 2024a ) incorporates a value head into the policy model, sharing a large-scale backbone. The main challenge in using value functions lies in accurately estimating them , especially in tasks with sparse rewards or high-dimensional outputs, such as large language model generation, where inaccuracies can significantly impact performance.
### 5.3 Search Strategies
We categorize search strategies into two types: tree search and sequential revisions . Tree search generates multiple answers simultaneously, acting as a global search that explores a broader range of solutions. In contrast, sequential revisions refine each attempt based on previous ones, functioning as a local search that may offer higher efficiency ( Snell et al., 2024 ) .
#### 5.3.1 Tree Search