Technische Samenvatting: Robuuste Policy-Optimalisatie om Catastrofaal Vergeten te Voorkomen

Probleemstelling

Grote Taalmodellen (LLMs) worden doorgaans getraind via meerstaps-pijplijnen, beginnend met Reinforcement Learning from Human Feedback (RLHF) om modellen af te stemmen op veiligheid en behulpzaamheid, gevolgd door fine-tuning voor specifieke downstream-taken (bijv. wiskundig redeneren, coderen). Een kritieke kwetsbaarheid in dit paradigma is catastrofaal vergeten: zelfs geringe downstream-updates kunnen eerder aangeleerde gedragingen, zoals veiligheidsbeperkingen, verslechteren.

Vorig onderzoek heeft zich grotendeels gericht op mitigatiestrategieën op het moment van downstream (bijv. herhaling, parameter-efficiënte fine-tuning, regularisatie of modelmerging) om capaciteiten achteraf te behouden. De auteurs betogen echter dat deze methoden ontoereikend zijn omdat ze afhankelijk zijn van specifieke downstream-procedures. Als een model later wordt onderworpen aan "gewone" fine-tuning die afwijkt van deze specifieke protocollen, blijven eerder aangeleerde gedragingen breekbaar. De kernuitdaging is ervoor te zorgen dat het basisbeleid zelf robuust is voor toekomstige aanpassingen, ongeacht de downstream-taak of het gebruikte algoritme.

Methodologie: Fine-tuning Robust Policy Optimization (FRPO)

De auteurs stellen Fine-tuning Robust Policy Optimization (FRPO) voor, een kader dat is ontworpen om het basisbeleid te optimaliseren voor stabiliteit tegen toekomstige downstream-verschuivingen.

1. Probleemformulering: Max-Min Robuustheid

De auteurs formaliseren het probleem als een max-min optimalisatie. In plaats van de verwachte beloning uitsluitend te maximaliseren bij het huidige beleid $\pi_\theta$, streeft de doelstelling ernaar de minimale beloning te maximaliseren die haalbaar is door elk downstream-beleid $Q$ dat ligt binnen een KL-gebonden omgeving van $\pi_\theta$.

Het optimalisatieprobleem is gedefinieerd als:
$$\max_{\pi_\theta} \left( \inf_{Q: \mathbb{E}[\text{KL}(Q \parallel \pi_\theta)] \le \rho} \mathbb{E}_{x,y \sim Q}[r(x,y)] - \beta \text{KL}(\pi_\theta \parallel \pi_{\text{ref}}) \right)$$

Hier vertegenwoordigt $\rho$ de straal van de KL-omgeving, wat standaardbeperkingen in downstream-fine-tuning weerspiegelt (bijv. vertrouwensgebieden, kleine leersnelheden of LoRA-beperkingen). De beperking wordt opgelegd in verwachting over de prompt-distributie, waardoor significante verschuivingen op prompts die relevant zijn voor de taak mogelijk zijn, terwijl de gemiddelde nabijheid tot het basisbeleid wordt behouden.

2. Theoretische Afleiding

Met behulp van dualiteitstheorie voor distributie-robuste optimalisatie (specifiek onder een $f$-divergentie-beperking) leiden de auteurs een duale vorm af van het binnenste minimaliseringsprobleem.

• Entropisch Risicodoel: De duale formulering onthult dat het robuuste doel gelijkwaardig is aan het optimaliseren van een entropisch risico-doel. De binnenste minimalisatie over $Q$ transformeert het probleem tot het maximaliseren van:
  $$-\lambda \log \mathbb{E}_{y \sim \pi_\theta} \left[ e^{-r(x,y)/\lambda} \right]$$
  waarbij $\lambda$ een instelbare hyperparameter is die risicosensitiviteit controleert.
• Interpretatie: Dit doel straft trajecten met lage beloning zwaarder dan standaard maximalisatie van verwachte beloning. Het zoekt effectief naar "vlakke" gebieden in het beloningslandschap binnen de beleidsruimte, zodat kleine verstoringen (downstream-fine-tuning) niet leiden tot scherpe dalingen in prestaties.
• Grenzgeval: Als $\lambda \to \infty$, herstelt het doel het standaard GRPO (Group Relative Policy Optimization)-doel, dat de verwachte beloning maximaliseert.

3. Algorithmische Implementatie

De auteurs passen het GRPO-algoritme aan om dit nieuwe doel te optimaliseren zonder extra rekenkosten te introduceren (bijv. geen geleerde criticus of extra forward passes).

• Importance Sampling: De verwachting onder het nieuwe beleid wordt geschat met behulp van steekproeven uit een vertraagd beleid $\pi_{\text{old}}$ via importance sampling.
• Clipping en Baselines: Om stabiliteit te waarborgen, past het algoritme PPO-stijl clipping toe op de importance ratios. Cruciaal leiden de auteurs een baseline-term af om de hoge variantie in de gradiënt van de log-partitiefunctie te cancelen wanneer $\lambda$ groot is.
• Bias-reductie: De Monte Carlo-schatting van de log-partitiefunctie introduceert een $O(1/G)$ bias (waarbij $G$ de groepsgrootte is). De auteurs maken gebruik van een jackknife-techniek (leave-one-out-schatting) om deze bias te reduceren tot $O(1/G^2)$, waardoor stabiele convergentie wordt gewaarborgd zelfs bij kleine groepsgroottes.

Belangrijkste Bijdragen

1. Robuust RLHF-kader: Het artikel introduceert een max-min formulering die rekening houdt met alle downstream-verschuivingen binnen een KL-bal. Het toont aan dat deze formulering wiskundig gelijkwaardig is aan het optimaliseren van een entropisch-risicodoel, wat een theoretische basis biedt voor beloningsvlakheid in de beleidsruimte.
2. FRPO-algoritme: Een praktische, directe vervanging voor GRPO die geen extra berekening vereist. Het integreert naadloos met bestaande RLHF-pijplijnen, bevat een afgeleide baseline voor stabiliteit, en herstelt standaard GRPO als $\lambda \to \infty$.
3. Empirische Validatie: De methode wordt geëvalueerd in diverse scenario's, waaronder veiligheidsafstemming en wiskundig redeneren, en toont een superieure behoud van eerdere capaciteiten aan in vergelijking met standaard fine-tuning en andere robuustheidsbaselines.

Experimentele Resultaten

De auteurs evalueren FRPO in twee primaire scenario's: Veiligheidstraining en Wiskundig Redeneren.

1. Veiligheidstraining (Veiligheidsbeperkingen)

• Opzet: Modellen (Mistral-v0.1 en Qwen2.5) werden getraind op een veiligheidsdataset met FRPO met variërende $\lambda$, en vervolgens gefine-tuned op downstream-taken (Alpaca voor instructievolgen, GSM8K voor wiskunde en UltraFeedback voor RL).
• Vondsten:
  • Met FRPO getrainde modellen ($\lambda \approx 0,2$) behielden na fine-tuning aanzienlijk hogere veiligheidsweigeringspercentages en lagere schadelijkheidsscores (StrongREJECT) in vergelijking met standaard GRPO.
  • De methode behield veiligheidsbeperkingen effectief, zelfs onder agressieve fine-tuning-instellingen (full-parameter SFT, verschillende leersnelheden, LoRA).
  • Trade-off-analyse: Hoewel een kleinere $\lambda$ de veiligheid beter behield, verbeterde een grotere $\lambda$ de behulpzaamheid enigszins ten koste van de veiligheid. De optimale $\lambda$ bleek taakafhankelijk te zijn, maar over het algemeen laag voor veiligheid.
  • Downstream-prestaties: Cruciaal compromitteerde FRPO de prestaties van downstream-taken (bijv. GSM8K-nauwkeurigheid, MMLU-scores) niet in vergelijking met GRPO, wat bevestigt dat het behoud van veiligheid niet het gevolg was van een falen om de nieuwe taak te leren.

2. Continual Learning (Wiskunde naar Code)

• Opzet: Modellen werden eerst getraind op wiskundig redeneren (MATH-dataset) en vervolgens gefine-tuned op codegeneratie (OpenCodeInstruct).
• Vondsten:
  • FRPO met $\lambda = 2,0$ (geoptimaliseerd voor de hogere variantie in wiskundebeloningen) behield de wiskundenauwkeurigheid aanzienlijk beter dan GRPO.
  • Specifiek behield FRPO 22% hogere nauwkeurigheid op de MATH500-benchmark na code-fine-tuning in vergelijking met GRPO, terwijl het vergelijkbare winsten boekte op coderingsbenchmarks (MBPP+).

3. Ablatiestudies

• Vergelijking met SAM: De auteurs vergeleken FRPO met Sharpness-Aware Minimization (SAM). Hoewel SAM (werkend in parameter-ruimte) enige verbetering bood ten opzichte van GRPO, was FRPO (werkend in beleids-/KL-ruimte) superieur. De auteurs merken op dat de L2-bal van SAM in parameter-ruimte downstream-beleidsverschuivingen niet volledig vastlegt, terwijl de KL-bal output-distributieverschuivingen direct beperkt.
• Rol van $\lambda$: Experimenten bevestigden dat $\lambda$ de trade-off tussen gemiddelde beloning en variantie controleert. De optimale $\lambda$ correleert met de standaardafwijking van beloningen en het KL-budget $\rho$.

Betekenis en Beweringen

Het artikel beweert dat robuustheid tegen downstream-fine-tuning direct in het basisbeleid tijdens RLHF moet worden opgenomen, in plaats van te vertrouwen op latere interventies. Door te optimaliseren voor beloningsstabiliteit binnen een KL-gebonden omgeving, identificeert FRPO beleidslijnen die zich bevinden in "belonings-vlakke" gebieden, waardoor ze inherent bestand zijn tegen catastrofaal vergeten.

De auteurs positioneren FRPO als een algemene oplossing die:

• Werkt in diverse domeinen (veiligheid, wiskunde, coderen).
• Geen extra rekenkosten vereist in vergelijking met standaard GRPO.
• Een principiële, wiskundig onderbouwde aanpak biedt voor continual learning in LLM's die agnostisch is ten opzichte van het specifieke downstream-fine-tuning-protocol.

Het werk suggereert dat het principe van optimalisatie voor robuustheid van toekomstige aanpassingen verder kan reiken dan RLHF naar pretraining en supervised fine-tuning, en een nieuwe richting biedt voor onderzoek naar afstemming en continual learning.