DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

Het artikel introduceert DyJR, een efficiënt framework voor versterkt leren dat de diversiteit van Large Language Models behoudt door een dynamische buffer en Jensen-Shannon-divergentie-regulering te gebruiken, waardoor het de prestaties van bestaande methoden zoals GRPO verbetert zonder in te leveren op trainingsefficiëntie.

De kern

Stel je voor dat een kunstmatige intelligentie (AI) een slimme student is die wiskundepuzzels moet oplossen. Om dit te leren, gebruikt de AI een methode genaamd "Reinforcement Learning" (Versterkend Leren). Dit werkt als volgt: de AI probeert een oplossing, krijgt een punt als het goed is, en past zich aan.

Het probleem met de huidige methoden is dat de AI vaak vergeten wat hij eerder heeft geleerd. Zodra hij een nieuwe poging doet, gooit hij de oude, succesvolle pogingen in de prullenbak. Het is alsof je elke dag een nieuwe wiskundetoets doet, maar de antwoorden van gisteren direct verbrandt. Je leert daardoor niet van je eerdere successen en verspilt veel tijd en energie.

Andere methoden proberen dit op te lossen door een "herinneringskast" (een buffer) te maken waar ze oude, goede antwoorden in opslaan. Maar deze methoden hebben twee grote nadelen:

1. Ze worden erg traag en zwaar voor de computer (te veel geheugen nodig).
2. Ze zorgen ervoor dat de AI te vastloopt. De AI kijkt alleen naar de één beste oplossing die hij ooit heeft gevonden en probeert die steeds opnieuw te kopiëren. Hierdoor stopt hij met zoeken naar andere, misschien nog betere manieren. Hij wordt een "kloon" van zichzelf en verliest zijn creativiteit.

De Oplossing: DyJR (De Slimme Herinneringsassistent)

De auteurs van dit papier hebben een nieuwe methode bedacht, genaamd DyJR. Ze zien het probleem anders: het doel van het opslaan van oude antwoorden is niet om de AI te dwingen precies hetzelfde te doen, maar om hem te helpen diversiteit te behouden. Ze willen voorkomen dat de AI in één spoor vastloopt.

Hier is hoe DyJR werkt, vertaald naar alledaagse analogieën:

1. De "Vluchtige Herinneringskast" (Dynamic Buffer)

Stel je voor dat je een dagboek bijhoudt.

• Oude methoden: Ze proberen elk stukje papier uit je hele leven op te slaan, van je geboorte tot nu. Dit wordt een enorme, rommelige stapel waar je geen raad mee weet.
• DyJR: DyJR is slimmer. Het zegt: "Ik bewaar alleen de notities van de afgelopen week."
  • Waarom? Omdat de AI in het begin van het leren heel creatief is en veel verschillende oplossingen probeert (hoge "entropie"). Later wordt hij steeds zekerder en kiest hij steeds vaker voor dezelfde oplossing.
  • DyJR houdt de kast dus dynamisch: hij is groot in het begin om die creatieve chaos vast te leggen, maar krimpt later. Hij gooit oude, verouderde notities weg (FIFO-principe: First In, First Out) zodat de AI alleen kijkt naar wat recent en relevant is. Dit bespaart enorm veel ruimte.

2. De "Veilige Rem" (Jensen-Shannon Regularisatie)

Dit is het meest creatieve deel.

• Oude methoden: Ze zeggen tegen de AI: "Kijk naar die oude goede oplossing en doe precies hetzelfde!" Dit dwingt de AI om die ene oplossing te kopiëren.
• DyJR: DyJR gebruikt een andere aanpak. Het zegt: "Kijk naar die oude goede oplossingen, maar zorg dat je niet te ver afwijkt van het gemiddelde van al die verschillende manieren."
  • Ze gebruiken een wiskundige maatstaf (Jensen-Shannon Divergentie) die fungeert als een veilige rem of een touw.
  • Het laat de AI vrij om nieuwe, creatieve oplossingen te vinden (zoals een kind dat speelt), maar het touw zorgt ervoor dat hij niet volledig de weg kwijtraakt of in één hoekje van de kamer blijft hangen. Het houdt de AI "in balans" tussen nieuw leren en het behouden van zijn brede denkvermogen.

Wat levert dit op?

In de proeven (met wiskundepuzzels en het schrijven van database-queries) bleek DyJR wonderen te doen:

• Beter presteren: De AI werd slimmer dan de standaardmethodes. Hij loste meer puzzels op, zelfs de heel moeilijke.
• Sneller en lichter: Omdat hij niet alles opslaat, heeft hij minder computerkracht nodig.
• Minder vastlopen: De AI bleef gedurende de hele training creatief. In plaats van na een paar dagen te stoppen met zoeken (omdat hij dacht dat hij het al wist), bleef hij verschillende wegen verkennen, wat uiteindelijk leidde tot betere resultaten.

Samenvattend in één zin

DyJR is als een slimme coach die een student niet dwingt om één oplossing te memoriseren, maar die een selectie van recente, diverse oplossingen bijhoudt om de student te helpen creatief te blijven en niet vast te lopen in één denkpatroon, allemaal zonder de computer te laten oververhitten.

Technische samenvatting

Probleemstelling

Reinforcement Learning (RL) met verifieerbare beloningen (RLVR) heeft de redeneercapaciteiten van Large Language Models (LLM's) aanzienlijk verbeterd, met name via algoritmen zoals GRPO (Group Relative Policy Optimization). Echter, deze on-policy algoritmen hebben een fundamenteel nadeel: ze zijn steekproef-inefficiënt. Omdat ze alleen gebruikmaken van de huidige policy, worden kostbare roll-out data (gesimuleerde antwoorden) na één update verworpen.

Bestaande oplossingen die "Experience Replay" toepassen (zoals RLEP en Ex-GRPO) proberen dit op te lossen door historische, accurate samples opnieuw te gebruiken voor directe policy-updates. Het paper identificeert twee kritieke tekortkomingen in deze aanpak:

1. Mode Collapse (Overfitting): Door historische data direct te maximaliseren, dwingt het model zich te focussen op specifieke oplossingspaden, wat leidt tot een verlies van exploratievermogen en een snelle daling van de entropie.
2. Resource-inefficiëntie: Traditionele replay-mechanismen vereisen enorme buffers om volledige trajecten op te slaan, wat leidt tot hoge GPU-geheugenniveaus.

De auteurs betogen dat historische data in RLVR niet primair moet dienen om nauwkeurigheid te versterken, maar als een regularisatiemechanisme om diversiteit te behouden.

Methodologie: DyJR

De auteurs stellen DyJR (Dynamic Jensen-Shannon Replay) voor, een framework dat twee innovaties combineert om diversiteit te behouden zonder de rekenkosten te verhogen.

1. Dynamische Data Constructie (Time-Sensitive Dynamic Buffer)

In plaats van brute-force opslag van alle historische data, introduceert DyJR een slimme bufferstrategie:

• FIFO en Max Age (M): De buffer gebruikt een "First-In-First-Out" protocol en behoudt alleen samples die binnen een bepaalde tijdspanne (bijv. de laatste 8 stappen) zijn gegenereerd. Dit zorgt voor data die temporair dicht bij de huidige modelstate ligt.
• Adaptieve Grootte: De buffer groeit tijdens de vroege, volatiele trainingsfase (waar de entropie snel daalt) om hoog-entropische redeneringspatronen vast te leggen, en krimpt later om geheugen te besparen.
• Bias-Aware Selectie: Samples worden geselecteerd op basis van "confidence" (het aantal correcte antwoorden in een groep). De buffer vult zich eerst met zeer betrouwbare samples en opent zich later voor moeilijkere taken om "catastrophic forgetting" te voorkomen.

2. Jensen-Shannon Divergentie Regularisatie

In plaats van de historische samples direct te gebruiken voor gradiëntupdates (wat de policy naar die specifieke samples duwt), gebruikt DyJR de Jensen-Shannon (JS) Divergentie als een regularisatie-term.

• De buffer fungeert als een dynamische referentieverdeling ($Q_B$).
• De loss-functie straft af hoe ver de huidige policy ($\pi_\theta$) afwijkt van deze gemengde verdeling van succesvolle historische paden.
• Waarom JS? In tegenstelling tot de Forward KL-divergentie (die neigt om alle modes te "dekken" en kan leiden tot over-gladde distributies) is JS symmetrisch en begrensd. Dit zorgt voor een robuustere regularisatie die diversiteit behoudt zonder de optimalisatierichting agressief te verstoren.

De totale loss-functie is:
$$L_{total}(\theta) = L_{GRPO}(\theta) + \alpha_{JS} \cdot L_{JS}(\theta)$$

Belangrijkste Bijdragen

1. Paradigmaverschuiving: Het redefineert de rol van Experience Replay van "nauwkeurigheids-optimalisatie" naar "diversiteits-regularisatie".
2. Temporeel Nabije Buffer: Een niet-uniforme bufferstrategie die alleen tijdelijk nabije, hoog-entropische samples bewaart, wat leidt tot optimale prestaties met minimale geheugenvraag.
3. JS als Regularisatie: Het vervangen van directe updates door een JS-divergentie constraint, wat effectief mode collapse voorkomt.
4. Uitgebreide Analyse: Gedetailleerde ablatiestudies die laten zien hoe sub-modules de training-dynamiek beïnvloeden, specifiek gemeten via de evolutie van Rank-k token-kansen.

Resultaten

DyJR werd getest op wiskundige redenering (Qwen3-4B) en Text-to-SQL taken (Llama-3-8B).

• Wiskundige Redenering: DyJR behaalde een gemiddelde nauwkeurigheid van 34,1% op zes benchmarks, wat een significante verbetering is ten opzichte van GRPO (29,8%) en andere state-of-the-art methoden zoals RLEP (31,7%) en Ex-GRPO (32,8%).
• Text-to-SQL: Op de BIRD en Spider datasets verbeterde DyJR de Pass@1 nauwkeurigheid met respectievelijk +3,3% en +5,0% ten opzichte van GRPO.
• Efficiëntie: In tegenstelling tot RLEP, dat enorme buffers nodig heeft, vereist DyJR theoretisch slechts een fractie van de opslag (ongeveer 1.638 sequenties in de experimenten) met een verwaarloosbare GPU-geheugenoverhead (<1GB). De trainingstijd blijft vergelijkbaar met de originele GRPO.
• Diversiteit: Analyse van de "Rank-1 token probability" toont aan dat GRPO snel convergeert naar een enkele oplossing (Rank-1 kans >90%), terwijl DyJR de kans over Rank-1, Rank-2 en Rank-3 tokens herverdeelt, wat wijst op behoud van exploratievermogen.

Betekenis en Impact

DyJR biedt een cruciale oplossing voor de schaalbaarheid van RL in complexe redeneertaken. Het paper toont aan dat het bewaren van vroege, diverse exploratiepatronen waardevoller is dan het herhalen van late, nauwkeurige trajecten. Door de focus te verleggen van pure accuraatheid naar diversiteitsbehoud via een lichtgewicht regularisatiemechanisme, maakt DyJR het mogelijk om LLM's robuuster te maken tegen mode collapse, zonder de rekenkosten van traditionele replay-methoden. Dit is een belangrijke stap naar het schalen van "long thinking" en complexe redenering in AI-systemen.