Scaling of Search and Learning: A Roadmap to Reproduce o1 from Reinforcement Learning Perspective — part 3/8
Cleaned-text code points [31269, 46861). Bibliography omitted here; source.md includes it.
The outcome reward involves assigning a score based on whether the outputs of a large language model (LLM) meet predefined expectations ( Cobbe et al., 2021 ; Shao et al., 2024 ) . For instance, Cobbe et al. (2021) uses outcome reward to assess whether the solution provided by the LLM to a mathematical problem is correct. Since ground truth is generally available for tasks such as solving mathematical problems, outcome reward is often straightforward.
Although the outcome reward is relatively easy to construct, it lacks supervision for intermediate steps. For instance, a solution to a mathematical problem that leads to the correct answer may still contain errors in the intermediate steps ( Lightman et al., 2024 ) . As a result, using the outcome reward may cause the LLM to generate incorrect solution steps, which could negatively affect performance. Moreover, outcome reward is sparse, as it does not reward intermediate steps, making it challenging to learn step-level policies. Despite these limitations, it remains widely used in reinforcement learning for LLMs due to its simplicity.
In contrast to the outcome reward, the process reward provides a reward signal not only for the final step but also for intermediate steps ( Lightman et al., 2024 ; Yin et al., 2024a ) . For example, Lightman et al. (2024) involves human annotators rewarding the intermediate steps in a mathematical solution generated by an LLM. The outcome reward can be seen as a special case of the process reward, where the rewards for intermediate steps are all set to zero.
Depending on the granularity of the actions, the process can be classified into token-level ( Rafailov et al., 2024 ) and step-level ( Lightman et al., 2024 ) . The step-level segmentation is flexible. Kuhn et al. (2023) uses information entropy to guide step segmentation while methods in mathematical problem-solving typically use newline characters as the split symbol ( Lightman et al., 2024 ) .
Although process rewards show promise, they are more challenging to learn than outcome rewards. For instance, the reward design in Lightman et al. (2024) depends on human annotators, making it costly and difficult to scale. There are some automatic ways to transform outcome rewards to process rewards, which is called reward shaping in traditional reinforcement learning. We introduce the methods of reward shaping in Section 4.2.3 .
### 4.2 Reward Design Methods
Since outcome reward can be seen as a special case of process reward, many reward design methods can be applied to the modeling of both outcome reward and process reward. The resulting models are typically referred to as the Outcome Reward Model(ORM) and the Process Reward Model(PRM). We review these reward design methods for LLMs, categorizing them based on whether or not they have direct access to the reward signal from the environment.
#### 4.2.1 Reward from Environment
The most straightforward approach to designing a reward is to directly utilize the reward signal from the environment, or learn a model to simulate the reward signal from the environment.
##### From Realistic Environment
Many environments can provide effective reward signals, for example, code generation can receive reward signals from a compiler or interpreter. Dou et al. (2024) ; Shojaee et al. (2023) ; Dong et al. (2024) have shown that the quality of code generation is improved with compiler feedback. Zhang et al. (2024e) generates test cases for input questions and uses the testing outcomes as the reward signal. ( Shi et al., 2022 ) measures the similarity of different programs based on their execution results. Zheng et al. (2024) uses the execution results of a program as guidance for refining a code generation model. For some real-world tasks, the environment is just a sandbox. In web shopping ( Liu et al., 2024e ) , the environment is represented by a real HTML page, and the reward is calculated based on the score the LLM obtains through clicks and searches. In more general scenarios, sandboxes are created to allow the policy model to interact directly with, such as in environments like MineDojo ( Fan et al., 2022 ) or Alfworld ( Shridhar et al., 2021 ) .
##### From Simulating Environment
While some environments can provide valid feedback, interacting with them to obtain the reward signal can be costly, or the feedback may not be available at test time. For example, during testing, we might not have test cases to validate whether the program generated by the LLM is correct. In such cases, a reward model is required to simulate the reward signal from the environment. For instance, Lightman et al. (2024) trains a verifier to predict whether the model’s solution to a mathematical problem is correct at test time.
Simulating the reward signal enables the LLM to obtain feedback at any time. However, using such a reward model during learning or search can lead to the problem of distribution shift. As the policy is updated during learning and search, the reward model, which is trained on data from interactions between the old policy and the environment, may fail to adapt to the new policy. This issue, known as reward optimization, has been discussed in Gao et al. (2023) ; Stroebl et al. (2024) . Thus, the reward model must be updated in tandem with the policy improvements. In more general domains, simulating the environment aligns with the concept of a world model, where the transition probabilities of states must be further simulated, allowing for a more accurate reward model. We discuss the world model in Section 4.5.2 .
##### From AI Judgment
AI Judgment involves using general AI assistants to provide the reward signal, serving as an alternative to relying on realistic environments, thereby avoiding the high costs associated with environment construction and interaction. For instance, it is common to use a powerful LLM, such as GPT-4, to assess the performance of an AI assistant ( Zheng et al., 2023b ) . While AI judgment can be viewed as a form of reward model, it does not face the issue of reward optimization, as it is not dependent on the policy model. As a result, AI judgment remains effective even when the policy model undergoes updates. The LLM is an implementation of the world model, which also highlights the effectiveness of building a world model to provide reward signals.
#### 4.2.2 Reward Modeling from Data
For some environment, the reward signal from environment is unavailable, which can not be simulated. For example, it is difficult to judge whether the response from an AI assistant is good or not. Fortunately, it is easier to collect the expert data or the preference data than giving a reward. With the expert data or preference data, we can also learn a model to provide effective rewards.
Learning rewards from preference data is widely recognized in the LLM community, particularly due to the success of RLHF ( Bai et al., 2022a ) . However, learning rewards from expert data, also known as Inverse Reinforcement Learning, has not been widely used for LLMs, which may represent a promising technique for o1.
##### Learning Rewards from Preference Data
Preference data is collected by ranking multiple responses from LLMs to the same question. Using the Bradley-Terry model, an outcome reward can be derived based on the pair-wise comparison ( Bradley & Terry, 1952 ) . Learning rewards from preference data for LLM is first proposed in Christiano et al. (2017) , which introduced the preference signal to tackle complex reinforcement tasks without access to a reward function, achieving significant success in Atari games and robotics. They further developed this approach for text summarization ( Stiennon et al., 2020 ) . This technique is then used for the development of ChatGPT, where preference signals are used for aligning policy behavior with human values ( Bai et al., 2022b ) .
Learning from preference data has been widely used in alignment of LLMs. However, it is crucial to construct the preference data that accurately reflects the actual performance of downstream tasks. For example, Wen et al. (2024) show that using human preferences as supervision may degrade the true performance of the model.
##### Learning Rewards from Expert Data
Inverse Reinforcement Learning (IRL) is a method for learning rewards from expert data, with the goal of recovering the reward function that the expert is optimizing. This is achieved by fitting the reward function to the trajectories generated by the expert and maximizing the recovered reward ( Abbeel & Ng, 2004 ; Ng & Russell, 2000 ) . Many IRL approaches integrate adversarial learning techniques. Garg et al. (2021) introduces the convex conjugate function to avoid the need for adversarial training. Compared with learning reward from preferences, the data for IRL is easier to be collected. However, IRL typically involve adversarial training, which makes the learning more complicated than learning rewards from preferences. Despite that IRL is well-known in reinforcement learning, there is no empirical evidence that it has been used for large-scale reinforcement learning for LLM.
#### 4.2.3 Reward Shaping
The reward signal from some environments may not be effective; for example, it could be an outcome reward rather than a process reward. In such cases, the reward can be reshaped to make it denser and more informative, a process known as reward shaping. For instance, Kumar et al. (2024) attempts to train an LLM to self-correct via reinforcement learning and discovers that proper reward shaping can prevent learning collapse. Wang et al. (2024c) reshaped the outcome reward by estimating the value function using Monte Carlo sampling, employing the Q-value Q π ( s t , a t ) Q_{\pi}(s_{t},a_{t}) as the reward during intermediate steps. However, since the value function depends on the policy π \pi , the value function estimate from one policy may not be a valid reward function for another policy. This is empirically validated by Xiong et al. (2024) , who found that using a value function from another policy for reward shaping is harmful.
Ng et al. (1999) defines potential-based reward shaping and proves that two reward functions can lead to the same optimal policy if and only if they satisfy the following equation: F ( s t , a t ) = r ( s t , a t ) + γ ϕ ( s t + 1 ) − ϕ ( s t ) , F(s_{t},a_{t})=r(s_{t},a_{t})+\gamma\phi(s_{t+1})-\phi(s_{t}), (1) where F ( s t , a t ) F(s_{t},a_{t}) and r ( s t , a t ) r(s_{t},a_{t}) are two reward functions that yield the same optimal policy, γ \gamma is the discount factor, and ϕ \phi is a potential function. Potential-based reward shaping suggests that the reward function r ( s t , a t ) r(s_{t},a_{t}) can be reshaped without altering the optimal policy, as long as the shaping satisfies Equation 1 . Setlur et al. (2024) uses the theory of potential-based reward shaping to transform outcome rewards into process rewards.
The reward learned from preference data can also be reshaped. For example, Rafailov et al. (2023) , Rafailov et al. (2024) and Zhong et al. (2024) demonstrate that DPO implicitly makes potential-based reward shaping ( Ng et al., 1999 ) on the reward learned from preference data.
While reward shaping can be advantageous, it may also be harmful. Improper reward shaping can negatively affect both learning and search processes. As a result, reward shaping requires careful design and often entails incorporating inductive bias.
### 4.3 Speculation About the Reward Design of o1
We now turn to discuss our assumptions regarding the reward design for o1. Given that o1 could handle multi-task reasoning, its reward model is likely to incorporate multiple reward design methods. For complex reasoning tasks, such as mathematics and code, where responses typically involve long chains of reasoning, a process-reward model is more likely to be employed to supervise the intermediate process, rather than ORM. Techniques like reward shaping can help derive process rewards from outcome rewards.
When no reward signal is available from the environment, we suspect that o1 will need to rely on learning from preference data ( Christiano et al., 2017 ) or expert data ( Garg et al., 2021 ) .
Given that o1 can be fine-tuned using few-shot examples, we suspect that it has a robust reward model trained on a large and diverse dataset spanning a wide range of domains. It can be adapted to a new domain easily through ground truth and solution pairs. Moreover, it is more likely to predict rewards by generating with LLM rather than through value heads.
### 4.4 Challenges of Reward Design for Reproducing o1
##### How to overcome distribution shift?
The reward model learns from an existing dataset distribution and the out-of-distribution problem still remains significant today ( Gao et al., 2023 ) , especially as the LLM continues to explore and learn from feedback. The reward provided by the proxy model deviates from the golden reward when the distribution of the policy model changes, as the trajectory becomes unseen during the training of the reward model and heavily depends on its generalization. Scaling the parameters of the reward model and increasing the amount of data can mitigate this issue but does not completely solve it. Iteratively training the reward model provides a more direct solution, but it still requires human involvement in the loop.
##### How to design fine-grained reward for language model?
Unlike in Atari games or robotic environments, language presents a unique challenge because the definition of a step or action can vary in granularity: token-level, step-level or solution-level. In many situations, such as fitting the model to human preferences, it is more natural to judge the entire solution rather than each token individually, as higher-order semantics in language emerge from the combination of tokens. However, using token combinations as actions results in an action space that is too large to define or learn a reward function for. The potential action space grows exponentially, leading to a long tail of actions. Meanwhile, the sparsity of the reward signal increases with the length of each step as we mentioned before.
##### How to choose data when modeling reward for complex tasks?
As the complexity of tasks increases, choosing the appropriate type of feedback becomes increasingly challenging. Recent studies have shown that using preference-based feedback for tasks like code generation or mathematical reasoning may actually degrade the performance of the policy model ( Wen et al., 2024 ) . Additionally, the question of how much data is necessary to accurately capture the intended behavior remains underexplored. The difficulty of evaluating whether a reward is effective also grows as the complexity of the task increases.
### 4.5 Generalization
The previous section primarily focused on reward design for specific tasks, such as math and coding. When addressing more general tasks, we need to create a broader environment. Practically, according to OpenAI’s 5-stages plan for AGI, o1 has been a strong reasoner, the next stage is to train an agent that can interact with the world and address real-world problems. To achieve this goal, we need a reward model to provide the reward signal for the agent taking actions in a realistic environment. In this section, we focus on how to construct a general reward signal, which can be divided into two components: the reward ensemble and the world model.
#### 4.5.1 Reward Ensemble