Technische Samenvatting: Reinforcement Learning met Rijke Feedback via Distributional DAgger

Probleemformulering

Het artikel adresseert de beperkingen van het dominante paradigma in de training van redeneermodellen: Reinforcement Learning from Verifiable Rewards (RLVR). Bij RLVR genereren modellen meerdere kandidaat-oplossingen, en aan elke oplossing wordt een enkele binaire beloning (correct/incorrect) toegekend op basis van het uiteindelijke antwoord. Deze aanpak lijdt onder twee kritieke gebreken:

1. Slechte Credit Assignment: De terminale beloning wordt uniform over alle tokens in een respons verspreid, waardoor het niet mogelijk is om te onderscheiden welke specifieke redeneerstappen hebben bijgedragen aan succes of falen.
2. Onbruikbaarheid voor Onverifieerbare Domeinen: RLVR vertrouwt op automatische verificatie van uiteindelijke antwoorden, wat het ongeschikt maakt voor taken waarbij grondwaarheid-antwoorden niet beschikbaar zijn of moeilijk te controleren zijn.

Hoewel veel settings "rijke feedback" bieden (bijv. executietrajecten, foutlogs, expertcorrecties of zelf-evaluaties), kampen huidige methoden die proberen dit te benutten via on-policy self-distillation (bijv. SDPO, OPSD) met theoretische en praktische hindernissen. Bestaande benaderingen minimaliseren doorgaans $f$-divergenties (zoals reverse KL of Jensen-Shannon) tussen een student-beleid en een feedback-geconditioneerde docent. De auteurs stellen dat deze methoden geen garanties bieden voor monotone beleidsverbetering en falen in het uitvoeren van sequentie-niveau credit assignment, wat vaak leidt tot suboptimale convergentie.

Methodologie: DistIL

De auteurs stellen DistIL (Distributional Imitation Learning) voor, een distributionele variant van het klassieke DAgger-algoritme, aangepast voor reinforcement learning met rijke feedback.

Kerndoelstelling

DistIL optimaliseert een forward cross-entropy verlies tussen het student-beleid ($\pi_\theta$) en een feedback-geconditioneerd docent-beleid ($\pi_T$) op staten die door de student worden bezocht:
$$\mathcal{L}_{\text{DistIL}}(\theta) := \mathbb{E}_{x \sim \rho, y \sim \pi_\theta(\cdot|x)} \left[ \sum_{t=1}^H H_\times(\pi_T(\cdot | s_t), \pi_\theta(\cdot | s_t)) \right]$$
waarbij $H_\times(p, q) = -\mathbb{E}_{y \sim p}[\log q(y)]$ de forward cross-entropy is. Cruciaal is dat de docent $\pi_T$ geconditioneerd is op rijke feedback $f$ (bijv. executielogs of grondwaarheid), maar vast wordt gehouden tijdens de student-update.

Gradiënt en Credit Assignment

In tegenstelling tot eerdere self-distillatie methoden die gradiënten benaderen met behulp van enkel lokale token-gewijze discrepanties (wat effectief de staat-distributie bevriest), gebruikt DistIL volledige gradiënten. De gradiënt ontbindt zich in twee termen:

1. Lokale Imitatie: Een door de docent gewogen term die de student bijwerkt om overeen te komen met de distributie van de docent in de huidige staat.
2. Toekomstbewuste Credit Assignment: Een term die de cross-entropy discrepantie waargenomen op toekomstige tijdstappen terugvoert naar eerdere beslissingen. Dit zorgt ervoor dat vroege token-keuzes die het traject naar een staat met een hoge docent-student discrepantie beïnvloeden, de passende credit of straf ontvangen.

Het algoritme ondersteunt black-box docenten (externe modellen of menselijke experts), aangezien het geen toegang vereist tot docent log-waarschijnlijkheden, maar enkel samples of de mogelijkheid om de doelstelling te schatten.

Belangrijkste Bijdragen

1. Theoretische Beperkingen van Bestaande Methoden

Het artikel biedt formele bewijzen die aantonen dat bestaande on-policy self-distillatie doelstellingen gebaseerd op $f$-divergenties (inclusclusief reverse KL en Jensen-Shannon) lijden onder twee fundamentele problemen:

• Geen Garantie voor Monotone Verbetering: Zelfs als de docent een strikt hogere verwachte beloning heeft dan de student, garandeert het minimaliseren van een $f$-divergentie niet dat een natuurlijke gradiëntstap de verwachte beloning van de student zal verhogen. De update-richting kan onbedoeld de waarschijnlijkheid van suboptimale acties verhogen door de weging door log-ratio's.
• Suboptimale Convergentie via Lokale Approximatie: Door de "toekomstige credit"-term (de impact van huidige acties op toekomstige staat-distributies) te negeren, kunnen lokale gradiënt-approximaties convergeren naar strikt suboptimale beleidsregels. De auteurs bewijzen dat in eenvoudige twee-staps MDP's, lokale updates kunnen falen in het toekennen van credit aan vroege beslissingen die cruciaal zijn voor het bereiken van hoog-belonende staten.

2. Theoretische Garanties van DistIL

DistIL bezit eigenschappen die optimalisatie afstemmen op beloningsmaximalisatie:

• Monotone Beleidsverbetering: Onder standaard aannames van lokale realiseerbaarheid garandeert de forward cross-entropy doelstelling dat het bewegen naar een betere docent de verwachte beloning van de student verhoogt.
• Regret Bounds: Het algoritme heeft een sublineaire regret bound, waarbij de convergentiesnelheden afhankelijk zijn van de concentratietoet coefficients en de variantie van de docent.
• Pass@N Optimalisatie: Er is bewezen dat de doelstelling een ondergrens maximaliseert van de docent-gewogen waarschijnlijkheid van succes. Dit biedt een principiële verklaring voor waarom DistIL de Pass@N metrieken (de waarschijnlijkheid van het genereren van minstens één correct antwoord in $N$ samples) verbetert over alle $N$.

3. Empirische Prestaties

DistIL werd geëvalueerd over drie domeinen:

• Wetenschappelijk Redeneren: Op SciKnowEval (Chemie, Natuurkunde, Biologie, Materiaalkunde) presteerde DistIL beter dan SDPO, GRPO (on-policy en off-policy), en behaalde hogere Best@16 en Maj@16 scores. Opvallend genoeg vertoonde DistIL een stabiele training, terwijl SDPO significante oscillaties en prestatie-instortingen vertoonde in bepaalde domeinen.
• Coderen: Op LiveCodeBench, gebruikmakend van executietrajecten als feedback, presteerde DistIL beter dan SDPO en GRPO in zowel exacte correctheid als partiële accuratesse.
• Moeilijk Wiskundig Redeneren: Op uitdagende problemen (AIME, HMMT) waar het basismodel een zero pass rate heeft, faalden RLVR-methoden (GRPO) door een gebrek aan advantage signalen. DistIL slaagde erin grondwaarheid-oplossingen te benutten om de prestaties te verbeteren, waarmee het self-distillatie baselines (SDPO, OPSD) en SFT overtrof.

Betekenis en Claims

Het artikel claimt dat DistIL de kloof tussen imitatie learning en reinforcement learning overbrugt door een theoretisch gefundeerd kader te bieden voor het leren van rijke feedback. De significantie ligt in:

1. Correctie van de Doelstelling: Het stapt af van $f$-divergentie minimalisatie (wat reward-misaligned kan zijn) naar forward cross-entropy, wat monotone verbetering garandeert wanneer de docent beter is.
2. Mogelijk maken van Sequentie-niveau Leren: Door volledige gradiënten te gebruiken, lost DistIL het credit assignment probleem op dat inherent is aan langdurige redenering, waardoor vroege beslissingen worden geëvalueerd op basis van hun downstream consequenties.
3. Praktische Robuustheid: De methode werkt met black-box docenten en diverse feedbacktypes (van binaire correctheid tot executielogs), wat het toepasbaar maakt voor een breder scala aan redeneertaken dan standaard RLVR.

De auteurs concluderen dat de combinatie van een reward-aligned doelstelling en een toekomstbewuste credit assignment essentieel is voor het effectief distilleren van rijke feedback in redeneermodellen, wat leidt tot consistente empirische winsten ten opzichte van state-of-the-art baselines.