Reverse-o1: Reverse Engineering of OpenAI o1(Part 2)

Reverse-o1: Key Elements of RL and How to Integrate RL with LLM

We start by deducing how o1 might integrate RL with LLM and call the deduced model Reverse-o1.

We will first analyze the key elements of RL in the Hidden COT scenario: State Space, Action Space, and Reward Model. As for the RL method, I speculate that it is likely to be similar to AlphaGo/AlphaZero, for several reasons:

Firstly, it is said that OpenAI employees read Sutton’s “Bitter Lesson” several times a day, which mentions that “universal methods that can leverage computing power, such as search and learning, will ultimately succeed.” The search mainly refers to the MCTS method of DeepMind’s AlphaGo. It would be strange if OpenAI employees, influenced by this, did not put search into practice.

Secondly, in the interview video of the main members of o1 recently announced by OpenAI, an employee mentioned that they have always been trying to integrate AlphaGo’s search method with LLM, which is also evidence.

So, we will briefly introduce the working principle of AlphaZero and try to integrate it with LLM to construct a complex logical reasoning system.

State Space in O1’s RL: Continuous State Space Composed of Token Sequences

Regarding the RL state space of o1, the first question is: Is this state space discrete or continuous? It is likely to be a continuous state space, or it is best to regard it as a continuous state space. O1 is a combination of LLM and RL. When a user inputs a question, naturally, these Token sequences that make up the question can be regarded as the first state (State1). The Token sequence of State1 serves as the input for the o1 model, o1 selects an action in the action space (the definition of the action space will be discussed later), regardless of what this action is. After selecting this action, o1 will output a Token sequence fragment (not a complete Hidden COT, it should be a fragment of it). Then, o1 appends the newly generated Hidden COT fragment content to State1 to form State2, which is used as the new input for o1 again. o1 selects a new action based on the new input and outputs a new Token sequence fragment… and so on, until the Hidden COT output is completed. It’s basically this process.

The RL state space of o1 is unlikely to be composed of discrete states. It is difficult to clearly divide a number of specific states. Of course, it can be said that in the extreme case, each Token forms a discrete state in the state space. However, this is basically not feasible in practice. Of course, it can be said that in the extreme case, each Token forms a discrete state in the state space. However, this is basically not feasible in practice. If each Token represents a state S, first, the state combination space is too large. Assuming the Token dictionary size is 100,000, the combination space of 2 Tokens is 100,000 squared, and the state space for a Token sequence of length n is 100,000 to the power of n, which is an astronomical number. Secondly, for o1’s RL, for each Token input, a certain action A needs to be selected, and the next Token representing the transition to another state S’ needs to be generated. If the RL process has search, it means that a search needs to be done for each Token, and we can see from many examples of o1 on the Internet that the Hidden COT is often very long, possibly tens or hundreds of K, and this calculation is basically unacceptable. So, looking at each Token as a discrete state is not impossible, but the granularity is too fine, and it feels difficult to apply in practice.

I think it is more appropriate to regard o1’s state space as a continuous state space composed of Token sequences. Although the example mentioned above refers to State1 or State2, which seems to be discrete states, this is just for the convenience of explaining the process (of course, if State1 is regarded as a point sampled in the huge Token combination space, there is no problem). It is similar to RL playing games or RL playing Go, where the RL input game (or Go) screen is composed of, for example, 1024*1024 different pixels (different pixels can be likened to different Tokens of LLM), and due to the huge combination space of pixels, it is difficult to clearly define what a discrete state is. Therefore, general RL playing games or Go treats the input image as a whole and regards it as a continuous state space, mapping it to a specific action through a neural network. The state space of O1 is similar to the image (refer to the figure above), and a Token fragment can be likened to a certain image input corresponding to RL playing games, regarded as a continuous state space composed of Token sequences, and mapped to a certain action in the action space through the LLM+RL neural network of o1.

From the above analysis, it can be seen that the RL technology used for playing games or Go is mostly based on continuous state space as the network input, and the output is mostly a certain action in the discrete action space. So it is obvious that these RL technologies are more suitable for being used as the solution for the RL part of o1, and RL models with discrete state spaces, such as MDP methods, are not very suitable.

Possible Action Space in O1’s RL: “Thought-Factor” Discrete Action Space

The key link in the RL technical solution of O1 is how to define the action (Action) space. The process of generating Hidden COT by OpenAI 01 is essentially imitating the thinking process of humans in solving complex problems, and when humans think about complex problems, there are relatively fixed and not too many “thinking modes” or “thinking factors”. For example, when we get a complex problem, we usually first clarify what the goal of the problem is, then decompose the complex problem into several links or steps, in order to get the solution of a specific step, we may propose a hypothesis, and then verify whether the hypothesis is established, if not, then continue to propose a new hypothesis until the sub-problem is solved… We may also calculate in the process and find that there are errors in some intermediate links, and correct the errors.

If you carefully analyze several Hidden COTs released by OpenAI’s official website, you can summarize some typical human thinking factors implied in thinking about complex problems (refer to the figure above, I have given some specific examples). I think if you regard Hidden COT as a Token, it is difficult to do RL (the state space of Hidden COT is already continuous and non-discrete, and if the action space is also non-discrete or the combination space is too large, RL is difficult to model. So the action space is discrete, this is very likely to be true, of course, how to define the discrete action space should have different methods), in my imagination, a reasonable method is to summarize the human thinking factors implied in thinking about complex problems, and use them as the candidate action set, such as “decompose the problem”, “rephrase the goal”, “check the result”, “correct the error”, “propose a hypothesis”, etc., the total number should not be too much, even if it is divided more finely, it is estimated to be a few dozen to a hundred. For each specific “thinking factor”, a Token fragment that conforms to the corresponding distribution probability can be generated, for example, if the behavior is the “propose a hypothesis” factor, then the probability of generating the Token “Alternatively” is relatively large (learned from the training data by PPO). So, the original content of Hidden COT may look like this:

<ACT_Proposer-Start> Alternatively, perhaps combine the numbers in some way. <ACT_Proposer-End> (Proposing a hypothesis)

<ACT_RephraseTarget-Start> Overall Task: Write a bash script that takes one argument (the string representing the matrix) and outputs its transpose in the same format. <ACT_RephraseTarget-End> (Rephrasing the goal)

That is to say, the original content or training data of OpenAI’s Hidden COT may be a two-level structure in form:

<Think_Start> (Hidden COT start mark)

……

<ACT-1_Start>token token token…..<ACT-1_End> (Thinking factor 1)

<ACT-2_Start>token token token…..<ACT-2_End> (Thinking factor 2)

<ACT-3_Start>token token token…..<ACT-3_End> (Thinking factor 3)

……

<ACT-n_Start>token token token…..<ACT-n_End> (Thinking factor n)

<Think_End> (Hidden COT end mark)

This hierarchical hidden COT structure can reflect the combination of advantages of RL and LLM. The action space of discrete actions, such as estimating what action to take given a state S, that is, the estimation of function Q(S,A), is something RL is good at, and the generation of Tokens in the thinking factor label is something LLM is good at. LLM can learn to adjust the generation probability of Tokens inside the factor label according to the type of “thinking factor”. The figure above shows the possible operation form of o1 after the two-level “thinking factor” discrete action space as described above. In the process of generating Hidden COT, the input and output both have the start and end symbols of ACT action Tokens.

First, O1 predicts the most likely “thinking factor” to be taken next based on the current question and the already generated Hidden COT fragment, to decide what specific thinking mode to take later, and then under the guidance of this “thinking factor”, LLM generates a specific Token sequence, using the end Token of “thinking factor” as the end mark of this thinking mode. And the Token sequence output this step is merged into the input to cycle through to generate the corresponding action and Token sequence for the next thinking… (Of course, the whole process is my imagination, there is no specific evidence).

You may ask: Why can’t I see the start and end Tokens corresponding to the “thinking factor” in the Hidden COT examples I give? The COT shown to users is likely to be a filtered version. You think, the start and end Tokens of Hidden COT (<Think_Start>/<Think_End>), these two Tokens are very likely to have, you haven’t seen it, right? It means that the output is a filtered COT, so there were “thinking factor” marks before, but they were filtered out when displayed, which is also possible.

Reward Model of O1’s RL Model

How to set Reward is crucial for RL. There are actually quite a few academic works on LLM+RL, and there are two common types of Reward models (refer to the figure above): Output Reward Model (ORM) and Process Reward Model (PRM).

The meaning of ORM is to train a model that only scores the final result, regardless of how many steps the derivation process takes. If you look at it in terms of Hidden COT, it means that only when o1 has completely written the Hidden COT, ORM will give a reward signal, and if the model result matches the standard answer, give a reward of 1, if the answer is wrong, give a reward of -1, something like that. It is clear that the advantage of ORM is that the feedback signal is accurate, for example, for math problems, the model is either right or wrong, very clear, so the feedback signal is precise; but the disadvantage of ORM is that the feedback signal is sparse, which is very intuitive, even if you write 10 pages of derivation process, anyway, there is only one feedback signal in the end. (The RM model of RLHF stage when OpenAI trains a large model belongs to ORM)

The meaning of PRM is to train a model that can give feedback signals for each intermediate process, so it is clear which step is wrong in the derivation process, and you don’t have to wait until the end, so its characteristic is that the feedback signal is rich and not sparse. But the problem comes, to train PRM, you need to have annotated data for each step, how to get so many annotated signals? The conventional approach is to rely on manual annotation, such as last year’s very popular OpenAI’s PRM work “Let’s Verify Step by Step”, which relied on manual annotation of 800,000 intermediate step feedback signals of math problem derivation process, and proved that PRM is better than ORM. So, the advantage of PRM is that the feedback is rich and the effect is good, but the cost of training data production is too high, and ordinary people can’t do it.

So is there a relatively low-cost method to annotate each step? Yes. The better approach I have seen is to do it this way (refer to the figure above): assuming we have a batch of math problems with complete derivation processes, we can first copy the first problem-solving step, and then use the MCTS tree to continue the derivation by search, we can do multiple derivations from this step, some of which will get the correct answer, and some will result in errors. In principle, the higher the proportion of derivations from this step that lead to the correct answer, the more important this copied step is for getting the correct answer, so it can be given a high score, and then the second problem-solving step can be copied, and so on… In this way, we can automatically annotate the quality of each derivation step. Then we can use this data to train the PRM model, and PRM can score each reasoning step. But it is clear that the precision of the PRM score trained by this data is definitely not as good as ORM.

So will OpenAI o1 use ORM or PRM during the training process? I guess both will be used. ORM is accurate, and PRM has rich feedback, each has its advantages, and the combination should be better. In addition, o1’s official website mentioned “Our large-scale reinforcement learning algorithm teaches the model how to think productively using its chain of thought in a highly data-efficient training process.”, the “data-efficient” here should refer to PRM.

The Basic Principle of AlphaZero

Here we will first introduce the basic working principle of AlphaZero. The Reverse-o1 plan we give later, the core is how to integrate RL and LLM, and the main framework mainly refers to the main idea of AlphaZero, so we will explain it here to facilitate the understanding of the subsequent content.

At the end of 2017, the general version of AlphaGo’s chess game, Alpha Zero, came out, not only Go, but also for other chess games such as international chess and Japanese chess, AlphaZero also defeated the strongest AI programs including AlphaGo with an overwhelming advantage.

AlphaZero has not made essential improvements in terms of technical means compared to AlphaGo, and the main body is still the MCTS Monte Carlo search tree plus neural network structure and RL training method, but the technical implementation is much simpler and elegant (refer to the figure above). The main changes include two aspects: one is to merge AlphaGo’s two prediction networks (policy network P and value network V, the policy network P is mainly used to predict the winning probability of each action a, that is, the possible chess piece position, under the current state S, that is, the function P(S,a); and the value network V mainly evaluates the overall probability of winning the current state S, that is, the function V(S), which is a value between 0 and 1, the larger the value of V(S), the higher the probability of winning from the current situation S.) into one network, and output two types of results P(S,a) and V (S) at the same time; the second is to upgrade the network structure from CNN to ResNet. AlphaZero completely abandoned learning from human chess games and started learning from a blank sheet of paper through self-play, and obtained far more than a thousand years of human accumulated Go experience in just three days of Self Play.

AlphaZero combines MCTS and RL, with MCTS as the main body, and RL plays a role in accelerating the search speed. During the Self Play process (refer to the figure a), for a certain AI chess player, it will use MCTS search, search for each possible chess piece (Action) under the current state S, and after searching each position, it can obtain the winning probability distribution pai of each chess piece position, and choose the position with the largest probability to place the chess piece, and then another AI chess player will also adopt a similar idea to place the chess piece… So back and forth until the winner is determined z indicates who is the winner, Reward signal).

So far, it seems that we have not seen the role of the neural network structure, in fact, it mainly plays a role in the MCTS search of a certain chess piece position. Because starting from a certain chess piece position to search, the searchable space is too large, and brute force search is definitely not feasible, so the role of the policy network P and the value network V (AlphaZero has already merged into one network, it is mainly to facilitate the explanation) is mainly to guide the search process, and prioritize the search of paths with high winning probability, and prune paths with low winning probability, so as to increase search efficiency.

During the search process, the neural network parameters are fixed, and when a game of chess is finished and the winner is determined, the corresponding training data can be generated for each state S passed through in the chess game. For the policy network P, the learning goal is the action probability distribution pai obtained by the MCTS search at the current state S, and for the value network V, the learning goal is the fact that “the final winner z has a high probability of winning”. Then, based on this training data, the neural network parameters can be adjusted, so that it will be stronger in the next game of chess (it can be seen that the reward model of AlphaZero is ORM). In this way, through the infinite repetition of the game process, the ability of AlphaZero becomes stronger and stronger.

It should be realized that for AlphaZero, its essence is still the MCTS Monte Carlo tree search. The reason why Go seems difficult to overcome is mainly that the search space is too large, and brute force search is completely unfeasible. If we assume that there is a machine that is infinitely powerful and can quickly traverse all search spaces, then even if we use MCTS tree search alone without relying on RL, the machine can achieve a perfect game state. AlphaGo Zero

mainly achieved better evaluation of the chessboard state (V) and the quality of the chess piece (P) through self-play and deep reinforcement learning, and gave up a large number of inferior paths, thus greatly reducing the search space needed. The self-evolution mainly reflects that the evaluation of the chessboard state (P and V) is more and more accurate, so it can find the chess piece with the highest winning probability more and more quickly. The reason why a large number of training data can be produced through self-play is that playing chess is a task with clear rules, and it can be won or lost at a certain state, it’s just that this final win or loss comes later, not visible with each chess piece.

Reverse-o1 Model Network Structure

o1 is different from playing chess in one point: in addition to RL, even Hidden COT, it is also output by a Token at a time behind, LLM must still be the main structure, but RL must also have a network structure to adjust the model parameters to gradually learn the internal thinking process. So, the first question we face is: how to integrate LLM and RL two models to obtain a complete network structure that simultaneously has the functions of both LLM and RL.

The figure above shows a structure I imagined: the main body is still a LLM model based on Transformer (Dense or MOE can be, the mini version should be a Dense structure), when the input is “question + already generated part of Hidden COT” (that is, the current state S composed of continuous Token sequences), it is encoded by the GPT network. On the LLM output Head, two sub-structures can be differentiated: one is used for the conventional LLM prediction of the Next Token, which is consistent with the usual LLM; on the Head, an RL model structure can be built, here referring to the idea of AlphaZero, one network with two outputs. For example, an FFN network structure can be used, on the one hand, it outputs the policy network P result (P(S,a)), representing the distribution probability pai of the next Action “thinking factor” under the current state S, the larger the probability of a “thinking factor”, the greater the possibility of this Action being selected for execution in the next step; on the other hand, it will output the value network V result (V(S)), representing the probability of the current state S leading to the final correct answer, the larger the probability, the higher the quality of the current state S, meaning that the overall quality of the part of Hidden COT that has been output so far is higher.

At this point, when Hidden COT is in a certain state S, after the network knows what action to take next, the current state S has a certain probability of leading to the correct answer. But there is still a part missing, that is, the corresponding Hidden COT series of Tokens after the next “thinking factor” behavior is known.

A simple method is to use the LLM part above the LLM head to continuously output the subsequent Tokens (when there is manual data training, PPO can be used to increase the output probability of the corresponding Tokens), and when outputting the subsequent Tokens, RL output is not considered, until LLM outputs to <ACT_i-End>, then judge the RL output to select actions… Continue this process, and the model that combines LLM and RL output Hidden COT can operate.

As we analyzed before, o1 is likely to use the process reward model PRM, and it is likely to be composed of multiple models. Under these two constraints, the model structure above can be transformed in this way (refer to the figure above): after knowing the next “thinking factor”, instead of having the main model generate subsequent Tokens, in order to increase the quality of the subsequent generated COT, the idea of Best-of-N Sampling can be adopted, and multiple copies of the Reverse-o1 model (different copies can set different temperature parameters to increase the diversity of the output) each give a Token sequence, and then the PRM trained offline is used as a judge to score, and choose the Token sequence with the highest score as the output Tokens of this “thinking factor”. After selecting the best content, it can be synchronized to the main model, and the main model can perform an operation similar to Prefill to synchronize the output of the best content, and then start the next round of output… This can be done in this way, and it is obvious that the quality of the generated Token sequence will be higher.

Reverse-o1 under MCTS Tree Search

We imitate AlphaZero and introduce the main structure MCTS, and its operation process is as follows (refer to the figure above): when a user inputs a question, Reverse-o1 uses the MCTS tree to search for each possible “thinking factor”, and uses the policy network P and the value network V to quickly find the optimal search path during the search, so as to obtain the distribution probability pai of all “thinking factors”, the larger the probability value, the higher the probability that this type of thinking leads to the correct answer. Then, the “thinking factor” with the largest probability is selected as the behavior of the current state, and as described in the previous section, the COT Tokens fragment is generated for this behavior by Reverse-o1. The COT Tokens fragment is merged with the user’s question to form a new state… and so on, until the answer to the question is produced, and compared with the standard answer, it is either right or wrong, and the corresponding Output Reward is obtained.

Imitating AlphaZero, starting from state S, when searching for the probability of a certain “thinking factor” leading to the correct answer, the optimal next state S’ is sought in the form of max(Q+U), and the Q function is positively correlated with the value network V(S), and the U function is positively correlated with the policy network P(S,a), so the meaning of Max(Q+U) is to find high-quality search paths guided by the value network and the policy network; when searching to the leaf node, node expansion will be carried out, and the strategy network and value network are used to estimate the initialization of related search parameters, and then the Q function corresponding to all states on the optimal path is updated from bottom to top. After each candidate “thinking factor” has been searched for a round, the distribution probability pai of all actions will be obtained, and the search step is completed. The difference between O1’s search and playing chess is that: if you want to transfer to the next state, you also need to generate the corresponding Hidden COT Tokens based on the current selected behavior, and this step can be completed by the Best-of-N Sampling strategy described in the previous section.

After generating the Hidden COT fragment step by step from the question, when it reaches the answer phase, the Output Reward is obtained, and a MCTS search process from the middle link to the answer is completed. If the answer is correct, Reward can be set to 1, and if the answer is wrong, Reward can be set to -1. Based on this, training data can be constructed for all intermediate states passed to train the policy network P and the value network V, and the learning goal of the policy network is the action probability distribution pai obtained by the corresponding state MCTS search, and the learning goal of the value network is the Output Reward.

In addition, for each “thinking factor” selected during the search process, the corresponding Hidden COT Tokens sequence obtained through Best-of-N Sampling (which can also obtain the corresponding Process Reward score given by PRM) can be used to adjust the LLM model parameters with PPO (the Reward of PRM is used as the Reward of PPO), so that LLM will increase the generation probability of these Tokens when encountering this “thinking factor” later.

So far, the journey of reverse engineering o1 can be basically ended, and some of the constraints mentioned earlier (o1 should be composed of multiple models, some or several tree searches should be used, RPM and ORM should be used, etc.) are basically reflected in the imagined Reverse-o1.

However, I personally think there is still a problem worth thinking about deeply: “Is the thinking factor necessary?” After all, this requires manual induction of human potential thinking modes, and there is still a strong trace of manual work, and it will increase the cost of manual annotation of data. I have been thinking about this problem for several days, and the conclusion seems to be: the whole framework can still be this framework, and it should be possible without introducing “thinking factors” insead of a “Hidden Thinking Factors”, because the article is already too long, and it is a bit complicated to explain, so let’s skip it for now.