Proximal Policy Optimization for LLMs

4.4.3 · AI-ML › Alignment, Prompting & RAG

---

HUM LLMs ke liye PPO kyun chahte hain?

Supervised fine-tuning (SFT) ke baad ek LLM human answers ki imitation kar sakta hai, lekin uske paas koi signal nahi hota ki kaunse answers better hain — sirf ye pata hota hai ki kaunse plausible hain. Hum ek preference reward $r_{\varphi }$ ko optimize karna chahte hain (jo human rankings se seekha gaya ho). Woh reward ek black box hai: hum "ek human ko ye pasand aaya" ke through backpropagate nahi kar sakte. Isliye hum text generation ko ek reinforcement learning problem ki tarah treat karte hain aur policy-gradient methods use karte hain.

• KYA optimize karte hain: ek policy ${\pi }_{\theta }(a\mid s)$ jahan "state" $s$ prompt + tokens so far hain, aur "action" $a$ next token hai.
• Vanilla policy gradient kyun nahi? Vanilla policy gradient (REINFORCE) bahut bade, noisy steps leta hai. Ek bura batch ${\pi }_{\theta }$ ko kisi bhi sensible cheez se door push kar sakta hai — ek 7B-parameter LLM ke liye iska matlab language ka catastrophic forgetting hai.
• PPO specifically kyun? Ye trust-region methods ki stability (chhote, safe steps) ko first-order simplicity (sirf SGD + ek clipped loss) ke saath deta hai.

---

Objective ko first principles se build karna

Step 1 — RL objective

Hum expected reward maximize karna chahte hain: $J(\theta )={\mathbb{E}}_{a\sim {\pi }_{\theta }}[r(s,a)].$

Ye step kyun? RL hamesha expected return maximize karta hai; yahan "return" reward-model score hai.

Step 2 — Policy gradient

Log-derivative trick: kyunki ${\nabla }_{\theta }{\pi }_{\theta }={\pi }_{\theta }{\nabla }_{\theta }\log {\pi }_{\theta }$, ${\nabla }_{\theta }J={\mathbb{E}}_{a\sim {\pi }_{\theta }}[r(s,a){\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)].$

Ye step kyun? Hum sampling ko differentiate nahi kar sakte, lekin hum $\log {\pi }_{\theta }$ ko differentiate kar sakte hain. Ye "probability ko high reward ki taraf push karo" ko ek aisa gradient bana deta hai jo hum compute kar sakein.

Step 3 — Baseline & advantage (variance kam karo)

Ek baseline $V(s)$ (ek seekha hua value/critic network) subtract karo. Raw reward ki jagah advantage $A(s,a)=Q(s,a)-V(s)$ use karo: ${\nabla }_{\theta }J=\mathbb{E}[A(s,a){\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)].$

Ye step kyun? Ek constant baseline add karna expected gradient ko nahi badalta (ye mean-zero hai) lekin variance ko dramatically kam karta hai — model "average se better/worse" seekhta hai, "bada/chhota number" nahi.

Step 4 — Importance sampling (purana data reuse karo)

Hum ek old policy ${\pi }_{{\theta }_{\text{old}}}$ se rollouts collect karte hain phir kai gradient steps lete hain. Estimate ko sahi rakhne ke liye hum probability ratio se reweight karte hain:

J = \mathbb{E}\big[\rho_t(\theta)\,A_t\big].$$ **Ye step kyun?** LLM rollouts (full generations) *expensive* hote hain. Importance sampling humein har baar resample karne ki jagah har batch se kai updates squeeze karne deta hai. ### Step 5 — Clip (PPO ka dil) Bina clip ke, ek bada $\rho_t$ ek huge update drive kar sakta hai. PPO **ratio ko** $[1-\epsilon,\,1+\epsilon]$ tak **clip karta hai** aur clipped vs unclipped ka pessimistic (min) leta hai: > [!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 clip humein protect karta hai:** - Agar $A_t>0$ (accha action): term $(1+\epsilon)A_t$ par cap hoti hai → $\rho_t$ ko $1+\epsilon$ se bahut upar push karne ka koi reward nahi. - Agar $A_t<0$ (bura action): $(1-\epsilon)A_t$ par cap → limited penalty push. - **Min** ensure karta hai ki hum hamesha *zyada conservative* estimate lete hain, isliye overshoot kabhi reward nahi hota. ![[4.4.03-Proximal-Policy-Optimization-for-LLMs.png]] ### Step 6 — Full LLM-RLHF objective LLMs ke liye hum ek **per-token KL penalty** add karte hain ek frozen reference (usually SFT model) ke liye taaki policy reward-hacked gibberish mein drift na kare: $$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}}$$ aur reward jo actually PPO ko per token feed hota hai woh hai: $$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)}.$$ **KL term kyun?** Reward model sirf SFT distribution ke *paas* accurate hota hai. Bahut door jaao toh RM fool ho jaata hai — KL leash tumhe in-distribution rakhta hai. --- ## Worked examples > [!example] Example 1 — Clip actually biting > Suppose $A_t = +2$ (great token), $\epsilon=0.2$, aur ratio $\rho_t=1.6$ nikla. > - Unclipped: $\rho_t A_t = 1.6\times 2 = 3.2$. > - Clipped: $\operatorname{clip}(1.6,0.8,1.2)=1.2$, toh $1.2\times2 = 2.4$. > - $\min(3.2, 2.4)=2.4$. **Kyun?** Hum ratio ko $1.2$ se aage move karne ka reward dene se mana kar dete hain; yahan gradient zero hai → update is token ko aur push karna band kar deta hai. > [!example] Example 2 — Bad action, clip protect NAHI karta > $A_t=-3$, $\rho_t = 0.5$ (policy ne is token ki prob already bahut gira di). > - 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$. **Kyun?** Negative advantage ke liye hum *zyada negative* (pessimistic) value lete hain, isliye model ko probability kam karte rehne diya jaata hai — clipping sirf *over-optimistic* moves ko limit karta hai, corrective waalon ko nahi. > [!example] Example 3 — KL leash in numbers > Ek response ko RM score $r_\phi = 1.0$ milta hai. Ek token par policy prob $0.9$ assign karti hai jabki reference ne $0.3$ diya tha: $\log(0.9/0.3)=\log 3\approx1.10$. $\beta=0.2$ ke saath, KL penalty $=0.2\times1.10=0.22$. > - Effective token reward at the end $= 1.0 - 0.22 = 0.78$. > **Kyun?** Model ek aisi token ke baare mein greedy ho gaya jiske baare mein reference unsure tha; PPO drift discourage karne ke liye reward ghata deta hai. --- > [!mistake] Classic errors ko steel-man karna > **Mistake A: "Clipping directly policy change ko bound karta hai."** > *Kyun sahi lagta hai:* clip ratio ko $[0.8,1.2]$ tak cap karta hai, toh surely policy sirf ~20% move kar sakti hai. *Fix:* clipping sirf **gradient ko zero karta hai** jab ratio reward-improving side par band se bahar jaata hai. Kai tokens/steps mein policy **phir bhi** bahut door ja sakti hai — isliye extra **KL penalty** chahiye. Clip = local step-size guard, KL = global leash. > > **Mistake B: "Advantage = reward."** *Kyun sahi lagta hai:* hum ultimately high reward chahte hain. *Fix:* raw reward use karne se massive variance aati hai aur model *relative* quality nahi bata sakta. Critic $V(s)$ subtract karna (ek valid baseline, gradient mein mean-zero) hi learning ko stable banata hai. > > **Mistake C: "Min sirf chhota number pick karta hai, toh ye hamesha hurt karta hai."** *Fix:* min *surrogate values* par hai, aur $A_t$ ke sign ke saath milkar ye **zero gradient** deta hai exactly tab jab hum already kaafi improve kar chuke hote hain — ye ek smart brake hai, penalty nahi. > > **Mistake D: "Reward hacking PPO mein ek bug hai."** *Fix:* ye ek *seekhe hue* reward ko optimize karne ki property hai. PPO ka KL term + reference model precisely mitigation hai, algorithm mein koi flaw nahi. --- > [!recall] Active recall — answers cover karo > - Hum reward model ke through directly $\pi_\theta$ mein backprop kyun nahi kar sakte? ⟶ generation discrete sampling hai; hum uski jagah log-derivative (policy gradient) trick use karte hain. > - **Clip** kya reward karta hai aur kya reward karne se mana karta hai? ⟶ reward-improving moves ko $1\pm\epsilon$ tak reward karta hai; band se aage push karne se mana karta hai. > - Kaunsa term model ko in-distribution rakhta hai? ⟶ KL-to-reference penalty ($\beta$). > - Advantage kya hai aur baseline kyun subtract karte hain? ⟶ $A=Q-V$; baseline mean-zero hai, variance kaata hai. > [!recall]- Feynman: ek 12-saal ke bachche ko explain karo > Socho tum ek dog train kar rahe ho. Har baar jab woh koi trick karta hai, ek judge (reward model) use score deta hai. Tum chahte ho ki dog zyada high-score tricks kare. Lekin agar tum ek lucky trick ke liye ek *bahut bada* treat dete ho, toh dog pagal ho jaata hai aur apni basic manners bhool jaata hai. Toh tum promise karte ho: "Main tumhe har baar sirf *thoda sa zyada* reward dunga — koi crazy jumps nahi." Ye 'sirf thoda sa zyada' rule hi **clip** hai. Aur tum dog ko ek **leash** par rakhte ho jo kal ke behavior (reference model) se bandha hai taaki woh sirf treats pakadne ke liye nonsense mein na bhaag jaye. Dheere dheere, safely, dog better hota jaata hai. > [!mnemonic] PPO ko **"CLARK"** se yaad karo > **C**lip the ratio, **L**each with KL, **A**dvantage (not raw reward), **R**atio = new/old prob, **K**eep steps small. *CLARK model ko civilized rakhta hai.* --- ### #flashcards/ai-ml Vanilla policy gradient (REINFORCE) ke comparison mein PPO kaunsi problem solve karta hai? ::: Ye trust-region-jaisi stability (chhote, safe updates) first-order simplicity ke saath deta hai, catastrophic large updates ko rokta hai. PPO clipped surrogate objective likho. ::: $L^{CLIP}=\mathbb{E}[\min(\rho_t A_t,\ \text{clip}(\rho_t,1-\epsilon,1+\epsilon)A_t)]$. Probability ratio $\rho_t$ define karo. ::: $\rho_t=\pi_\theta(a_t\mid s_t)/\pi_{\theta_{old}}(a_t\mid s_t)$. Baseline kyun subtract karte hain / advantage $A=Q-V$ kyun use karte hain? ::: Ek constant/state baseline gradient mein mean-zero hota hai isliye use bias nahi karta, lekin ye variance slash karta hai, learning ko stable banata hai. PPO mein min operator kya achieve karta hai? ::: Ye pessimistic (conservative) surrogate leta hai, taaki ratio ko overshoot karna kabhi reward na ho (gradient → 0). RLHF-PPO mein reference model ko KL penalty kyun add karte hain? ::: Reward model sirf SFT distribution ke paas reliable hota hai; KL policy ko in-distribution rakhta hai aur reward hacking rokta hai. $A_t>0$ ke liye, clipped term maximum kitna contribute karta hai? ::: $(1+\epsilon)A_t$ — $\rho_t$ ko aur upar push karne se koi extra reward nahi milta. Kya LLMs mein total policy drift bound karne ke liye clipping akela kaafi hai? ::: Nahi; clip sirf har token par locally gradients ko zero karta hai. KL-to-reference penalty global leash hai. Log-derivative trick kya hai? ::: $\nabla_\theta\pi_\theta=\pi_\theta\nabla_\theta\log\pi_\theta$, jisse hum policy gradient ko $\mathbb{E}[A\,\nabla\log\pi_\theta]$ ki tarah likh sakte hain. Clip parameter $\epsilon$ ki typical value kya hai? ::: 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)]] (explicit PPO loop se bachta hai) - [[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] ```