Bootstrapped Mixed Rewards for RL Post-Training: Injecting Canonical Action Order

Die Studie zeigt, dass das Mischen einer skalaren Belohnung für die Einhaltung einer kanonischen Lösungsreihenfolge mit der eigentlichen Aufgabenbelohnung beim Reinforcement-Learning-Post-Training von Transformern auf Zebra-Rätseln die Leistung verbessert, indem es das Modell zu kanonischen Lösungspfaden lenkt, ohne die Architektur oder die überwachungsdaten zu ändern.

Das Wesentliche

Stell dir vor, du möchtest einem sehr intelligenten, aber etwas chaotischen Schüler beibringen, ein komplexes Logikrätsel (wie das berühmte "Zebra-Rätsel") zu lösen. Normalerweise belohnst du ihn nur dafür, wenn er am Ende die richtige Lösung hat. Aber was, wenn du ihm auch eine kleine, unsichtbare Handreichung gibst, die ihm sagt: "Hey, du hast die Lösung zwar richtig, aber du hast die Schritte in einer seltsamen Reihenfolge gemacht. Versuche es doch mal so, wie ein erfahrener Detektiv es tun würde?"

Genau das ist die Idee hinter diesem Forschungsprojekt für die ICLR 2026. Hier ist die Erklärung in einfachen Worten:

1. Das Problem: Der "Richtige Weg" ist oft unsichtbar

Wenn man KI-Modelle trainiert, geben wir ihnen oft nur eine einfache Belohnung: "Richtig = 1 Punkt, Falsch = 0 Punkte". Das ist wie ein Lehrer, der nur das Endergebnis einer Matheaufgabe ansieht und ignoriert, wie der Schüler dorthin gekommen ist.

Das Problem: Ein KI-Modell kann die richtige Antwort finden, indem es zufällig herumrät oder einen sehr umständlichen Weg geht. Es lernt nicht unbedingt, logisch zu denken, sondern nur, das Ziel zu erreichen.

2. Die Lösung: Ein "Geisterkompass"

Die Forscher haben eine clevere Methode entwickelt, um dem Modell beizubringen, nicht nur das Ziel zu sehen, sondern auch den Weg zu verstehen.

• Der chaotische Start: Zuerst trainieren sie das Modell mit Lösungen, die völlig durcheinander gewürfelt sind. Das Modell lernt also, die Antwort zu finden, aber ohne eine feste Reihenfolge.
• Der "Geisterkompass" (RL Post-Training): Danach kommt der eigentliche Trick. Sie nutzen eine Technik namens "Reinforcement Learning" (Bestärkendes Lernen). Hier geben sie dem Modell zwei Arten von Belohnungen:
  1. Die Hauptbelohnung: "Hast du das Rätsel gelöst?" (Ja/Nein).
  2. Die "Ordnungs-Belohnung": "Hast du die Schritte in der Reihenfolge gemacht, die ein echter Detektiv (ein 'Solver') verwenden würde?"

Die Metapher: Stell dir vor, du läufst durch einen Labyrinth.

• Ohne den Kompass läufst du einfach wild herum. Wenn du am Ausgang ankommst, bekommst du einen Keks.
• Mit dem Kompass bekommst du auch einen Keks, wenn du am Ausgang bist. Aber: Du bekommst zusätzliche kleine Bonuspunkte, wenn du dich in die richtige Richtung bewegst, auch wenn du noch nicht am Ziel bist.

3. Der Trick: Die "Bootstrapped" Mischung

Ein großes Problem bei solchen Belohnungen ist: Die "Lösungs-Belohnung" ist riesig (1 Punkt), während die "Ordnungs-Belohnung" winzig ist (vielleicht 0,01 Punkte). Das Modell würde dann nur auf den großen Keks achten und den Kompass ignorieren.

Die Forscher haben einen cleveren Trick angewendet (die "Bootstrapped Scaling"):
Sie haben das Modell kurz getestet, um zu sehen, wie groß die Belohnungen im Durchschnitt sind. Dann haben sie die Waage so justiert, dass beide Belohnungen am Anfang gleich stark gewichtet sind, genau so, wie sie es wollten (z. B. 99 % für die Lösung, 1 % für die Ordnung).

Das ist wie ein Koch, der zwei Zutaten mischt: Er stellt sicher, dass das Salz (die Ordnung) und der Zucker (die Lösung) im Verhältnis 99:1 gemischt werden, damit das Salz nicht vom Zucker "überschattet" wird, obwohl es eigentlich viel weniger davon gibt.

4. Das Ergebnis: Wenig Ordnung, viel Erfolg

Das Überraschende an den Ergebnissen:

• Das Modell hat niemals gesehen, wie ein Detektiv die Rätsel löst. Es hat nur durcheinander gewürfelte Lösungen gelernt.
• Aber durch den winzigen "Ordnungs-Hinweis" (nur 1 % der Belohnung) hat sich das Modell plötzlich verbessert!
• Die beste Leistung erzielte das Modell, wenn die Belohnung zu 99 % für die Lösung und zu 1 % für die richtige Reihenfolge stand.

Die Moral der Geschichte:
Selbst wenn du einem KI-Modell nur eine winzige, unscharfe Andeutung gibst, wie ein "richtiger" Denkprozess aussieht, kann es lernen, diesen Weg zu bevorzugen. Es muss nicht von Grund auf neu trainiert werden, und man braucht keine neuen, perfekten Daten. Ein kleiner "Kompass" reicht aus, um den chaotischen Lernprozess in eine geordnete, logische Richtung zu lenken.

Zusammenfassung in einem Satz

Die Forscher haben gezeigt, dass man KI-Modelle viel besser machen kann, indem man ihnen nicht nur sagt "Das ist richtig", sondern ihnen auch ganz leise flüstert: "Das ist auch der richtige Weg, auf dem man dorthin kommt." – und das funktioniert sogar, wenn das Modell den Weg vorher noch nie gesehen hat.

Technische Zusammenfassung

Titel: Bootstrapped Mixed Rewards für RL-Post-Training: Einbringen kanonischer Aktionsreihenfolgen

1. Problemstellung

Das Training von Sprachmodellen mit Reinforcement Learning (RL) im Post-Training-Stadium optimiert typischerweise ein einziges skalares Ziel (z. B. die Erfolgswahrscheinlichkeit einer Aufgabe). Dabei wird die strukturelle Information der Umgebung, insbesondere die Reihenfolge, in der Zwischenschritte oder Aktionen ausgeführt werden, weitgehend ignoriert.

Die Autoren untersuchen, ob eine skalare Hinweisgröße (Scalar Hint), die auf eine kanonische Lösungsreihenfolge (Solver-Order) hinweist, die Leistung verbessern kann, selbst wenn das Modell zuvor auf randomisierten Lösungssequenzen feinabgestimmt (Fine-Tuning) wurde. Das Ziel ist es, zu prüfen, ob RL-Post-Training durch solche strukturellen Signale in Richtung einer logisch korrekten Denkweise (Reasoning Trajectory) gelenkt werden kann, ohne die Trainingsdaten oder die Architektur ändern zu müssen. Als Testumgebung dienen Zebra-Rätsel (auch bekannt als Einstein-Rätsel), die als deterministische Umgebungen mit latenten Zustandsübergängen modelliert werden.

2. Methodik

Datensatz und Vorverarbeitung:

• Es wird der Zebra-Rätsel-Datensatz verwendet (basierend auf Shah et al., 2024).
• Jede Lösung besteht aus einer Sequenz von 9 Aktionen (Triplets: Zeile, Spalte, Wert).
• Zwei Varianten der Lösungssequenzen werden definiert:
  1. Solver-Order: Die chronologische Reihenfolge, in der ein deterministischer, menschlicher Solver die Zellen füllt (logisch zwingende Reihenfolge).
  2. Random-Order: Dieselben Triplets, aber zufällig gemischt.

Trainingspipeline:

1. Standard Fine-Tuning: Ein GPT-2-ähnlicher Transformer (4 Schichten, 4 Attention-Heads, Hidden Size 256) wird von Grund auf neu trainiert (nicht von einem vortrainierten Checkpoint) auf dem Zebra-Datensatz mit randomisierten Lösungsreihenfolgen. Das Ziel ist hier nur die korrekte Vorhersage der Inhalte, nicht der Reihenfolge.
2. RL Post-Training (GRPO): Das feinabgestimmte Modell wird mit Group Relative Policy Optimization (GRPO) weiter trainiert.

Reward-Design (Belohnungsfunktionen):
Das Kernstück der Arbeit ist die Kombination zweier Reward-Signale:

• $R_{solve}$ (Sparse Task Reward): Ein binärer Reward (1 oder 0). Er ist nur 1, wenn das Modell die gesamte Lösung korrekt und vollständig liefert. Er ignoriert die Reihenfolge der Vorhersagen.
• $R_{order}$ (Ordering Reward): Ein Reward, der misst, wie gut die vom Modell generierte Reihenfolge der Zellen mit der kanonischen Solver-Reihenfolge übereinstimmt, unabhängig davon, ob die Werte korrekt sind.
  • Formel: $r(r, c) = \frac{1}{1 + |\pi^*(r, c) - \hat{\pi}(r, c)|}$, wobei $\pi^*$ der Index im Solver und $\hat{\pi}$ der Index der ersten Ausgabe des Modells ist.
  • Dies dient als Reward-Shaping, um das Modell in Richtung einer solver-ähnlichen Trajektorie zu lenken.

Bootstrapped Scaling (Skalierung):
Da die Magnituden von $R_{solve}$ (oft 0 oder 1) und $R_{order}$ (ein Durchschnittswert zwischen 0 und 1) unterschiedlich sein können, wird eine Bootstrapped Reward Scaling eingeführt:

• Vor dem GRPO-Training wird das feinabgestimmte Modell auf einem Validierungsset evaluiert.
• Die mittleren Rewards ($\bar{R}_{solve}$ und $\bar{R}_{order}$) werden berechnet.
• Globale Skalierungsfaktoren werden festgelegt, um sicherzustellen, dass die gewichtete Summe $R_{total} = \alpha \cdot R_{solve} + (1-\alpha) \cdot R_{order}$ bei der Initialisierung genau dem gewünschten Mischungsverhältnis $\alpha$ entspricht.
• Dies verhindert, dass ein Reward-Term aufgrund seiner rohen Skalierung den anderen dominiert, und ermöglicht einen sauberen Vergleich verschiedener Mischungen.

3. Wichtige Beiträge

1. Scalar Reward Hinting: Eine Methode, um Solver-Reihenfolgen in GRPO zu injizieren, ohne die Fine-Tuning-Daten zu ändern oder die Architektur anzupassen. Die Information wird ausschließlich als skalare Belohnung während des RL-Phasen bereitgestellt.
2. Bootstrapped Scaling-Verfahren: Ein Verfahren zur Normalisierung heterogener Reward-Magnituden, das kontrollierte Mischungsstudien (Mixture Studies) ermöglicht, indem die relativen Beiträge der Reward-Komponenten bei der Initialisierung stabilisiert werden.
3. Empirische Evidenz: Nachweis, dass grobe Ordnungssignale, wenn sie mit dem Korrektheits-Reward gemischt werden, die Genauigkeit des RL-Post-Trainings signifikant verbessern, selbst wenn das Modell nur auf zufälligen Sequenzen gelernt hat.

4. Ergebnisse

• Baseline: Das auf randomisierten Reihenfolgen feinabgestimmte Modell erreicht eine Puzzle-Accuracy von 0,279.
• Einfluss der Reward-Mischung: Modelle, die mit GRPO nachtrainiert wurden, zeigten konsistent bessere Ergebnisse als das reine Task-Optimierung (nur $R_{solve}$, Verhältnis 1:0).
  • Das beste Ergebnis wurde bei einem Mischungsverhältnis von 0,99 : 0,01 (Solve : Order) erzielt, mit einer Accuracy von 0,363.
  • Selbst ein sehr geringer Anteil an Order-Reward (1%) führte zu einem deutlichen Gewinn gegenüber der reinen Task-Optimierung (0,326).
  • Andere Mischungen (z. B. 0,9 : 0,1 oder 0,75 : 0,25) erreichten ebenfalls signifikante Verbesserungen (ca. 0,355).
• Schlussfolgerung: Die Mischung aus Rewards lenkt das Modell effektiv zu solver-ähnlichen Trajektorien, ohne dass das Modell jemals Solver-Reihenfolgen während des überwachten Trainings gesehen hat. Das Ordnungssignal wirkt als effektiver „leichter Formungsterm" (light shaping term).

5. Bedeutung und Fazit

Die Arbeit demonstriert, dass strukturelle Informationen (wie die logische Reihenfolge von Schritten in einem Rätsel) durch skalare Reward-Signale effektiv in RL-Post-Training integriert werden können. Dies ist besonders relevant, da es eine kostengünstige und modulare Möglichkeit bietet, Bias in Richtung korrekter Denkprozesse einzubringen, ohne neue, aufwendig kuratierte Datensätze erstellen oder Modelle von Grund auf neu trainieren zu müssen.

Einschränkungen und Ausblick:

• Die Studie beschränkt sich derzeit auf eine einzige Aufgabe (Zebra-Rätsel) und eine einzelne Architektur (GPT-2-Style).
• Die verwendeten Skalierungsfaktoren sind statisch (bootstrapped). Da sich die Performance der Reward-Komponenten während des Trainings unterschiedlich schnell entwickeln kann, könnten diese Faktoren im Laufe der Zeit unkalibriert werden.
• Zukünftige Arbeiten sollten untersuchen, ob eine periodische Aktualisierung der Skalierungsfaktoren die Stabilität verbessert und ob sich diese Ergebnisse auf andere Aufgaben und größere Modelle verallgemeinern lassen.

Zusammenfassend zeigt das Paper, dass das Hinzufügen eines „kanonischen Ordnungs-Hinweises" als Reward eine effektive Strategie ist, um die Reasoning-Fähigkeiten von Modellen zu verbessern, die ansonsten nur auf chaotischen Daten trainiert wurden.