PPO for Language Models

GPT-3 could write fluent prose about almost anything, yet early user studies found it frequently produced responses that were unhelpful, dishonest, or subtly toxic. The gap between "predicts plausible tokens" and "behaves as intended" is not closed by more pretraining data. It is closed by optimisation against a signal of human preference, and the algorithm that made this practical at scale is Proximal Policy Optimisation (PPO).

What PPO actually optimises

Standard policy gradient methods compute the gradient of expected reward with respect to policy parameters. The update rule is:

∇J(θ) = E_t [ ∇ log π_θ(a_t | s_t) · A_t ]

where A_t is an advantage estimate (how much better action a_t was than average). The problem is that large gradient steps can collapse the policy: one bad update and the distribution shifts so far that all subsequent rollouts are off-distribution, producing a death spiral.

TRPO (Schulman et al., 2015) handled this with a hard KL constraint:

maximise  E_t [ π_θ(a_t | s_t) / π_θ_old(a_t | s_t) · A_t ]
subject to  KL(π_θ_old || π_θ) ≤ δ

PPO (Schulman et al., 2017) achieves the same protective effect more cheaply by clipping the probability ratio r_t(θ) = π_θ / π_θ_old directly in the objective:

L^CLIP(θ) = E_t [ min( r_t(θ) · A_t,  clip(r_t(θ), 1-ε, 1+ε) · A_t ) ]

When A_t > 0 (action was good), the min prevents over-crediting the action beyond the 1+ε boundary. When A_t < 0 (action was bad), the clip prevents punishing it harder than the 1-ε boundary allows. A typical value is ε = 0.2. This single change eliminates the need for a second-order constrained optimisation at every step, making the algorithm GPU-friendly for billion-parameter models.

How the language modelling setup maps onto the RL abstraction

Because a single token is an "action" and episodes are short (typically under 1024 tokens), the RL horizon is compact enough for on-policy PPO rollouts to be tractable. The SFT model is the starting point; PPO nudges it toward higher reward without letting it drift too far from sensible language.

The full KL-regularised objective used in InstructGPT (Ouyang et al., 2022) is:

R_total(x, y) = r_φ(x, y) - β · KL[ π_θ(y|x) || π_ref(y|x) ]

where r_φ is the learned reward model score and β is a coefficient that trades off reward maximisation against proximity to the reference distribution. Without the KL term the policy can find sequences that fool the reward model while producing gibberish or degenerate repetition. With it, the policy is forced to stay in a region of text that the reference model finds probable, which acts as a rough proxy for grammaticality and coherence.

In practice the KL penalty is sometimes applied as an extra reward term at every token position (per-token KL), not just at the end of the response. This provides a denser training signal for the value function and tends to produce more stable training.

The four-model dance of RLHF with PPO

A full RLHF run with PPO keeps four models in memory simultaneously:

1. Policy model (π_θ): the LLM being trained, updated every gradient step.
2. Reference model (π_ref): the frozen SFT checkpoint; used only to compute the KL penalty.
3. Reward model (r_φ): a separate model (often the same architecture with a scalar head) that outputs a single value for a complete response. Trained on human preference pairs before the RL phase starts.
4. Value (critic) model (V_ψ): estimates the expected future reward from each token position; used to compute generalised advantage estimates (GAE). Sometimes initialised from the reward model.

All four must fit on the same (or tightly coordinated) GPU cluster during training. For a 70B-parameter LLM, this means roughly 4 x 70B x 2 bytes = 560 GB of weights in bf16 before accounting for optimiser states. This is the primary engineering constraint that drove the search for lighter alternatives such as GRPO and REINFORCE-leave-one-out variants.

The training loop is:

for each PPO epoch:
    sample prompts from dataset
    rollout: policy generates responses (no grad)
    score: reward model scores each response
    compute per-token KL vs reference
    compute advantages via GAE with value model
    update policy with clipped PPO objective
    update value model with MSE loss

Multiple gradient steps are taken on each batch of rollouts (typically 2-4 PPO epochs per batch), amortising the cost of rollout generation.

Advantage estimation and the credit-assignment problem

One subtlety specific to language models is that the reward signal is sparse: a single scalar appears at the end of a 500-token response. The value function must estimate, at each token position, how much reward the partial sequence is on track to accumulate. Generalised Advantage Estimation (GAE) with λ ≈ 0.95 and γ ≈ 1.0 is the standard choice. With γ = 1 all tokens in an episode are treated as equally discounted (no future discounting), which is appropriate when the "episode" is a single response rather than a temporally extended task.

The value model shares most weights with the policy in many implementations. This keeps parameter count manageable but creates a tension: the policy loss and the value loss pull the shared layers in different directions.

When it falls down

Reward over-optimisation (Goodhart's law). The reward model is an imperfect proxy for human preference. Gao, Schulman, and Hilton (2022) showed empirically that proxy reward increases monotonically with KL budget, but ground-truth reward peaks and then degrades. The policy finds "reward model weaknesses" rather than genuinely better responses. Common symptoms: repetitive qualifiers, verbose hedging, sycophantic preamble, or suspiciously high-scoring outputs that humans find hollow. Mitigations include using an ensemble of reward models, iteratively refreshing the reward model, or capping the KL budget aggressively.

Four-model memory pressure. At scales above roughly 30B parameters, fitting all four models on a single node requires tensor parallelism across devices and careful activation checkpointing. Training throughput drops substantially compared to SFT. This is why research groups with limited compute gravitate to GRPO or Direct Preference Optimisation (DPO), which eliminate the value model entirely.

Policy collapse. If β (the KL coefficient) is set too low, the policy can rapidly shift into a degenerate mode (e.g., always outputting a single safe response). If β is too high, the RL phase produces negligible improvement over SFT. Adaptive KL controllers that target a desired KL divergence (as used in the OpenAI RLHF codebase) help, but require careful initialisation.

Off-policy drift within a batch. PPO is technically on-policy, but reusing rollouts for multiple gradient updates introduces some off-policy error. The clip ratio ε limits this, but with very large models where generating rollouts is expensive and practitioners push for more update steps per rollout batch, the policy can drift enough to invalidate the advantage estimates.

Sparse rewards and long responses. When responses run to several thousand tokens, the per-token advantage estimates become noisy. Reward shaping (e.g., adding per-token length penalties or rule-based partial rewards) helps but must be designed carefully to avoid introducing its own gaming incentives.

Further reading

• Schulman, J. et al. (2017). "Proximal Policy Optimization Algorithms." arXiv:1707.06347. The original PPO paper; Section 3 contains the clipped objective derivation. https://arxiv.org/abs/1707.06347
• Ouyang, L. et al. (2022). "Training language models to follow instructions with human feedback." arXiv:2203.02155. The InstructGPT paper; Appendix C details the four-model PPO setup and KL controller. https://arxiv.org/abs/2203.02155
• Stiennon, N. et al. (2020). "Learning to summarize from human feedback." arXiv:2009.01325. The first large-scale demonstration that PPO-RLHF outperforms SFT for a concrete NLP task. https://arxiv.org/abs/2009.01325
• Gao, L., Schulman, J., and Hilton, J. (2022). "Scaling Laws for Reward Model Overoptimization." arXiv:2210.10760. Quantitative analysis of how proxy reward and ground-truth reward diverge as a function of KL budget. https://arxiv.org/abs/2210.10760