Policy Optimization with PPO

Policy optimization fine-tunes a language model by leveraging a reward model, $r_\phi(x, y)$, which estimates human preferences for different responses $y$ given a prompt $x$. This process involves adjusting the parameters $\theta$ of the language model policy, $\pi_\theta(y|x)$, to generate responses that receive high scores from the reward model. This effectively aligns the model's behavior with the learned preferences. Proximal Policy Optimization (PPO) is the most commonly used algorithm for this policy optimization phase in Reinforcement Learning from Human Feedback (RLHF).

PPO is an on-policy reinforcement learning algorithm designed to achieve stable and efficient policy updates. In the context of LLMs, "on-policy" means that the updates are based on data generated by the current version of the policy $\pi_\theta$ we are trying to improve. The core idea is to maximize the expected reward predicted by $r_\phi$, while simultaneously ensuring the updated policy $\pi_\theta$ doesn't stray too far from a reference policy, typically the initial supervised fine-tuned (SFT) model, denoted as $\pi_{\text{ref}}$. This constraint is important to prevent the model from collapsing into generating repetitive, high-reward but nonsensical text, a phenomenon sometimes called "reward hacking," and to maintain the model's general language capabilities.

The RLHF PPO Objective Function

The optimization process in RLHF typically involves maximizing an objective function that combines the reward signal with a penalty term. For a given prompt $x$ sampled from a dataset $D$, and a response $y$ sampled from the current policy $\pi_\theta(y|x)$, the objective for a single step can be formulated as:

$$\text{Objective}(\theta) = \mathbb{E}_{(x,y) \sim D, \pi_\theta} [ \underbrace{r_\phi(x,y)}_{\text{Reward}} - \underbrace{\beta \cdot \text{KL}(\pi_\theta(y|x) || \pi_{\text{ref}}(y|x))}_{\text{KL Penalty}} ]$$

Let's break down this objective:

1. Reward Term ($r_\phi(x,y)$): This is the score assigned by the trained reward model to the prompt-response pair $(x, y)$. Maximizing this term directly encourages the policy $\pi_\theta$ to generate responses that are predicted to be preferred by humans.
2. KL Penalty Term: This term measures the Kullback-Leibler (KL) divergence between the current policy $\pi_\theta$ and the reference policy $\pi_{\text{ref}}$. The KL divergence quantifies how much the probability distribution of responses generated by $\pi_\theta$ differs from that of $\pi_{\text{ref}}$. Multiplying by the coefficient $\beta$ controls the strength of this penalty.
   • A high $\beta$ forces the updated policy to stay very close to the original SFT model, potentially limiting alignment gains.
   • A low $\beta$ allows the policy to deviate more significantly, potentially leading to better alignment but risking degradation in language quality or coherence if the reward model is imperfect.
   • The reference policy $\pi_{\text{ref}}$ is usually kept frozen during PPO training. Its role is important: it acts as an anchor, preventing the optimized policy $\pi_\theta$ from drifting too far into regions of the policy space that might yield high reward signals but correspond to unnatural or repetitive language.

In practice, PPO includes additional components like value functions and advantage estimation to stabilize training, but the core objective remains balancing the reward maximization against the KL divergence constraint. The full PPO algorithm uses techniques like clipping the policy ratio to prevent excessively large updates in a single step, contributing further to stability.

The PPO Training Loop for LLMs

The PPO training process for LLMs typically proceeds iteratively as follows:

1. Sample Prompts: Draw a batch of prompts $x$ from the prompt dataset $D$.
2. Generate Responses: For each prompt $x$, generate a response $y$ by sampling from the current policy $\pi_\theta(y|x)$.
3. Calculate Rewards: Compute the reward for each pair $(x, y)$ using the frozen reward model: $r = r_\phi(x, y)$.
4. Compute KL Divergence: Calculate the KL divergence between the current policy's output distribution $\pi_\theta(y|x)$ and the reference policy's output distribution $\pi_{\text{ref}}(y|x)$ for each generated response $y$.
5. Estimate Advantages (Simplified View): Determine how much better the obtained reward $r$ is compared to an expected baseline (often estimated using a value function trained alongside the policy).
6. Compute PPO Objective: Calculate the loss based on the PPO objective function, incorporating the rewards, KL penalty, and potentially clipped advantage terms.
7. Update Policy: Perform a gradient update on the parameters $\theta$ of the policy $\pi_\theta$ to maximize the objective (or minimize the negative objective/loss).
8. Repeat: Continue the loop for a specified number of steps or until convergence.

Iterative process of policy optimization using PPO in RLHF. The policy network $\pi_\theta$ generates responses, which are evaluated by the reward model $r_\phi$. The KL divergence relative to a reference model $\pi_{\text{ref}}$ acts as a constraint. These components form the PPO objective, used to update $\pi_\theta$.

Practical Aspects

Implementing PPO for large language models requires careful handling of computational resources. Because it's an on-policy algorithm, new responses need to be generated frequently using the current policy $\pi_\theta$. This involves running inference with the large model within the training loop, which can be computationally intensive.

• Model Synchronization: Keeping the policy $\pi_\theta$, the reference policy $\pi_{\text{ref}}$, the reward model $r_\phi$, and potentially a value function model synchronized and accessible during training requires significant memory, often necessitating distributed training setups (like DeepSpeed ZeRO or FSDP).
• Batching: Prompts are processed in batches, and responses are generated and evaluated together to leverage parallel computation.
• Hyperparameter Sensitivity: The performance of PPO in RLHF is quite sensitive to hyperparameters, particularly the KL coefficient $\beta$, the learning rate, and parameters related to advantage estimation and clipping in the PPO algorithm itself. Tuning these often requires significant experimentation.

By carefully applying PPO, we can steer the LLM's behavior to align better with desired characteristics captured by the human preference data and encoded in the reward model, moving from simple supervised fine-tuning towards more detailed behavioral adjustments. This step is fundamental for creating models that are not only capable but also helpful, harmless, and honest according to specified criteria.