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

Das Paper stellt DyJR vor, einen effizienten Regularisierungsrahmen für das Reinforcement Learning, der durch einen dynamischen Puffer und Jensen-Shannon-Divergenz die Vielfalt der Trainingsdaten erhält und so die Überanpassung sowie den Modus-Kollaps bei der Verbesserung von Large Language Models verhindert.

Das Wesentliche

Stell dir vor, du lernst ein neues, schwieriges Fach, wie Mathematik oder das Schreiben von komplexen Datenbankabfragen (SQL). Du hast einen sehr klugen Tutor (das KI-Modell), der versucht, die besten Lösungen zu finden.

Das Problem bei den aktuellen Methoden (wie GRPO) ist folgendes: Der Tutor probiert viele Lösungen aus, behält aber nur die eine perfekte Lösung, die er gerade gefunden hat, und wirft den Rest weg. Wenn er dann wiederholt, was er gelernt hat, konzentriert er sich so stark auf diesen einen Weg, dass er vergisst, wie man auch andere, vielleicht kreative Wege geht. Er wird stur und findet keine neuen Lösungen mehr. Das nennt man "Mode Collapse" (ein Zusammenbruch der Vielfalt).

Andere Methoden versuchen, alte Lösungen in einer riesigen Bibliothek zu speichern und sie immer wieder zu lesen. Das kostet aber extrem viel Platz (Gedächtnis) und Zeit, und oft führt es dazu, dass der Tutor sich nur noch auf die alten, perfekten Lösungen versteift, anstatt zu lernen, flexibel zu bleiben.

Die Lösung: DyJR (Dynamischer Jensen-Shannon Replay)

Die Autoren dieses Papers haben eine neue Methode namens DyJR entwickelt. Man kann sich das wie einen weisen, dynamischen Coach vorstellen, der zwei geniale Tricks anwendet:

1. Der "Frische-Salat"-Ansatz (Dynamischer Puffer)

Stell dir vor, der Coach hat einen Kühlschrank für alte Lösungen.

• Die alten Methoden: Sie füllen den Kühlschrank mit jeder Lösung, die je gefunden wurde, egal wie alt sie ist. Der Kühlschrank wird riesig, voller alter, verrotteter Lebensmittel, und es ist schwer, das Frische zu finden.
• DyJR: Dieser Coach ist sehr wählerisch. Er sagt: "Ich behalte nur die Lösungen, die ich in den letzten paar Tagen gefunden habe." Er wirft alles Alte sofort raus (FIFO-Prinzip: First In, First Out).
• Warum? Weil das Gehirn des KIs sich schnell verändert. Was vor 100 Schritten noch eine geniale Lösung war, ist heute vielleicht veraltet. DyJR konzentriert sich nur auf das, was gerade frisch und relevant ist. Besonders am Anfang, wenn das Modell noch viel experimentiert, füllt er den Kühlschrank großzügig, um viele verschiedene Ideen zu speichern. Sobald das Modell stabiler wird, macht er den Kühlschrank kleiner, um Platz zu sparen.

2. Der "Vielfalts-Gürtel" (Jensen-Shannon Regularisierung)

Stell dir vor, der Coach möchte, dass der Schüler nicht nur eine Lösung perfekt beherrscht, sondern viele verschiedene Wege kennt, um ans Ziel zu kommen.

• Die alten Methoden: Sie sagen zum Schüler: "Mach es genau so wie in diesem alten Buch!" (Direktes Kopieren). Das führt dazu, dass der Schüler stur wird.
• DyJR: Der Coach sagt: "Schau dir die Vielfalt der Lösungen an, die wir gerade gespeichert haben. Versuche nicht, eine davon perfekt zu kopieren, sondern bleibe in der Nähe dieser ganzen bunten Mischung."
• Die Metapher: Statt den Schüler zu zwingen, genau wie ein einzelner Meister zu sein, gibt DyJR ihm einen "Vielfalts-Gürtel". Solange der Schüler innerhalb dieses Gürtels bleibt (also eine gewisse Vielfalt an Lösungen behält), darf er weiter lernen. Er wird nicht bestraft, wenn er einen neuen Weg versucht, solange er nicht völlig von den erfolgreichen Mustern abweicht. Dies verhindert, dass der Schüler in einer einzigen Denkweise gefangen bleibt.

Das Ergebnis

Durch diese zwei Tricks passiert Magie:

• Effizienz: Der Coach braucht keinen riesigen, teuren Kühlschrank mehr. Er spart Speicherplatz und Rechenzeit.
• Bessere Ergebnisse: Das Modell wird nicht starr. Es behält seine Kreativität und findet auf schwierigen Prüfungen (wie Mathematik-Wettbewerben oder SQL-Abfragen) deutlich bessere Lösungen als die alten Methoden.
• Stabilität: Das Modell lernt schneller, ohne in Sackgassen zu laufen.

Zusammenfassend:
DyJR lehrt die KI nicht einfach nur, "die eine richtige Antwort" auswendig zu lernen. Es lehrt sie, vielfältig zu denken. Es behält nur die frischesten, besten Ideen im Gedächtnis und sorgt dafür, dass die KI nicht vergisst, wie man verschiedene Wege geht. So wird sie nicht nur schlauer, sondern auch flexibler und robuster.

Technische Zusammenfassung

1. Problemstellung

Das Paper adressiert ein zentrales Problem beim Training von Large Language Models (LLMs) für logisches Schließen (Reasoning) mittels Reinforcement Learning mit verifizierbaren Belohnungen (RLVR).

• Ineffizienz von On-Policy-Algorithmen: Aktuelle State-of-the-Art-Methoden wie GRPO (Group Relative Policy Optimization) sind „On-Policy". Das bedeutet, dass Rollout-Daten (generierte Antwortpfade) nach einer einzigen Aktualisierung verworfen werden. Dies führt zu einer extremen Ineffizienz bei der Datennutzung und einem hohen Ressourcenverbrauch.
• Fehler bestehender Replay-Ansätze: Versuche, Experience Replay (Wiederverwendung historischer Daten) einzuführen, scheitern oft an zwei Hauptproblemen:
  1. Modus-Kollaps (Mode Collapse): Durch das direkte Maximieren der Likelihood historischer, korrekter Pfade neigen Modelle dazu, sich auf spezifische Lösungspfade zu spezialisieren und ihre Explorationsfähigkeit zu verlieren (Overfitting).
  2. Hoher Speicherbedarf: Traditionelle Methoden speichern oft riesige Mengen an historischen Daten, was den GPU-Speicherbedarf massiv erhöht.
• Kernthese der Autoren: Historische Daten sollten nicht primär dazu dienen, die Genauigkeit zu verstärken, sondern die Vielfalt (Diversity) der Lösungspfade aufrechtzuerhalten. Die wertvollsten Daten für die Diversität stammen oft aus der frühen Trainingsphase, in der die Entropie noch hoch ist.

2. Methodik: DyJR (Dynamic Jensen-Shannon Replay)

Die Autoren schlagen DyJR als einen Regularisierungsrahmen vor, der zwei Hauptinnovationen kombiniert, um die Diversität zu erhalten, ohne die Recheneffizienz zu opfern.

A. Dynamischer Puffer mit Zeit-Sensitivität (Time-Sensitive Dynamic Buffer)

Anstatt alle historischen Daten zu speichern, nutzt DyJR eine nicht-uniforme Strategie:

• FIFO-Prinzip (First-In-First-Out): Der Puffer behält nur zeitlich nahe Proben bei. Daten, die älter als ein bestimmter Schwellenwert ($M$, „Max Age") sind, werden verworfen. Dies synchronisiert den Puffer mit der aktuellen Entwicklung des Modells und verhindert das Lernen aus veralteten Verteilungen.
• Adaptive Größe: Die Kapazität des Puffers wird dynamisch angepasst. Während der frühen „Warm-up"-Phase (hohe Entropie, schnelle Transition) wird der Puffer vergrößert, um hoch-entropische Muster zu erfassen. Sobald das Modell stabilisiert ist, wird die Größe reduziert, um Speicher zu sparen.
• Bias-Awaree Auswahl: Es wird eine Strategie zur Auswahl von Daten basierend auf der „Konfidenz" ($C_{id}$) verwendet. Für einfache Aufgaben werden hoch-konfidente Lösungen bevorzugt; für schwierige Aufgaben werden auch seltenere Lösungen zugelassen, um Datenmangel zu vermeiden.

B. Jensen-Shannon-Divergenz-Regularisierung (JS-Divergence Regularization)

Statt historische Daten direkt für Gradienten-Updates zu nutzen (was zu Overfitting führt), wird die Jensen-Shannon-Divergenz (JS) als Regularisierungsterm eingeführt:

• Verteilungs-Anker: Der Puffer repräsentiert eine Mischung aus historischen erfolgreichen Strategien.
• Regularisierungsziel: Das aktuelle Modell $\pi_\theta$ wird so optimiert, dass es die JS-Divergenz zu dieser gemischten Verteilung minimiert.
• Vorteil gegenüber KL-Divergenz: Im Gegensatz zur Forward-KL-Divergenz, die dazu neigt, alle Modi abzudecken und zu glätten (was zu instabilem Training führen kann), ist die JS-Divergenz symmetrisch und beschränkt. Sie wirkt als flexibler Anker, der das Modell daran hindert, zu weit von diversen erfolgreichen Pfaden abzuweichen, ohne die Optimierungsrichtung aggressiv zu verändern.
• Effiziente Berechnung: Da die direkte Berechnung der Mischungsverteilung bei autoregressiven Modellen unmöglich ist, wird ein varianzarmes generatives Schätzverfahren verwendet, das auf den gespeicherten Log-Wahrscheinlichkeiten basiert.

Die Gesamt-Loss-Funktion lautet:
$$\mathcal{L}_{total}(\theta) = \mathcal{L}_{GRPO}(\theta) + \alpha_{JS} \cdot \mathcal{L}_{JS}(\theta)$$

3. Hauptbeiträge

1. Paradigmenwechsel: Die Autoren definieren Experience Replay neu: Das Ziel ist nicht die Genauigkeitsoptimierung, sondern die Diversitäts-Regularisierung.
2. Dynamische Datenkonstruktion: Einführung einer zeitnahen Pufferstrategie, die den Speicherbedarf drastisch senkt, indem sie nur relevante, aktuelle Daten speichert und die Größe an die Trainingsphase anpasst.
3. Umfassende Evaluation und Analyse:
   • Nachweis der Überlegenheit auf mathematischen Benchmarks (z. B. AIME, AMC, HMMT) und Text-to-SQL-Aufgaben (BIRD, Spider).
   • Detaillierte Analyse der Token-Wahrscheinlichkeitsentwicklung (Rank-k), die zeigt, wie DyJR den Kollaps auf den Rang-1-Token verhindert.

4. Ergebnisse

Die Experimente wurden auf Modellen wie Qwen3-4B und Llama-3.1-8B durchgeführt.

• Leistungssteigerung: DyJR übertrifft den GRPO-Baseline deutlich. Auf mathematischen Benchmarks erreichte DyJR eine durchschnittliche Genauigkeit von 34,1 % (vs. 29,8 % bei GRPO), was einer Verbesserung von +4,3 % entspricht.
• Vergleich mit SOTA: DyJR schlägt andere Replay-Methoden wie RLEP und Ex-GRPO sowie statische JS-Ansätze wie DPH-RL.
  • RLEP schnitt schlecht ab, da es zu viel Speicher benötigt und zu starre Updates vornimmt.
  • DPH-RL (statisch) war weniger effektiv als DyJR, da es nicht auf die sich ändernde Komplexität der Daten reagiert.
• Effizienz: DyJR benötigt nur einen Bruchteil des Speichers (ca. 2k Paare vs. 28k bei RLEP) und hat eine Rechenzeit, die mit dem Original-GRPO vergleichbar ist.
• Robustheit: Die Methode funktioniert sowohl bei einfachen als auch bei extrem schwierigen Aufgaben konsistent gut.

5. Analyse der Trainingsdynamik

Eine tiefgehende Analyse der Token-Wahrscheinlichkeiten (Rank-k) liefert wichtige Erkenntnisse:

• GRPO: Zeigt einen schnellen Entropie-Verlust. Die Wahrscheinlichkeit für den Rang-1-Token steigt schnell auf >90 % und bleibt dort, was auf einen Modus-Kollaps hindeutet.
• DyJR: Hält die Entropie über die gesamte Trainingsdauer hinweg höher. Die Wahrscheinlichkeit für Rang-1-Tokens sinkt kontrolliert, während die Wahrscheinlichkeiten für Rang-2 und Rang-3 erhalten bleiben. Dies zeigt, dass das Modell weiterhin alternative Lösungspfade erkundet, anstatt sich auf einen einzigen Pfad festzulegen.
• JS vs. Forward KL: Die JS-Regularisierung führt zu einer stabileren und diverseren Verteilung als die Forward-KL-Regularisierung, die in diesem Kontext zu stark glättet.

6. Bedeutung und Fazit

DyJR stellt einen bedeutenden Fortschritt im Bereich des Reinforcement Learning für LLMs dar. Es löst das Dilemma zwischen der Effizienz von Experience Replay und dem Risiko des Modus-Kollaps.

• Skalierbarkeit: Durch den geringen Speicherbedarf und die hohe Recheneffizienz ist die Methode für das Training großer Modelle auf komplexen Reasoning-Aufgaben skalierbar.
• Theoretischer Einsicht: Das Paper zeigt, dass der größte Wert historischer Daten in den frühen, hoch-entropischen Explorationsphasen liegt, nicht in den späteren, hoch-akkuraten Trajektorien.
• Praxisrelevanz: Die Methode ermöglicht es, Modelle robuster und vielseitiger zu machen, ohne die Trainingskosten signifikant zu erhöhen, was sie zu einem starken Kandidaten für zukünftige Reasoning-Systeme macht.