RLHF & PPO: From text continuation to aligned assistant

Grab some popcorn 🍿! Today we’ll check how OpenAI turned the monster one gets from training on raw internet data into something useful. We’ll walk through three key stages:

1. Text-continuation model (GPT-3, June 2020)
2. Instruction-following aligned model (InstructGPT, Jan 2022)
3. Chatbot (ChatGPT, November 2022).

What do I understand by “text-continuation” model?

I am referring to the models we get by “just” optimizing for next-token prediction on a huge corpus of generic text. The loss function used is categorical cross-entropy on the correct token.

Unless applying heavy prompt engineering or finetunning (like we will see in this post), these models aren’t good at question-answering and aren’t necessarily aligned with their users. As a result of being trained on huge amounts of unsupervised data from the internet, they often return untruthful, toxic, or biased outputs.

In this post, we’ll see how these problems got addressed.

Visual representation of LLM finetuning.

2017: Reinforcement Learning from Human Preferences

Before directly jumping into LLMs let’s first go back and review the paper RL from human preferences from 2017¹ (context: Tip 1).

¹ This might look a bit disconnected from the main topic but bare with me.

Tip 1: State of RL in 2017

For me this was the golden age of Reinforcement Learning (RL) Tip 2:

• New algorithms were being published left and right to solve OpenAI’s Gym environment tasks.

• AlphaGo² from Deepmind had just defeated Lee Sedol at GO: A game with a search space so big that it was considered almost impossible for a machine to achieve human-level performance.

• DeepMind and OpenAI where battling to master StarCraft and Dota respectively with RL-based algorithms.

• OpenAI was about to publish some research on impressive dexterity learning skills, and agents learning to play hide and seek.

² I can’t recommend enough the linked documentary, super emotional portrait of human vs machine capabilities.

³ I’m actually not sure how conscious they were about it on that moment…

⁴ For instance, think of gymnastics: for most of us it is very hard to demonstrate what a good move looks like 🤸‍♀️. However, it is very easy to tell apart a pathetic one from a decent one.

RL from human preferences laid the foundations of what would become key to later align language models³. In particular, they focus on hard-to-specify but easy-to-judge tasks: Oftentimes we can recognize a desired behavior but not necessarily demonstrate it⁴.

Tip 2: What is Reinforcement Learning?

Reinforcement Learning (RL) is a branch of Machine Learning concerned on how artificial agents ought to take actions in an environment in order to maximize some notion of reward. The objective for the agent is to learn an optimal policy (how to act in the given environment so as reward is maximized).

The typical framing of a Reinforcement Learning (RL) scenario: an agent takes actions in an environment, which is interpreted into a reward and a representation of the state, which are fed back into the agent.

Arguably, this is what living beings also evolved to do: Act in a way that maximizes personal well-being and species continuation.

During my masters I became a bit obsessed with RL (and the reward hypothesis) and I have a lot of posts, paper summaries, and experiments around it in my old blog: CampusAI/RL.

How does it work?

Consider a basic RL problem where we don’t have a specified reward function for it, for whatever reason⁵.

⁵ Such as it being too hard, too sparse or too open-ended.

⁶ Good action $\to$ high reward. Bad action $\to$ low reward.

⁷ Remember that in the context of Rl a trajectory $\tau$ is a sequence of state-action pairs: $\tau =\{(s_0,a_0),...,(s_T,a_T)\}$

One solution would be to have a human manually giving feedback (aka rewards) to the agent⁶. While this would work, it isn’t sample efficient: it takes too many human inputs, plus it might be hard to consistently give an absolute value to a particular trajectory⁷.

Their proposal is to train a new model (the reward predictor) which approximates the reward function⁸ in a supervised way. Its inputs are pairs of agent trajectories, and the learned outputs are human “preferability score”. The agent then updates its policy according to the reward given by this reward predictor model. This generalizes with few training examples and the human doesn’t need to label as much data 🥳.

⁸ Mapping from trajectories to desirability.

This is the overall idea:

How does the agent integrate human preferences?

Tip 3: From paired comparisons to an absolute scalar

This is all very nice conceptually, but a question remains, how does one map “I prefer this one over that one” to an actual scalar that can be used as a reward?

Simple, we assume that this scalar exists and we approximate it: You can think of it as some latent factor $r$ explaining the human preference associated with each trajectory. Imagine we get a dataset of paired trajectories and the human preference:

$$\mathcal{D}=\{({\sigma }_i^1,{\sigma }_i^2,{\mu }_i){\}}_{i=1:n}$$

Where ${\sigma }^1$ is the first trajectory, ${\sigma }^2$ is the second one, and $\mu \in [0,1]$ indicating the preference (human label). For instance if the human marked ${\sigma }^1\succ {\sigma }^2$ (1 is better than 2), then $\mu =0$. If the human says both are the same, then $\mu =\frac{1}{2}$. The human can also mark the segments as incomparable, the sample is removed from the dataset.

We can then write the probability of preference as a softmax of the latent factor (expected reward):

$$\hat{P}[{\sigma }^1\succ {\sigma }^2]=\frac{e^{\hat{r}({\sigma }^1)}}{e^{\hat{r}({\sigma }^1)}+e^{\hat{r}({\sigma }^2)}}$$

To learn the reward (latent factor) we can do cross-entropy loss:

$$\text{loss}(\hat{r})=-\sum _{({\sigma }^1,{\sigma }^2,\mu )\in D}\mu (1)\log \hat{P}[{\sigma }^1\succ {\sigma }^2]+\mu (2)\log \hat{P}[{\sigma }^2\succ {\sigma }^1].$$

SUMMARY: We are doing standard binary classification and using the logits as the reward value.

Tip 4: Other neat details about this paper

• The reward predictor and the main agent’s policy are trained simultaneously. Thus, to train the agent they focus on RL methods which are robust to changes in the reward function: such as policy gradient methods. These methods directly learn the best policy, so a change in reward isn’t as critical as in the methods that attempt to learn the “value” of states or actions, such as value-based methods, or model-based methods.

• Consider a simulated robot whose goal is to learn how to do a backflip. Interestingly, it was able to learn this trick with 1 hour of non-expert human feedback. On the contrary, it took 2 hours of an expert human to write a reward function which achieved the same goal.

2021: From text-continuation to instruction-following

Alright, fast forward to the Large Language Model (LLM) era (aka 2021). We now have very powerful LLMs which do next-token prediction very well (s.a. GPT-3). These are great to be finetuned to different tasks, one of which being question-answering.

How do we make sure their answers are factually correct, unbiased and aligned with our intentions? Actually this falls into the realm of things that are hard to define but (relatively) easy to evaluate. Wait!? Didn’t we just talk about a possible solution to these type of problems? 😱🤯 What were the chances??⁹

⁹ Actually quite high, given I planned this minimally.

Anyway, this is how they do it:

• Step 0: Create a dataset of instructions and desired answers (humans do this)¹⁰.
• Step 1: Good old supervised finetunning (SFT) of the Language Model (LM) using the previous dataset.
• Step 2:
  • Sample instructions from the dataset and obtain different answers from the model¹¹.
  • Have a human rank those answers¹².
  • Train the reward model: An adaptation of the LM that learns to rank answers of a given question.
• Step 3:
  • Sample questions from the dataset and generate new answers using the finetuned model.
  • Forward these question-answer pairs through the reward model, obtaining a “score” (aka reward) for each one.
  • Adjust the LM weights so that high-score answers are more likely than low-score ones using an RL algorithm. The algorithm they used is Proximal Policy Optimization (PPO) Tip 5.

¹⁰ Allegedly (and sadly not-very-surprisingly) under bad working conditions.

¹¹ I assume they do this by randomizing a bit the model outputs. For instance, by playing with the slack given by temperature or other sampling techniques (i.e. not greedily choosing the most likely next token).

¹² Same as if they were agent trajectories in an environment as we saw previously.

Recipe to get an aligned instruction-finetuned LM.

Tip 5: PPO, the general case

Alright, let’s put the first gear… This is the steepest part of all this. It might seem like going a bit too deep for this post but I still wanted to refresh my memory on those algorithms (they are very cool).

The hidden truth behind RL.

Basics

Consider a standard RL problem defined by a:

• State space: $\mathcal{S}=\{s_i{\}}_i$. $$ Set of all possible states the agent can be in.

• An action space: $\mathcal{A}=\{a_j{\}}_j$. $$ Set of all possible actions the agent can take in the given environment.

• Reward function: $R(s_t,a_t)$. $$ Reward given for taking a given action $a$ at time $t$ at state $s$.

Policy Gradient Methods

Consider an explicit, differentiable, parametrized (by $\theta$ ¹³) policy, directly mapping a given state to a probability distribution over the action space:

$${\pi }_{\theta }:\mathcal{S}\times \mathcal{A}\to \mathbb{R}$$

Usually we write:

$${\pi }_{\theta }(s,a)\equiv {\pi }_{\theta }(a\mid s)=p(a\mid s)$$

Policy gradient methods find the optimal policy parameters ${\theta }^{\star }$ directly differentiating over the following objective function (the reward expectation):

$$J(\theta )=E_{\tau \sim {\pi }_{\theta }(\tau )}[R(\tau )]$$

Where $\tau$ is a trajectory (succession of state-action pairs $\{(s_0,a_1),...,(s_n,a_n)\}$¹⁴) in the given environment sampled by running ${\pi }_{\theta }$.

Exploiting the log gradient trick we obtain that: $${\mathrm{\nabla }}_{\theta }J(\theta )=E_{\tau \sim {\pi }_{\theta }(\tau )}[{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(\tau )R(\tau )]$$

Then, by plugging in ${\pi }_{\theta }(\tau )=p(s_1)\prod _{t=1}^T{\pi }_{\theta }(a_t|s_t)p(s_{t+1}|s_t,a_t)$ into the $\log$ and observing that the gradient is zero for all the terms not depending on $\theta$, we obtain:

$${\mathrm{\nabla }}_{\theta }J(\theta )=E_{\tau \sim {\pi }_{\theta }(\tau )}[(\sum _{t=1}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t))R(\tau )]$$

This is how the pseudocode of this looks like:

The basic REINFORCE algorithm implements this idea.

Actor-Critic Methods

Purely policy-gradient-based methods such as the REINFORCE suffer from high variance in gradient estimates in practice. This makes learning unstable and slow.

Oftentimes, they get combined with actor-critic methods. You won’t believe it, but those methods actually have two components:

• The actor: Policy function that selects the actions based on current state
• The critic: Function evaluating how good the state/action being taken is. Usually $V(s)$ if evaluating state-value, $Q(s,a)$ if evaluating action-value.

In the case of TRPO and PPO, they make use of the advantage function as a critic (A2C paper):

$$A^{\pi }(s,a)=Q^{\pi }(s,a)-V^{\pi }(s)$$

It models how advantageous each action is in a given state.

Turns out we can measure how advantageous some parameters ${\theta }^{\prime }$ are wrt some others $\theta$ by doing something like that:

$$J(\theta )-J({\theta }_{\text{OLD}})=E_{\tau \sim {\pi }_{\theta }(\tau )}[A^{{\pi }_{{\theta }_{\text{OLD}}}}(\tau )]$$

Iteratively finding better parameters maximizing this difference is known as Policy Improvement. It can be shown that this improvement objective is actually the same one as in Policy Iteration, thus guarantying the global optima. Thus, we want to optimize:

$$\underset{\theta }{max}{E}_{\tau \sim {\pi }_{\theta }(\tau )}[A^{{\pi }_{{\theta }_{\text{OLD}}}}(\tau )]$$

TRPO (Trust Region Policy Optimization)

TRPO (2015) was one of the first very successful policy gradient methods, achieving much better sampling efficiency than its value-based counterparts. Its main contribution is to turn the algorithm from an on-policy method into an off-policy method¹⁵. They do so by taking advantage of importance sampling:

$$E_{x\sim p(x)}[f(x)]=E_{x\sim q(x)}[\frac{p(x)}{q(x)}f(x)]$$

The objective function then becomes:

$$\underset{\theta }{max}{E}_{\tau }[\frac{{\pi }_{\theta }(\tau )}{{\pi }_{{\theta }_{\text{OLD}}}(\tau )}{A}^{{\pi }_{{\theta }_{\text{OLD}}}}(\tau )]$$

Subject to:

$$E_t\left[KL({\pi }_{{\theta }_{\text{OLD}}}(\cdot \mid s_t),{\pi }_{\theta }(\cdot \mid s_t))\right]<\delta$$

Where $KL$ is the Kullback–Leibler divergence between the two distributions. We are trying to improve the old parameters in the direction that maximizes the advantage (reward difference). However, we don’t allow the new policy to deviate too much from the old one to ensure stability. This is controlled by $\delta \in \mathbb{R}$.

This algorithm is pretty good, but it is hard to implement the constraint optimization.

PPO (Proximal Policy Optimization)

PPO (2017) addresses the previous issue by introducing a slight change in the objective function. Let’s first define:

$$r_{\tau }(\theta ):=\frac{{\pi }_{\theta }(\tau )}{{\pi }_{{\theta }_{\text{OLD}}}(\tau )}$$

The objective function of PPO is:

$$L^{\text{CLIP}}(\theta )=E_{\tau }[min\{r_{\tau }(\theta )A_{\tau }^{{\pi }_{{\theta }_{\text{OLD}}}},\text{clip}(r_{\tau }(\theta ),1-ϵ,1+ϵ)A_{\tau }^{{\pi }_{{\theta }_{\text{OLD}}}}\}]$$

Where $ϵ$ limits the incentive of “policy change amount”, typically around $0.2$.

The idea is not to encourage the model to make big changes by limiting the obtained reward (done by the $\text{clip}$ function), while still heavily penalizing worse reward directions (done by the $min$ function)¹⁶. In ablation studies they show that in practice it works the same or better than TRPO, while being much simpler to implement.

PPO pseudocode

¹⁶ If you ae trying to parse this, keep in mind that $A_{\tau }$ could be negative. it is useful to split both cases

¹⁵ Essentially allowing the model to learn from trajectories sampled by a policy different from the currently being trained on.

¹⁴ I’m simplifying some things like considering finite state-spaces, action-spaces, and trajectories.

¹³ E.g. by the weights of an ANN.

Tip 6: PPO for aligning LLMs

Ok, so we learned how the PPO algorithm works for agents interacting in a generic environment Tip 5. Let’s see now how this maps into today’s problem (aligning LLMs):

• State space: Input prompt given to the model.

• An action space: Response given by the model.

• Reward function: $R(s,a)$. This is provided by the reward estimator, given a prompt and a model output, it returns a scalar of “how good” it things the pairing is.

• Policy: Now the parametrized policy is the LLM itself: An “agent” which takes and input state (prompt) and outputs an action (response).

Notice that now there isn’t the temporal aspect from before, we are dealing with a one-step (static) problem. We don’t have a notion of “trajectories” we need to maximize the reward over. With all this in mind, we have that the reward expectation is something like that:

$$J(\theta )=E_{a\sim {\pi }_{\theta }(s)}[R(s,a)]$$

I assume they get ${\pi }_{\theta }(s,a)$ evaluating the likelihood of that answer by the model (similar as one would get perplexity). Computing $r$ is quite simple then:

$$r_{\tau }(\theta ):=\frac{{\pi }_{\theta }(\tau )}{{\pi }_{{\theta }_{\text{OLD}}}(\tau )}$$

Thus, we can just do a gradient step as we would in the PPO case Tip 5. If by now you are thinking something like “Wait… What are we even doing here? Can whatever-this-is be called RL??” Don’t worry, you are not alone. Here you have some twitter drama.

2022: From text-continuation to chatbot

Ok nice, we now know how to align a model to output less terrible answers. How do we make it be able to follow a conversation though?

Actually super easy, barely an inconvenience 😉. We use the same technique as before but now we add conversations to the dataset. To do so, humans mimic conversations (in which they play both sides)¹⁷.Everything else is the same.

¹⁷ They get suggestions from the LLM to more efficiently provide answers.

After ChatGPT a lot of other LLMs used ChatGPT answers to create aligning datasets, reducing thus the amount of work provided by humans.

Epilogue

As I write this, OpenAI announced a new model: o1. They claim “it spends more time thinking through problems before it responds” but don’t give much information on how they do it. There is some speculation which looks interesting. I might add an entry about it in the future if I find reliable insights on how they do it. That’s it for today folks.

Tip 7: References

To create this blog I “consumed” and summarized these amazing resources:

• OpenAI blog: Learning from human preferences
• OpenAI blog: Aligning language models to follow instructions
• OpenAI blog: Introducing ChatGPT
• Deep reinforcement learning from human preferences paper
• PPO video by Edan Meyer.