Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

Titel: Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

Auteurs: Jiacheng Wang, Ping Jian, et al. (Beijing Institute of Technology)

---

1. Het Probleem: Catastrophical Forgetting bij Off-Policy SFT

Grote Taalmodellen (LLMs) presteren uitstekend in algemene taken, maar wanneer ze worden aangepast aan specifieke domeinen (zoals wiskundig redeneren) via Supervised Fine-Tuning (SFT), treedt vaak catastrophical forgetting op. Dit betekent dat het model zijn eerder verworven algemene vaardigheden verliest.

De kern van het probleem ligt in de distributiemismatch:

• Off-Policy Aard: De downstream data (expert demonstraties) vertonen een aanzienlijke verschuiving ten opzichte van de natuurlijke generatiedistributie van het basismodel.
• Instabiliteit: Wanneer het model wordt getraind op data die onder de huidige beleidspolicy ($\pi_\theta$) een lage waarschijnlijkheid heeft, leidt dit tot instabiele updates en grote gradiëntvariaties.
• Beperkingen van Bestaande Oplossingen: Bestaande methoden voor data rewriting (het herschrijven van trainingsdata om deze beter te laten aansluiten bij het model) gebruiken vaak vaste prompt-sjablonen. Dit resulteert in:
  1. Een mismatch met de natuurlijke "QA-stijl" (vraag-antwoord) generatie van het model.
  2. Een verlies aan diversiteit (diversity collapse), waarbij het model in herhalende, sjabloon-achtige patronen terechtkomt.

2. Methodologie: Een RL-gebaseerde Rewriting Agent

De auteurs stellen een nieuwe aanpak voor die data-rewriting vormgeeft als een beleidsleerprobleem (policy learning) in plaats van een statische prompt-engineering taak.

Het Framework

Het systeem bestaat uit twee fasen:

1. Fase I: Training van de Rewriting Agent ($R_\phi$)

   • Architectuur: Een geïntegreerde agent wordt getraind bovenop een bevroren basismodel ($\pi_0$) met behulp van LoRA (Low-Rank Adaptation). Dit zorgt voor een parameter-efficiënte "patch" die lokale correcties toepast zonder de basisdistributie volledig te verstoren.
   • Optimalisatie: De agent wordt getraind met GRPO (Group Relative Policy Optimization), een variant van PPO die geen aparte waarde-netwerk vereist.
   • Doel: De agent leert een herschreven antwoord ($\tilde{y}$) te genereren dat voldoet aan drie criteria: taakconsistentie, distributie-uitlijning en diversiteit.

2. Fase II: Dataset Constructie en Downstream SFT

   • De getrainde agent genereert een herschreven dataset ($D_R$).
   • Er wordt een Generate-Verify-Fallback pipeline gebruikt: Als een herschreven antwoord faalt in de taakverificatie, wordt het originele expert-antwoord behouden.
   • Het doelmodel ($\pi_\theta$) wordt vervolgens standaard getraind (SFT) op deze geoptimaliseerde dataset.

De Unificatie van Doelstellingen (Reward Functie)

De kern van de methode is een unified reward die drie componenten combineert, waarbij taakconsistentie fungeert als een harde poort (hard gate):

1. Taakconsistentie ($r_{task}$): Een harde binaire poort. Een herschreven antwoord moet zowel het juiste eindantwoord hebben als geldig redeneren (gecontroleerd via een regelgebaseerde verifier en een LLM-judge). Alleen als dit slaagt, worden de andere beloningen toegepast.
2. Distributie-uitlijning ($r_{dist}$): De herschreven antwoorden moeten waarschijnlijk zijn onder de QA-stijl distributie van het basismodel ($\pi_0(\cdot | x)$), zonder de extra prompt-context van de herschrijftaak. Dit wordt gemeten via de genormaliseerde negatieve log-likelihood (NLL).
3. Diversiteit ($r_{div}$): Om mode-collapse te voorkomen, wordt een beloning gegeven voor semantische diversiteit binnen een groep van kandidaat-antwoorden voor dezelfde input. Dit wordt berekend op basis van embedding-afstanden.

De totale beloning is:
$$r = r_{task} + r_{task} \cdot (\lambda_{dist} r_{dist} + \lambda_{div} r_{div})$$
Dit zorgt ervoor dat alleen geldige antwoorden worden beloond op basis van hun kwaliteit en diversiteit.

3. Belangrijkste Bijdragen

• Inzicht in Distributiemismatch: De auteurs identificeren dat bestaande rewriting-methoden falen omdat ze conditioneren op een prompt-gedreven distributie in plaats van de natuurlijke QA-generatiedistributie van het model.
• RL-gebaseerd Rewriting Agent: Ze introduceren een agent die via Reinforcement Learning (GRPO) leert om supervisie-data te herschrijven die beter aansluit bij het basismodel, terwijl taakconsistentie en diversiteit worden gewaarborgd.
• Unificatie van Doelen: Een nieuw objectief dat taakconsistentie (als harde constraint), distributie-uitlijning en diversiteit simultaan optimaliseert.
• Empirische Validatie: Uitgebreide experimenten tonen aan dat de methode prestaties op de doeltaak behoudt die vergelijkbaar zijn met standaard SFT, maar het verlies aan algemene vaardigheden (catastrophical forgetting) aanzienlijk reduceert.

4. Resultaten

De methode werd getest op drie verschillende backbone-modellen (Llama-3.2-1B, Llama-3.2-3B, Mistral-7B) met wiskundig redeneren als downstream taak.

• Prestaties op Doeltaak: De "Rewriting-agent" behaalde prestaties op wiskundebenchmarks (Math500, AMC23, etc.) die vergelijkbaar waren met, of zelfs licht beter waren dan, standaard SFT (bijv. +55.23% verbetering voor Mistral-7B vs +54.19% voor SFT).
• Vermindering van Vergeten: De methode reduceerde de daling in prestaties op algemene benchmarks (MMLU, AGIEval) aanzienlijk.
  • Voor Mistral-7B: De daling in algemene vaardigheden was slechts 7.35% bij de Rewriting-agent, vergeleken met 19.17% bij standaard SFT.
  • Gemiddeld werd het vergeten op niet-downstream benchmarks met 12.34% gereduceerd.
• Trainingstabiliteit: De training van het doelmodel op de herschreven data startte met een lagere loss en convergeerde sneller en stabieler dan bij andere methoden, wat wijst op een betere distributie-uitlijning.
• Ablatie Studies:
  • Het verwijderen van de distributie-uitlijning ($r_{dist}$) leidde tot de grootste daling in algemene vaardigheden, wat bevestigt dat uitlijning met de QA-distributie cruciaal is.
  • Het gebruik van een "harde poort" (alleen belonen voor correcte antwoorden) bleek superieur aan "soft shaping" (belonen voor incorrecte antwoorden), omdat dit ruis in het leerproces voorkomt.

5. Significantie en Conclusie

Dit paper biedt een fundamentele oplossing voor het stabiliseren van Supervised Fine-Tuning onder omstandigheden van distributieverandering. Door data-rewriting te formuleren als een RL-probleem, kunnen onderzoekers de "gap" tussen expert-data en de natuurlijke generatie van het model dynamisch dichten.

De belangrijkste implicatie is dat het mogelijk is om modellen effectief te specialiseren voor specifieke taken zonder hun algemene intelligentie op te offeren. De methode biedt een schaalbare, data-centric alternatief voor complexe optimisatiemethoden (zoals PPO met conservatieve updates) en benadrukt dat de kwaliteit en distributie van de trainingsdata net zo belangrijk zijn als de trainingsalgoritme zelf. De code is openbaar beschikbaar, wat de reproduceerbaarheid en adoptie in de gemeenschap faciliteert.