Large Language Models (LLMs) are trained on enormous amounts of data from the internet during a “pre-training” phase designed to give them general knowledge. They aren’t inherently helpful at this point. They’re essentially “fancy parrots.” Given a sequence of words (a prompt), they are designed to spit out words that, based on their training dataset (the internet), are most likely to appear after your prompt. That means it might generate more questions based on your question in the prompt, mention a tangential fact, etc. This isn’t super helpful for us, but it happens because it’s modeling the distribution of internet text rather than “being a helpful assistant.”
We want to train LLMs to be helpful, friendly, non-harmful, etc. The way we solve this is through Reinforcement Learning from Human Feedback (RLHF). It’s a method for aligning LLMs trained on the internet to be helpful and safe.
Supervised Fine-Tuning
Supervised Fine-Tuning (SFT) is usually the first step in training an LLM to become more helpful. This process involves collecting many quality answers to a diverse range of questions a user might ask. SFT is very similar to the training process in pre-training. The main difference is that the goal is to shape behavior rather than give LLMs general knowledge. Thus, during SFT you train the LLM to predict the response on high-quality instruction -> response pairs rather than the general internet. To learn more about the pre-training process, check out this deep dive where I explain how to train an LLM from scratch!
Proximal Policy Optimization
After SFT, you collect many instruction/response pairs. Each preference example is a triple where is the instruction, is the worse response, and is the “better” response. You train a reward model, , to maximize the Bradley-Terry objective. In other words, the reward model gives a score for each instruction/response (higher means better response). The Bradley-Terry objective then takes an instruction/response pair and the score for each response and computes the probability that is better than : .
Given the reward model, we can now apply proximal policy optimization. If you are unfamiliar with reinforcement learning, review some of my notes in the Appendix that go over it. Essentially, reinforcement learning provides us tools to learn a policy that maximizes the sum of expected rewards (also known as return) in an environment, and proximal policy optimization is a specific, sample-efficient method of learning such a policy.
To understand how proximal policy optimization can help us, we need to convert the problem of training an LLM into an RL problem. In this case, is the LLM and it is parameterized by the weights . The action space is the set of tokens the LLM can output. Once the response is generated, we can plug it into the reward model to get a final reward (note: in this setup the rewards are sparse, you only get a reward at the end and not after each token is produced).
Given this setup, proximal policy optimization updates the policy (LLM) to choose actions (tokens) that maximize the reward (producing a better, more helpful response). But maximizing reward on its own is dangerous. RL is prone to reward hacking, where the policy drifts into strange outputs that score well under the reward model without actually being good responses. To prevent this, RLHF subtracts a KL-divergence term between the new policy and the SFT policy from the objective:
Here controls how much we penalize straying from the SFT policy. A higher keeps the new policy closer to the SFT policy we trust.
Direct Preference Optimization
While PPO is effective at teaching LLMs to be more helpful, it also brings in all of the overhead of RL: agent, trajectories, reward model, value functions, etc. This added complexity is often undesirable which is where Direct Preference Optimization (DPO) comes in. It allows you to take a set of preference pairs (instruction/response pair where response is better than response ), and train directly using those pairs with a cross entropy loss.
The important insight that DPO makes is that maximizing rewards while not diverging too far from the reference (SFT) policy has a closed form solution.
First, remember what PPO was trying to maximize:
This optimization has a known closed-form solution:
where is a normalizing constant that makes sure is a valid probability distribution. Overall, this result shows that you can just re-weight the original SFT policy by the exponentiated reward function to get the optimal policy.
Given this closed form, we can reparameterize to solve for the reward in terms of the policy:
Now we can plug this into our Bradley-Terry objective:
And we can convert the overall goal of maximizing the Bradley-Terry objective into minimizing a supervised loss:
Now we can just treat this as a supervised fine-tuning problem over preference pairs rather than bringing all the complexities of RL in (and training a reward model). So the natural question is why not use DPO for everything? There are a few reasons:
- PPO can explore by generating new responses and seeing how the reward model grades them. DPO is fixed to learning from the preference pairs.
- The reward model is reusable and gives a richer signal. The reward model can score any generated response, and can be reused/iterated on.
- Perhaps one of the biggest reasons is that well-tuned PPO can still outperform DPO empirically.
Appendix
Reinforcement Learning
The main players in Reinforcement Learning (RL) are an agent and the environment. The environment is completely described at timestep by a vector . The agent described by policy (often stochastic) performs some action in the environment (discrete or continuous) which takes you to a new state . The state transition is determined by the current state and action and can be probabilistic . You then collect a reward based on the transition . The trajectory is the entire sequence of states and actions that occurred in the world: . An agent’s objective is to maximize the expected sum of rewards collected over a trajectory .
To make the agent’s goal concrete, let’s observe the following. First, we can represent the probability of a trajectory under policy as:
where is the probability of occurring from the list of possible start states.
The expected cumulative rewards under policy is:
RL tries to find a policy that maximizes :
Value, Action-Value, Advantage Functions
The on-policy value function gives you the expected return at state under policy :
The on-policy action-value function gives you the expected return at state if you take action and then listen to the policy for all future actions:
The advantage function tells you how much better an action is compared to other actions on average:
Standard Policy Iteration
In traditional machine learning, we have algorithms like gradient descent to iteratively improve the weights of a network to converge on a better network (reduce the loss). Policy iteration is a similar concept. We iterate on a policy to converge to a better policy (increase the expected return). This iterative process is done by taking small steps in the gradient of with respect to :
In practice, we collect many sample trajectories , and estimate using the following:
This weights each action by the full return , but we can instead weight it by the advantage without changing the gradient of , which dramatically lowers its variance (see this derivation for why swapping in a baseline and the reward-to-go leaves the gradient unbiased).
Proximal Policy Optimization
While Standard Policy Iteration provides an iterative way to converge on a policy that maximizes expected return, it is very sample inefficient. We must collect a bunch of trajectories to perform just a single policy update. Proximal Policy Optimization is a method that relaxes some of the constraints given by Standard Policy Iteration to improve sample efficiency.
For the sake of completeness, there are two forms of PPO: PPO-Penalty and PPO-Clip. This article focuses on PPO-Clip since it’s the most common form. PPO-Clip relies on a specialized clipping of the objective function to remove incentives for each step of policy iteration to get too far from the old policy. Intuitively, by limiting how far each step changes the old policy, we can use the old policy’s data for longer.
In PPO-Clip, our objective changes from maximizing expected rewards: to maximizing a clipped version at each iteration :
where
This is a complicated expression, but notice that when , this equals the standard policy iteration (using advantages) that we described above. When the policy changes, the objective changes to maximize the advantage within a boundary. The intuition is that as our policy diverges from the policy that data was collected on, we trust our update less and less. To prevent our policy from drifting too far from the original policy, the part of the objective stops giving more advantages after you change the policy too much (effectively flatlining the gradient and removing the incentive for a big policy change).
The part of the objective is used to undo the effect of clipping in exactly two cases:
- The probability of a worse action increases significantly (, )
- The probability of a good action decreases significantly (, )
In these cases, clipping would hide just how bad the new policy is and make the gradient of our objective for sufficiently bad updates. The helps us restore just how bad the policy is and also restores our gradient, so that way PPO can push the policy in a more optimal direction. Effectively, the caps how good a move can look, but ensures a bad move can’t look good because of the .
The reason this algorithm is much more sample efficient is because of the component. It enables you to use the samples repeatedly to make many updates using the gradient of the objective, as long as you don’t stray too far away from the original policy.