PPO for LLMs Math

Generation as a reinforcement-learning problem

PPO reframes text generation as a sequential decision process. At step t the state s_t is the prompt plus everything generated so far, s_t = (x, y_1, …, y_{t-1}). The action a_t = y_t is the next token drawn from the vocabulary. The policy is the language model itself: π_θ(a_t | s_t) is exactly the softmax over next-token logits. An episode is one full response, running until an end-of-sequence token or a length cap.

This mapping is what lets an RL algorithm designed for robots and games drive a transformer. The environment is deterministic — appending a token to a prefix always yields the same next prefix — so all the stochasticity lives in the policy’s sampling. The catch that shapes everything downstream is that the reward is sparse and delayed: you cannot judge a half-written answer, so the real signal arrives only when the response is complete. PPO’s machinery — a critic, advantages, and a trust region — exists largely to route that single terminal signal back to every token that helped earn it.

The reward: reward-model score minus a KL penalty

The per-token reward has two parts. A separately trained reward model r_RM(x, y) scores the whole response and is added at the final token only. Every token also pays a KL penalty for straying from a frozen reference policy π_ref (the supervised-fine-tuned model you started from):

r_t = r_RM(x, y) · 1[t = T]  −  β · ( log π_θ(y_t | s_t) − log π_ref(y_t | s_t) )

The bracketed term is a single-sample estimate of the per-token KL divergence KL(π_θ || π_ref), and β sets how hard the leash pulls. Without it, PPO would happily reward-hack — find degenerate, repetitive, or gibberish text that the imperfect reward model rates highly — while destroying fluency. The KL term makes the objective a constrained one: maximize reward while staying close to the reference. Many implementations adapt β on the fly, raising it when measured KL overshoots a target, keeping the divergence in a healthy band throughout training.

The value head: a critic bolted onto the model

To turn one terminal reward into a useful per-token learning signal, PPO needs a baseline: how good is a state on average? That is the job of the value function V_φ(s_t), the expected total future reward from state s_t. In practice it is a critic head — a single linear layer projecting the transformer’s hidden state to one scalar — sitting on top of the policy backbone, sometimes sharing weights with it and sometimes a full separate copy.

The critic is trained by regression toward the observed returns. If R_t is the return target (below), the value loss is a simple mean-squared error:

L_VF(φ) = ( V_φ(s_t) − R_t )^2

A good critic makes credit assignment tractable: by subtracting V(s_t) from returns we ask not ‘was the outcome good?’ but ‘was this token better than expected from here?’ — a far lower-variance question.

GAE: turning rewards into advantages

The advantage A_t measures how much better action a_t was than the critic’s baseline. Generalized Advantage Estimation (GAE) computes it from the per-step TD residual and an exponentially weighted sum:

δ_t = r_t + γ · V(s_{t+1}) − V(s_t)
A_t  = δ_t + (γλ) δ_{t+1} + (γλ)^2 δ_{t+2} + … = Σ_{l≥0} (γλ)^l δ_{t+l}

Here γ is the discount (often 1.0 for short responses) and λ ∈ [0,1] trades bias against variance: λ = 0 gives the low-variance, critic-dependent one-step estimate δ_t, while λ = 1 gives the unbiased but high-variance Monte-Carlo return. Values around 0.95 sit in the sweet spot. The return target the critic regresses on falls straight out: R_t = A_t + V(s_t). Advantages are then normalized to zero mean and unit variance across the batch, which keeps gradient magnitudes stable regardless of the reward model’s arbitrary scale.

The clipped surrogate objective

Vanilla policy gradients take one gradient step per batch of experience, which is wasteful. PPO wants several epochs of updates on the same rollouts — but reusing data means the policy being optimized, π_θ, drifts from the policy that collected the data, π_θ_old. PPO corrects for this with an importance-sampling ratio and then clips it to forbid overlarge steps:

ratio_t(θ) = π_θ(a_t|s_t) / π_θ_old(a_t|s_t)

L_CLIP(θ) = E_t[ min( ratio_t · A_t ,
                          clip(ratio_t, 1−ε, 1+ε) · A_t ) ]

with ε ≈ 0.2. The intuition: when A_t > 0 (a good action) the objective wants to raise its probability, but the min caps the gain once ratio_t exceeds 1+ε — no reward for moving too far. When A_t < 0 the clip floors the ratio at 1−ε. Either way the update is confined to a trust region near the old policy — the ‘Proximal’ in PPO. The full loss adds the value term and a small entropy bonus: L = L_CLIP − c_1 L_VF + c_2 · entropy.

A worked example, end to end

Take a two-token response with β = 0.1, γ = 1, λ = 0.95. The reward model scores the finished text r_RM = 1.5. Suppose the log-probabilities give KL log-ratios of +0.2 at token 1 and −0.2 at token 2, and the critic predicted V(s_1)=0.8, V(s_2)=1.0, terminal V(s_3)=0. Per-token rewards:

r_1 = −0.1·(0.2)               = −0.02
r_2 = 1.5 − 0.1·(−0.2)   =  1.52

δ_2 = 1.52 + 0 − 1.0            =  0.52
δ_1 = −0.02 + 1.0 − 0.8   =  0.18

A_2 = δ_2                     = 0.52
A_1 = δ_1 + 0.95·δ_2   = 0.18 + 0.494 = 0.674

The return targets are R_1 = 0.674 + 0.8 = 1.474 and R_2 = 1.52, which the critic regresses toward. Now the clip: token 1 has A_1 = 0.674 > 0, so PPO wants to raise its probability. If a step pushes ratio_1 = 1.3, the clip caps it at 1.2, so the objective uses min(1.3·0.674, 1.2·0.674) = min(0.876, 0.809) = 0.809. The excess step earns nothing — the trust region held the line, exactly as designed.

The training loop in one turn of the crank

One PPO iteration is a four-beat cycle. Rollout: sample a batch of prompts and let π_θ_old generate responses, caching each token’s log-probability. Score: run the reward model on the finished responses and the reference model on every token to build the per-token reward, then the critic to get values. Estimate: compute GAE advantages and return targets, and normalize the advantages. Optimize: for a few epochs, take minibatch gradient steps on the clipped surrogate plus value and entropy terms.

A subtlety worth naming: π_θ_old is not a fifth network. It is simply the policy weights snapshotted at rollout time — in practice just the cached log-probabilities from the generation step. Once the epochs finish, the updated π_θ becomes the new π_θ_old and the crank turns again. Because fresh data is generated from the current policy each iteration, PPO is squarely on-policy — the property that makes it powerful and, as the next section shows, expensive.

The four-model memory bill

The signature cost of PPO-RLHF is that four models sit in memory at once:

The asymmetry is the real story. The two frozen models (reference, reward) need only weights and activations for a forward pass. The two trainable models (policy, critic) each also carry gradients and Adam moment buffers — roughly several times their parameter count in optimizer state alone. So ‘four models’ understates the trainable side, but the headline holds: PPO-RLHF needs far more memory than ordinary fine-tuning, which is why it is heavy for small teams and out of reach on a CPU-class SLM budget without aggressive sharing.

PPO versus DPO: with and without the loop

Direct Preference Optimization (DPO) targets the same goal — align a model to human preferences — but deletes the RL loop entirely. It starts from a dataset of preferred vs rejected response pairs and shows, via the same KL-constrained objective PPO optimizes, that the optimal policy can be reached by a supervised classification loss on those pairs. There is no sampling during training, no reward model, no critic, no advantages, no clip.

The contrast is stark. DPO keeps only two models in memory (policy and reference), trains on a fixed offline dataset, and is stable and cheap — which is why it has become the default for resource-constrained alignment. PPO’s price buys one thing DPO gives up: because it generates fresh responses and scores them online, PPO can keep improving against a reward model on distributions its preference data never covered, and can optimize signals — a verifier, a tool-use checker, an execution reward — that have no natural pairwise form. This article covers that on-policy PPO pipeline; DPO is the loop-free alternative for when the four-model bill is too steep.

Pitfalls and practical notes

PPO-RLHF is famously finicky, and most failures trace to a handful of causes. Reward hacking is the constant enemy: the reward model is a proxy, and the policy will exploit its blind spots — the KL penalty is the main defense, so watch measured KL like a vital sign and adapt β if it runs away. Value collapse, where the critic fits poorly and poisons every advantage, shows up as unstable or exploding policy loss; advantage normalization and value-loss clipping help. Setting ε too large defeats the trust region and lets the policy lurch; too small and learning crawls.

The deeper lesson is that PPO’s complexity is not incidental — each piece earns its place. The critic and GAE exist to route a sparse terminal reward to individual tokens; the clip exists to reuse expensive rollouts safely; the KL penalty and reference model exist to keep a capable base model from being wrecked in pursuit of a flawed score. Understanding why each component is there is what lets you debug a run that has quietly started producing confident nonsense.