Technische Samenvatting: Reflectieve On-Policy Self-Distillatie (ROSD)

Probleemstelling

Hoewel Reinforcement Learning met Verifieerbare Beloningen (RLVR) effectief is gebleken voor het optimaliseren van Large Language Models (LLMs) in verifieerbare domeinen, vertrouwen standaardalgoritmen zoals GRPO op schaarse, uitkomstgerelateerde beloningen. Dit gebrek aan token-niveau supervisie voor tussenliggende redeneerstappen beperkt het vermogen van het model om de kwaliteit van zijn redenering te verfijnen.

On-Policy Self-Distillatie (OPSD) lost dit op door dichte token-niveau supervisie te bieden vanuit een "zelfleraar" die is geconditioneerd op een correcte oplossing. Echter, bestaande OPSD-methode lijden aan twee kritieke beperkingen die zowel stabiliteit binnen het domein als generalisatie buiten het domein (OOD) belemmeren:

1. Imitatie in plaats van Correctie: Het conditioneren van de zelfleraar op een geverifieerde oplossing moedigt de student aan om de specifieke patronen en domeinspecifieke uitdrukkingen van de referentietrajecten na te bootsen, in plaats van te leren om zijn eigen specifieke fouten te corrigeren.
2. Overschrijven van Geldige Redenering: Het toepassen van distillatieverlies op het volledige antwoord dwingt het model om zich te aligneren met de lerardistributie over alle tokens. Dit kan geldige redeneringsvoorvoegsels overschrijven en training-domeinvoorkeuren injecteren in correcte stappen, wat leidt tot overfitting en slechte generalisatie.

Methodologie: ROSD

De auteurs stellen Reflectieve On-Policy Self-Distillatie (ROSD) voor, een raamwerk dat is ontworpen om self-distillatie te verschuiven van imitatie van referentieoplossingen naar gerichte correctie van redenering. ROSD werkt via twee primaire mechanismen:

1. Foutgerichte Zelfreflectie

In plaats van direct een correcte oplossing te gebruiken als context voor de leraar, maakt ROSD gebruik van een zelfreflector (met dezelfde modelgewichten als de student/leraar) om de discrepantie tussen een rollout van de student en een correcte oplossing te analyseren.

• Voor Foutieve Rollouts: De reflector vergelijkt een foutieve rollout ($y^-$) met de kortste correcte rollout ($y^*$) uit dezelfde groep. Het genereert twee outputs:
  • Correctief Idee ($e$): Een uitleg waarom de rollout faalde en hoe de logica kan worden hersteld.
  • Foutcitaat ($q$): Een exacte substring uit de foutieve rollout die het eerste segment markeert waar de redenering fout gaat.
• Voor Correcte Rollouts: De reflector vat de geldigheid van de redenering en het kernidee achter het correcte antwoord samen, waardoor directe imitatie van de volledige oplossing wordt vermeden.

2. Geciteerde Lokalisatie van Self-Distillatie

De zelfleraar is geconditioneerd op het correctief idee ($e$) in plaats van de volledige referentieoplossing. Dit zorgt ervoor dat de leraar fout-specifieke begeleiding biedt.

• Lokalisatie: Het foutcitaat ($q$) wordt teruggekoppeld naar de oorspronkelijke studentrollout om de starttokenindex ($k$) van de fout te identificeren.
• Selectief Verlies: Een masker wordt toegepast op de distillatieverliesfunctie. Het model wordt alleen getraind om de lerardistributie vanaf token $k$ (het foutieve achtervoegsel) te matchen. Het geldige redeneringsvoorvoegsel (tokens $< k$) wordt behouden en uitgesloten van de verliesberekening.
• Doel: Het trainingsdoel minimaliseert de token-niveau Jensen-Shannon-divergentie (JSD) tussen de student en de gemaskerde lerardistributie, zodat updates beperkt blijven tot de segmenten die correctie vereisen.

Belangrijkste Bijdragen

1. ROSD-raamwerk: Een nieuwe reflectieve on-policy self-distillatiemethode die overgaat van het imiteren van referentieoplossingen naar het uitvoeren van selectieve correcties van redenering.
2. Ontwerp met Dubbele Techniek:
   • Een foutgerichte zelfreflector die een correctief idee extraheren en het eerste foutieve segment precies lokaliseert.
   • Foutgelokaliseerde distillatie die supervisie-updates beperkt tot het achtervoegsel dat begint bij de geïdentificeerde fout, waardoor geldige redeneringsvoorvoegsels behouden blijven.
3. Empirische Validatie: Uitgebreide experimenten die aantonen dat ROSD de prestaties van redenering binnen het domein verbetert, terwijl het tegelijkertijd de instorting van OOD-generalisatie die bij standaard OPSD-methoden wordt waargenomen, aanzienlijk vermindert.

Experimentele Resultaten

De auteurs hebben ROSD geëvalueerd op Qwen3-4B en Qwen3-8B backbones over vijf datasets: vier wetenschappelijke QA-benchmarks (Chemie, Fysica, Biologie, Materiaalkunde) en een Tool-Use dataset.

• Prestaties binnen het Domein: ROSD behaalde de hoogste gemiddelde scores voor beide modelgroottes (72,83% op 4B en 73,45% op 8B), en presteerde beter dan zowel GRPO als de standaard SDPO-baseline. Verbeteringen waren consistent over de meeste domeinen, wat aangeeft dat er geen significante afweging is in specifieke gebieden.
• Generalisatie buiten het Domein: ROSD toonde aanzienlijk betere OOD-robustheid in vergelijking met SDPO.
  • Terwijl SDPO ernstige degradatie leed wanneer trainings- en testdomeinen significant verschilden (bijvoorbeeld trainen op ToolUse en testen op Science QA, waarbij SDPO daalde tot 2,88% gemiddeld), behield ROSD sterke prestaties (41,31%).
  • Dit suggereert dat ROSD effectief de ruis en overfitting reduceert die geassocieerd zijn met distillatie van volledige antwoorden.
• Trainingsdynamiek: ROSD convergeerde sneller dan baselines en behield stabiele prestaties gedurende de training. Daarentegen toonde SDPO vaak vroeze winsten gevolgd door een daling in prestaties, wat wijst op overfitting op oplossingsstijlen van het trainingsdomein.
• Ablatiestudies: Het verwijderen van ofwel de reflectiecomponent ofwel het gelokaliseerde distillatiemechanisme resulteerde in prestatiedalingen, wat bevestigt dat zowel foutgerichte begeleiding als selectieve supervisie noodzakelijk en complementair zijn.

Betekenis en Claims

Het artikel stelt dat ROSD het belang benadrukt van foutbewuste, selectief toegepaste supervisie voor effectieve post-training van redenering. Door gebrekkige redeneringsstappen te corrigeren terwijl geldige voorvoegsels behouden blijven, vermijdt ROSD de valkuilen van het imiteren van domeinspecifieke referentietrajecten.

De auteurs positioneren ROSD als een stap richting robuustere redeneringsmodellen die over domeinen kunnen generaliseren. Ze merken echter bescheiden op dat ROSD, hoewel het de OOD-prestaties verbetert ten opzichte van standaard OPSD, in alle instellingen nog niet volledig de OOD-robustheid van GRPO bereikt, wat aangeeft dat verdere vooruitgang nodig is om cross-domein generalisatie volledig te versterken. Het werk benadrukt dat sterke redeneercapaciteiten een strenge veiligheidsalignatie vereisen om potentieel misbruik te voorkomen.