4.4.3Alignment, Prompting & RAG
Proximal Policy Optimization for LLMs
2,019 words9 min readdifficulty · medium
WHY do we even need PPO for LLMs?
After supervised fine-tuning (SFT) an LLM can imitate human answers, but it has no signal about which answers are better — only which are plausible. We want to optimize a preference reward (learned from human rankings). That reward is a black box: we can't backpropagate through "a human liked this". So we treat text generation as a reinforcement learning problem and use policy-gradient methods.
- WHAT we optimize: a policy where a "state" is the prompt + tokens so far, and an "action" is the next token.
- WHY not plain policy gradient? Vanilla policy gradient (REINFORCE) takes huge, noisy steps. One bad batch can push far from anything sensible — for a 7B-parameter LLM that means catastrophic forgetting of language itself.
- WHY PPO specifically? It gives the stability of trust-region methods (small, safe steps) with first-order simplicity (just SGD + a clipped loss).
Building the objective from first principles
Step 1 — The RL objective
We want to maximize expected reward:
Why this step? RL always maximizes expected return; here "return" is the reward-model score.
Step 2 — Policy gradient
The log-derivative trick: since ,
Why this step? We can't differentiate the sampling, but we can differentiate . This turns "push probability toward high reward" into a gradient we can compute.
Step 3 — Baseline & advantage (reduce variance)
Subtract a baseline (a learned value/critic network). Replace raw reward with the advantage :
Why this step? Adding a constant baseline doesn't change the expected gradient (it's mean-zero) but shrinks variance dramatically — the model learns "better/worse than average" instead of "big/small number".
Step 4 — Importance sampling (reuse old data)
We collect rollouts from an old policy then do several gradient steps. To keep the estimate correct we reweight by the probability ratio:
J = \mathbb{E}\big[\rho_t(\theta)\,A_t\big].$$ **Why this step?** LLM rollouts (full generations) are *expensive*. Importance sampling lets us squeeze several updates out of each batch instead of resampling every time. ### Step 5 — The clip (the heart of PPO) Unclipped, a large $\rho_t$ can drive a huge update. PPO **clips the ratio** to $[1-\epsilon,\,1+\epsilon]$ and takes the pessimistic (min) of clipped vs unclipped: > [!formula] PPO clipped surrogate > $$L^{\text{CLIP}}(\theta)=\mathbb{E}_t\Big[\min\big(\rho_t A_t,\ \operatorname{clip}(\rho_t,1-\epsilon,1+\epsilon)\,A_t\big)\Big]$$ > Typically $\epsilon\approx 0.2$. **HOW the clip protects us:** - If $A_t>0$ (good action): the term is capped at $(1+\epsilon)A_t$ → no reward for pushing $\rho_t$ way above $1+\epsilon$. - If $A_t<0$ (bad action): capped at $(1-\epsilon)A_t$ → limited penalty push. - The **min** ensures we always take the *more conservative* estimate, so overshooting is never rewarded. ![[4.4.03-Proximal-Policy-Optimization-for-LLMs.png]] ### Step 6 — Full LLM-RLHF objective For LLMs we add a **per-token KL penalty** to a frozen reference (usually the SFT model) so the policy doesn't drift into reward-hacked gibberish: $$L = \underbrace{L^{\text{CLIP}}}_{\text{maximize reward}}\ -\ \underbrace{c_1\,\mathbb{E}[(V_\theta-V_\text{target})^2]}_{\text{critic/value loss}}\ +\ \underbrace{c_2\,H[\pi_\theta]}_{\text{entropy bonus}}$$ and the reward actually fed to PPO per token is: $$r_t = \underbrace{r_\phi(\text{full response})\,[t=T]}_{\text{terminal RM score}} \;-\; \beta\,\log\frac{\pi_\theta(a_t\mid s_t)}{\pi_\text{ref}(a_t\mid s_t)}.$$ **Why the KL term?** The reward model is only accurate *near* the SFT distribution. Wander too far and the RM is fooled — the KL leash keeps you in-distribution. --- ## Worked examples > [!example] Example 1 — Clip actually biting > Suppose $A_t = +2$ (great token), $\epsilon=0.2$, and the ratio came out $\rho_t=1.6$. > - Unclipped: $\rho_t A_t = 1.6\times 2 = 3.2$. > - Clipped: $\operatorname{clip}(1.6,0.8,1.2)=1.2$, so $1.2\times2 = 2.4$. > - $\min(3.2, 2.4)=2.4$. **Why?** We refuse to reward moving the ratio past $1.2$; gradient here is zero → update stops pushing this token further. > [!example] Example 2 — Bad action, clip does NOT protect > $A_t=-3$, $\rho_t = 0.5$ (policy already dropped this token's prob a lot). > - Unclipped: $0.5\times(-3)=-1.5$. > - Clipped: $\operatorname{clip}(0.5,0.8,1.2)=0.8$, $0.8\times(-3)=-2.4$. > - $\min(-1.5,-2.4) = -2.4$. **Why?** For negative advantage we take the *more negative* (pessimistic) value, so the model is still allowed to keep decreasing the probability — clipping only limits *over-optimistic* moves, not corrective ones. > [!example] Example 3 — KL leash in numbers > A response gets RM score $r_\phi = 1.0$. On a token the policy assigns prob $0.9$ while the reference gave $0.3$: $\log(0.9/0.3)=\log 3\approx1.10$. With $\beta=0.2$, KL penalty $=0.2\times1.10=0.22$. > - Effective token reward at the end $= 1.0 - 0.22 = 0.78$. > **Why?** The model got greedy about a token the reference was unsure of; PPO shaves reward to discourage drift. --- > [!mistake] Steel-manning the classic errors > **Mistake A: "Clipping bounds the policy change directly."** > *Why it feels right:* clip caps the ratio to $[0.8,1.2]$, so surely the policy can only move ~20%. *The fix:* clipping only **zeroes the gradient** once the ratio leaves the band on the *reward-improving* side. Across many tokens/steps the policy **can** still move far — that's why the extra **KL penalty** is needed. Clip = local step-size guard, KL = global leash. > > **Mistake B: "Advantage = reward."** *Why it feels right:* we ultimately want high reward. *The fix:* using raw reward gives massive variance and the model can't tell *relative* quality. Subtracting the critic $V(s)$ (a valid baseline, mean-zero in the gradient) is what makes learning stable. > > **Mistake C: "The min just picks the smaller number, so it always hurts."** *The fix:* the min is over *surrogate values*, and combined with the sign of $A_t$ it yields **zero gradient** exactly when we've already improved enough — it's a smart brake, not a penalty. > > **Mistake D: "Reward hacking is a bug in PPO."** *The fix:* it's a property of optimizing a *learned* reward. PPO's KL term + reference model is precisely the mitigation, not a flaw in the algorithm. --- > [!recall] Active recall — cover the answers > - Why can't we backprop directly through the reward model into $\pi_\theta$? ⟶ generation is discrete sampling; we use the log-derivative (policy gradient) trick instead. > - What does the **clip** reward and what does it refuse to reward? ⟶ rewards reward-improving moves up to $1\pm\epsilon$; refuses further pushing past the band. > - Which term keeps the model in-distribution? ⟶ the KL-to-reference penalty ($\beta$). > - What is the advantage and why subtract a baseline? ⟶ $A=Q-V$; baseline is mean-zero, cuts variance. > [!recall]- Feynman: explain to a 12-year-old > Imagine training a dog. Every time it does a trick, a judge (the reward model) gives it a score. You want the dog to do more of the high-score tricks. But if you give a *giant* treat for one lucky trick, the dog goes wild and forgets its basic manners. So you promise: "I'll only reward you a *little bit more* each time — no crazy jumps." That 'only a little bit more' rule is the **clip**. And you keep the dog on a **leash** tied to how it behaved yesterday (the reference model) so it never runs off into nonsense just to grab treats. Slowly, safely, the dog gets better. > [!mnemonic] Remember PPO with **"CLARK"** > **C**lip the ratio, **L**eash with KL, **A**dvantage (not raw reward), **R**atio = new/old prob, **K**eep steps small. *CLARK keeps the model civilized.* --- ### #flashcards/ai-ml What problem does PPO solve compared to vanilla policy gradient (REINFORCE)? ::: It gives trust-region-like stability (small, safe updates) with first-order simplicity, preventing catastrophic large updates. Write the PPO clipped surrogate objective. ::: $L^{CLIP}=\mathbb{E}[\min(\rho_t A_t,\ \text{clip}(\rho_t,1-\epsilon,1+\epsilon)A_t)]$. Define the probability ratio $\rho_t$. ::: $\rho_t=\pi_\theta(a_t\mid s_t)/\pi_{\theta_{old}}(a_t\mid s_t)$. Why subtract a baseline / use advantage $A=Q-V$? ::: A constant/state baseline is mean-zero in the gradient so it doesn't bias it, but it slashes variance, stabilizing learning. What does the min operator achieve in PPO? ::: It takes the pessimistic (conservative) surrogate, so overshooting the ratio is never rewarded (gradient → 0). Why add a KL penalty to a reference model in RLHF-PPO? ::: The reward model is only reliable near the SFT distribution; KL keeps the policy in-distribution and prevents reward hacking. For $A_t>0$, what is the max the clipped term contributes? ::: $(1+\epsilon)A_t$ — pushing $\rho_t$ higher gives no extra reward. Is clipping alone enough to bound total policy drift in LLMs? ::: No; clip only zeroes gradients locally per token. The KL-to-reference penalty is the global leash. What is the log-derivative trick? ::: $\nabla_\theta\pi_\theta=\pi_\theta\nabla_\theta\log\pi_\theta$, letting us write the policy gradient as $\mathbb{E}[A\,\nabla\log\pi_\theta]$. Typical value of the clip parameter $\epsilon$? ::: About $0.2$. --- ### Connections - [[Reinforcement Learning from Human Feedback (RLHF)]] - [[Reward Modeling from Human Preferences]] - [[KL Divergence]] - [[Policy Gradient Methods & REINFORCE]] - [[Generalized Advantage Estimation (GAE)]] - [[Direct Preference Optimization (DPO)]] (avoids explicit PPO loop) - [[Supervised Fine-Tuning (SFT)]] - [[Reward Hacking & Specification Gaming]] ## 🖼️ Concept Map ```mermaid flowchart TD SFT[SFT model] -->|lacks quality signal| RM[Reward model r_phi] RM -->|black box reward| RL[Frame as RL problem] RL -->|policy pi_theta| OBJ[Maximize expected reward J] OBJ -->|log-derivative trick| PG[Policy gradient] PG -->|subtract baseline V s| ADV[Advantage A s,a] ADV -->|reduces variance| STABLE[Stable learning signal] PG -->|reuse old rollouts| IS[Importance sampling ratio rho] IS -->|prevents huge updates| CLIP[Clipped objective] CLIP -->|trust-region stability| PPO[PPO update] ADV --> CLIP PPO -->|avoids| COLLAPSE[Catastrophic forgetting] ``` ## 🔊 Hinglish (regional understanding) > [!intuition]- Hinglish mein samjho > Dekho, SFT ke baad LLM ko yeh nahi pata ki kaunsa answer "behtar" hai — bas kaunsa "plausible" hai. Toh hum ek reward model banate hain (humans ki ranking se) jo har response ko score deta hai. Ab problem: reward model ke through direct gradient nahi jaa sakta kyunki text generate karna discrete sampling hai. Isliye hum RL use karte hain — policy gradient (log-derivative trick) se hum "high reward wale tokens ki probability badhao" ko ek differentiable gradient mein badal dete hain. > > PPO ka asli jaadu hai **clipping**. Agar hum bina rok-tok reward maximize karein, toh ek hi batch model ko itna door dhakel sakta hai ki woh gibberish likhne lage (7B parameters, catastrophic forgetting). Toh PPO probability ratio $\rho = \pi_{new}/\pi_{old}$ ko $[0.8, 1.2]$ ke beech clip kar deta hai aur min leta hai — matlab agar action accha hai ($A>0$) toh ratio ko $1.2$ se zyada push karne pe koi extra reward nahi milta, gradient zero ho jaata hai. Yeh ek smart brake hai, safe chhote-chhote steps. > > Lekin sirf clip kaafi nahi. Clip toh per-token, local guard hai; poore training mein model phir bhi door ja sakta hai. Isliye hum ek **KL penalty** lagate hain reference (SFT) model se — yeh ek leash hai jo model ko in-distribution rakhti hai, taaki woh reward ko "hack" na kare. Aur variance kam karne ke liye raw reward ki jagah **advantage** ($A = Q - V$, ek critic network) use karte hain — model seekhta hai "average se behtar ya bura", na ki koi random bada number. > > Yaad rakho: **CLARK** — Clip, Leash (KL), Advantage, Ratio, Keep steps small. Yeh paanch cheezein PPO ko civilized aur stable banati hain, aur yahi RLHF ka core engine hai jisse ChatGPT jaise models train hote hain. ![[audio/4.4.03-Proximal-Policy-Optimization-for-LLMs.mp3]]