Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

Technische Zusammenfassung: Reformulierung des Reinforcement Learning für LLMs zur effizienten Modellierung unter Black-Box-Diskrepanz

1. Problemstellung

Reinforcement Learning (RL) hat sich zu einem entscheidenden Post-Training-Paradigma entwickelt, um die Denkfähigkeiten von Large Language Models (LLMs) zu verbessern. Das RL-Post-Training leidet jedoch häufig unter unvorhersehbarem suboptimalem Verhalten oder vollständigen Trainingskollapsen. Jüngste Untersuchungen führen diese Ausfälle auf eine allgegenwärtige Trainings-Inferenz-Diskrepanz (oder Diskrepanz) zurück.

Diese Diskrepanz entsteht, weil industrielle Anforderungen Praktiker oft dazu zwingen, unterschiedliche zugrunde liegende Engines und Präzisionsformate (z. B. verschiedene Kernel, Parallelisierungsstrategien oder numerische Präzisionen wie BF16 vs. FP8) für exakt dieselben Modellgewichte während des Trainings und der Inferenz zu verwenden. Dies erzeugt eine verborgene Abweichung zwischen den Aktionswahrscheinlichkeitsverteilungen der Trainings-Policy ($\pi_{Train}$) und der Inferenz-Policy ($\pi_{Inf}$).

Bestehende Minderungsstrategien weisen erhebliche Einschränkungen auf:

• Einheitliche Präzisionsanpassung: Das bloße Herabstufen der Trainingspräzision, um sie an die Inferenz anzupassen (z. B. auf FP16), verkleinert den darstellbaren Wertebereich und verschärft das Risiko von Gradientenexplosionen.
• Heuristische Interventionen: Methoden wie das Maskieren von Gradienten bei Token mit hoher Diskrepanz unterdrücken lediglich die Auswirkungen verrauschter Daten. Sie adressen nicht die Ursache und vernachlässigen das ultimative Ziel: die Optimierung der Policy für den Einsatz in der Inferenz, nicht nur für den Trainingsprozess.
• Infrastruktur-Anpassung: Der Versuch, Low-Level-Operatoren und Hardware-Stacks anzugleichen, ist arbeitsintensiv und oft nicht machbar, da es eine Vielzahl von Randfällen gibt.

Die zentrale Herausforderung besteht darin, die Trainings-Policy in die Lage zu versetzen, ihre Diskrepanz zur Inferenz-Policy dynamisch wahrzunehmen und aktiv zu neutralisieren, ohne das primäre Ziel der Belohnungsmaximierung zu opfern.

2. Methodik: Discrepancy-Constrained MDP (DCMDP)

Die Autoren schlagen einen Paradigmenwechsel von Standard-Markov-Entscheidungsprozessen (MDP) hin zu einem Discrepancy-Constrained MDP (DCMDP) vor. Dieses Framework behandelt die Trainings-Inferenz-Diskrepanz als Black-Box-Optimierungsproblem, wodurch das System heterogene Trainings- und Inferenz-Setups nahtlos handhaben kann.

Kernkomponenten:

• Diskrepanz-Metrik:
  Die Autoren evaluieren empirisch drei Token-Level-Diskrepanzsignale: die Wahrscheinlichkeitsdifferenz ($\delta_{diff}$), das Wahrscheinlichkeitsverhältnis ($\delta_{ratio}$) und die Log-Wahrscheinlichkeitsdifferenz ($\delta_{log}$). Sie identifizieren die absolute Wahrscheinlichkeitsdifferenz ($\delta_{diff} = |p_{Train} - p_{Inf}|$) als das effektivste Straffignal. Sie bietet eine „plastische“ Regularisierung, die driftspezifische Anpassungen für die Bereitstellung korrigiert, ohne die belohnungsgesteuerte Exploration durch zu restriktive Bedingungen (wie bei verhältnisbasierten Strafen, die für Token mit geringer Wahrscheinlichkeit zu konservativ sein können) übermäßig einzuschränken.

• Die Diskrepanz-Toleranzregion:
  Eine entscheidende empirische Erkenntnis ist die Existenz einer Diskrepanz-Toleranzregion. Eine aggressive Bestrafung geringfügiger, vernachlässigbarer Diskrepanzen in einem frühen Stadium des Trainings unterdrückt die Exploration der Policy und reduziert die Lerneffizienz. Umgekehrt verbessert die Reduzierung übermäßiger Diskrepanzen außerhalb dieser Region die Optimierungskonsistenz.

  • Mechanismus: Die DCMDP-Formulierung wendet Strafen nur an, wenn die Diskrepanz einen vordefinierten Schwellenwert $c$ überschreitet. Innerhalb dieser Toleranzregion darf die Policy frei explorieren, um Belohnungen zu maximieren.

• Dual-Objective-Optimierung via Lagrangian-Relaxation:
  Das Problem wird als Maximierung der erwarteten Belohnung unter einer Beschränkung der durchschnittlichen Token-Level-Diskrepanz formuliert:
  $$\max_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum r(o_t) \right]$$
  $$\text{s.t. } \frac{1}{|\tau|} \sum |\pi_{Train}(o_t) - \pi_{Inf}(o_t)| \leq c$$

Um dies zu lösen, verwenden die Autoren Lagrangian-Relaxation. Ein lernbarer Multiplikator $\lambda$ passt das Gewicht der Diskrepanzstrafe dynamisch an:
* Wenn die Diskrepanz die Toleranzgrenze überschreitet ($>c$), erhöht sich $\lambda$, was die Strafe verstärkt, um die Policy zurück in die sichere Zone zu ziehen.
* Wenn die Diskrepanz innerhalb der Grenze liegt, geht $\lambda$ gegen Null, sodass die Policy sich auf die Maximierung der Belohnung konzentrieren kann.

Das finale Ziel (Dual-GRPO) injiziert diese adaptive Strafe direkt in den gruppen-normalisierten Vorteil (group-normalized advantage):
$$\hat{A}^{DCMDP}_{i,t} = \hat{A}_{i,t} - \lambda_k \cdot \delta_{i,t}$$

3. Kernbeiträge

1. Autonome Diskrepanzkompensation: Die Arbeit zeigt, dass die Integration einer robusten Black-Box-Strafe (Token-Level absolute Wahrscheinlichkeitsdifferenz) der Trainings-Policy die Fähigkeit verleiht, das Risiko eines durch Diskrepanz induzierten Trainingskollapses autonom zu eliminieren.
2. Paradigmenwechsel zu DCMDP: Die Autoren formulieren das LLM-RL als Discrepancy-Constrained MDP um. Dieser Ansatz balanciert optimal die dualen Ziele der Maximierung der Aufgabenbelohnung und der Minimierung der Trainings-Inferenz-Diskrepanz, wodurch die naive Skalarisierung von Zielen überwunden wird.
3. Heterogenes Trainingsparadigma: Die Methode ermöglicht ein Setup, in dem Modelle in hochpräzisen, ressourcenintensiven Umgebungen (z. B. BF16-Präzision) optimiert werden, während sie explizit auf kostengünstige, ressourcenbeschränkte Inferenz-Deployments (z. B. FP8) abgestimmt werden. Dies macht die infrastrukturelle Heterogenität von einer Belastung zu einem steuerbaren Optimierungsparameter.

4. Experimentelle Ergebnisse

Die Methode wurde an den Modellen Qwen3-8B (dense) und Qwen3-30B-A3B (Mixture-of-Experts) über sechs mathematische Reasoning-Benchmarks (AIME24/25, AMC23, MATH-500, Minerva, OlympiadBench) evaluiert.

• BF16 Training/Inferenz: DC-GRPO übertraf die klassische GRPO-Baseline konsistent über alle Benchmarks für beide Modelle und erreichte höhere Durchschnittswerte (z. B. verbesserte Qwen3-8B von 67,3 auf 71,0). Dies bestätigt, dass die Durchsetzung von Diskrepanz-Beschränkungen die primäre Belohnungsoptimierung stärkt, statt sie zu behindern.
• Heterogenes Training (BF16) + Deployment (FP8): In diesem Szenario litt die klassische GRPO aufgrund der verstärkten Trainings-Inferenz-Lücke erheblich unter Leistungsabfall oder Kollaps. DC-GRPO konnte dies erfolgreich abmildern und die Performance der reinen BF16-Baseline bewahrte oder übertraf diese sogar.
• Ablationsstudien: Die Methode erwies sich als robust gegenüber Variationen des initialen Lagrangian-Multiplikators ($\lambda_0$) und der dualen Lernrate ($\eta_\lambda$). Es wurde festgestellt, dass das Toleranzbudget ($c$) kritisch ist; Ein zu kleines $c$ führte zu Überregulierung, während ein zu großes $c$ dazu führte, dass die Methode in das klassische GRPO degenerierte.

5. Bedeutung und Ansprüche

Das Paper behauptet, dass DCMDP einen vereinheitlichten, algorithmischen Rahmen zur Milderung der Trainings-Inferenz-Instabilität bietet, ohne dass eine erschöpfende, maßgeschneiderte Optimierung der zugrunde liegenden Hardware- oder Software-Stacks erforderlich ist. Durch die Behandlung der Diskrepanz als Black-Box-Feedbacksignal ermöglicht die Methode:

• Stabiles Training: Verhinderung der häufigen abrupten Kollapse im RL-Post-Training.
• Effiziente Exploration: Erlaubt der Policy, innerhalb einer sicheren Diskrepanzregion frei zu explorieren.
• Ressourceneffizienz: Ermöglicht die Nutzung hochpräziser Trainings-Setups, um sich an beschränkte Inferenz-Umgebungen anzupassen und zu optimieren, was so den effizienteren Real-World-Einsatz von Modellen erleichtert.

Die Autoren räumen Einschränkungen ein und merken an, dass ihr empirischer Umfang derzeit auf BF16- und FP8-Präzisionen sowie spezifische Modellfamilien begrenzt ist, wobei zukünftige Arbeiten notwendig sind, um die Methode auf andere Backend-Kombinationen und Modellarchitekturen auszuweiten.