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.