Policy Optimization with PPO

Policy optimization is a primary stage in the Reinforcement Learning from Human Feedback (RLHF) pipeline. It follows the training of a reward model $r_\theta(x, y)$, which is designed to capture human preferences. The aim here is to employ this reward model to enhance the language model's policy, $\pi_\phi(y|x)$. Specifically, the goal is to fine-tune the LLM to generate responses $y$ (for a given prompt $x$) that achieve higher scores according to the reward model. This fine-tuning must also preserve the coherent and helpful language capabilities learned during pre-training and initial fine-tuning, which are often referred to as the reference policy, $\pi_{ref}$.

Directly optimizing the LLM using supervised learning on high-reward samples isn't ideal, as it doesn't explore the space of possible outputs efficiently or directly optimize the preference objective. Instead, we frame this as a reinforcement learning problem where the LLM policy $\pi_\phi$ is the agent, the context $x$ is the state, the generated response $y$ is the action (or sequence of actions), and the reward model $r_\theta$ provides the reward signal.

The Reinforcement Learning Objective for RLHF

The objective is to find policy parameters $\phi$ that maximize the expected reward from the reward model, penalized by a Kullback-Leibler (KL) divergence term. The KL penalty prevents the policy $\pi_\phi$ from moving too far away from the initial reference policy $\pi_{ref}$, which helps maintain language quality and prevents the model from finding degenerate solutions that exploit the reward model.

The optimization objective can be formulated as:

$$\text{maximize}_\phi \quad J(\phi) = \mathbb{E}_{x \sim D, y \sim \pi_\phi(y|x)} \left[ r_\theta(x, y) - \beta D_{KL}(\pi_\phi(y|x) || \pi_{ref}(y|x)) \right]$$

Let's break down this objective:

• $x \sim D$: Prompts $x$ are sampled from a distribution $D$ (often the same distribution used for collecting preference data).
• $y \sim \pi_\phi(y|x)$: Responses $y$ are generated by the current LLM policy $\pi_\phi$ that we are optimizing.
• $r_\theta(x, y)$: The scalar reward assigned to the prompt-response pair $(x, y)$ by the trained reward model. This is the primary signal indicating alignment with human preferences.
• $D_{KL}(\pi_\phi(y|x) || \pi_{ref}(y|x))$: The KL divergence between the distribution of responses generated by the current policy $\pi_\phi$ and the reference policy $\pi_{ref}$. This term measures how much the current policy has diverged from the original one. The reference policy is typically the model before RLHF fine-tuning, such as the supervised fine-tuned (SFT) model.
• $\beta$: A hyperparameter that controls the strength of the KL penalty. A higher $\beta$ discourages deviation from $\pi_{ref}$, while a lower $\beta$ allows for more aggressive optimization towards the reward model's preference.

This objective encourages the model to generate outputs that the reward model prefers, but not at the cost of generating text that is statistically very different from what the initial, well-behaved model $\pi_{ref}$ would produce.

Why Proximal Policy Optimization (PPO)?

While standard policy gradient methods like REINFORCE could theoretically be used, they often suffer from high gradient variance and instability, especially when dealing with the large parameter spaces and complex action spaces (sequences of tokens) of LLMs. Updates can be destructive, where a single large gradient step significantly harms the policy's performance.

Proximal Policy Optimization (PPO) has emerged as a popular and effective algorithm for RLHF. It belongs to the family of trust region optimization methods, which aim to improve stability by constraining the amount the policy changes in each update step. PPO achieves this stability without the heavy computational cost of some other trust region methods like TRPO.

The main advantages of PPO in this context include:

1. Stability: PPO uses a clipped surrogate objective function that discourages large policy updates, leading to more stable training.
2. Sample Efficiency: It allows for multiple epochs of gradient updates on the same batch of sampled data, improving data utilization compared to single-update methods.
3. Simplicity: While sophisticated, its implementation is relatively straightforward compared to alternatives like TRPO.

PPO's Clipped Surrogate Objective

At the foundation of PPO is its objective function. Let $\pi_{\phi_{old}}$ be the policy before the update (the policy used to generate the data for the current optimization step). PPO optimizes the policy $\pi_\phi$ based on the following objective:

$$L^{CLIP}(\phi) = \mathbb{E}_{(x, y) \sim \pi_{\phi_{old}}} \left[ \min(r(\phi) \hat{A}(x, y), \text{clip}(r(\phi), 1-\epsilon, 1+\epsilon) \hat{A}(x, y)) \right]$$

Let's unpack the components relevant to RLHF:

• Probability Ratio: $r(\phi) = \frac{\pi_\phi(y|x)}{\pi_{\phi_{old}}(y|x)}$. This ratio measures how likely the response $y$ is under the new policy $\pi_\phi$ compared to the old policy $\pi_{\phi_{old}}$. In practice, we compute this using log probabilities: $\exp(\log \pi_\phi(y|x) - \log \pi_{\phi_{old}}(y|x))$.
• Advantage Estimate: $\hat{A}(x, y)$. This estimates how much better the response $y$ is compared to the average response expected from prompt $x$. In the context of RLHF, a common approach is to use the KL-penalized reward as the advantage: $$\hat{A}(x, y) \approx r_\theta(x, y) - \beta D_{KL}(\pi_\phi(y|x) || \pi_{ref}(y|x))$$ Sometimes, a baseline value function $V(x)$ (estimating the expected reward given prompt $x$) is subtracted from $r_\theta(x, y)$ to reduce variance, but simpler implementations often use the reward model score directly (potentially normalized or combined with the KL term). Let's denote the combined reward signal used for optimization as $R(x,y) = r_\theta(x, y) - \beta (\log \pi_\phi(y|x) - \log \pi_{ref}(y|x))$. So, $\hat{A}(x,y) \approx R(x,y)$.
• Clipping: The $\text{clip}(r(\phi), 1-\epsilon, 1+\epsilon)$ function constrains the probability ratio $r(\phi)$ to stay within the interval $[1-\epsilon, 1+\epsilon]$. $\epsilon$ is a small hyperparameter (e.g., 0.2).
  • If the advantage $\hat{A}$ is positive (meaning the response $y$ was better than expected), the objective increases as $r(\phi)$ increases, encouraging the policy update. However, the $\min$ and the upper clipping boundary $1+\epsilon$ prevent $r(\phi)$ from becoming too large, thus limiting the update step size.
  • If the advantage $\hat{A}$ is negative (the response was worse than expected), the objective decreases as $r(\phi)$ increases. The lower clipping boundary $1-\epsilon$ prevents $r(\phi)$ from becoming too small, again limiting how much the policy can change in a single step.

This clipping mechanism effectively creates a pessimistic bound on the policy update, ensuring that the new policy doesn't deviate too drastically from the old one in a way that could destabilize training.

The PPO Loop for LLM Fine-Tuning

Applying PPO to fine-tune an LLM involves an iterative process:

The RLHF PPO optimization cycle. Prompts are sampled, responses generated by the current policy, then scored by both the reward model and the reference policy to calculate a KL-penalized reward. This data batch is used to update the policy via the PPO algorithm.

1. Sample Prompts: Draw a batch of prompts $x$ from the dataset $D$.
2. Generate Responses: For each prompt $x$, generate a response $y$ using the current policy $\pi_\phi$. Store the generated sequences $(x, y)$ and their log-probabilities $\log \pi_\phi(y|x)$. This $\pi_\phi$ becomes $\pi_{\phi_{old}}$ for the optimization step.
3. Score Responses: Use the fixed reward model $r_\theta$ to compute the reward score $r_\theta(x, y)$ for each generated pair.
4. Compute Reference Log-Probabilities: Calculate the log-probabilities of the generated responses $y$ under the fixed reference policy $\pi_{ref}$, yielding $\log \pi_{ref}(y|x)$.
5. Calculate Final Reward: Compute the reward signal used for optimization, typically incorporating the KL penalty: $R(x, y) = r_\theta(x, y) - \beta (\log \pi_\phi(y|x) - \log \pi_{ref}(y|x))$. The KL term is estimated using the difference in log-probabilities between the current/old policy and the reference policy.
6. Optimize Policy: Perform multiple epochs of gradient ascent on the PPO clipped surrogate objective $L^{CLIP}(\phi)$ using the collected batch of experiences $(x, y, R(x, y), \log \pi_{\phi_{old}}(y|x))$. An optimizer like Adam or AdamW is commonly used to update the parameters $\phi$ of the language model.

This cycle repeats, iteratively refining the policy $\pi_\phi$ to produce responses that achieve higher rewards according to $r_\theta$ while staying close to $\pi_{ref}$.

Practical Notes

• KL Coefficient ($\beta$): Selecting an appropriate value for $\beta$ is significant. If $\beta$ is too small, the policy might over-optimize for the reward model, potentially leading to repetitive or nonsensical outputs that exploit reward model weaknesses, or diverging too far from the general language capabilities captured in $\pi_{ref}$. If $\beta$ is too large, the policy updates will be overly constrained, and the model might not improve sufficiently in terms of desired alignment characteristics. $\beta$ can be a fixed hyperparameter or sometimes adjusted dynamically during training based on the observed KL divergence.
• Model Copies: Efficient implementation requires managing several copies or states of the model: the policy being actively trained ($\pi_\phi$), the policy used for generating the current batch of experience ($\pi_{\phi_{old}}$), and the frozen reference policy ($\pi_{ref}$).
• Computational Resources: The PPO step is computationally demanding. It involves:
  • A forward pass through $\pi_\phi$ to generate responses.
  • A forward pass through the (potentially large) reward model $r_\theta$.
  • A forward pass through the reference model $\pi_{ref}$.
  • Multiple forward and backward passes through $\pi_\phi$ during the PPO optimization epochs.

In summary, PPO provides an effective and widely adopted method for the policy optimization phase of RLHF. By using the signal from a learned reward model and incorporating a KL divergence penalty within a stable optimization framework, PPO enables fine-tuning large language models to better align with human preferences while preserving their core language understanding and generation abilities.