Technische Zusammenfassung: Robuste Policy-Optimierung zur Verhinderung katastrophalen Vergessens

Problemstellung

Large Language Models (LLMs) werden typischerweise über mehrstufige Pipelines trainiert, beginnend mit Reinforcement Learning from Human Feedback (RLHF), um Modelle mit Sicherheit und Hilfsbereitschaft in Einklang zu bringen, gefolgt von Fine-Tuning für spezifische nachgelagerte Aufgaben (z. B. mathematisches Schlussfolgern, Programmieren). Eine kritische Schwachstelle in diesem Paradigma ist das katastrophale Vergessen: Selbst geringfügige nachgelagerte Aktualisierungen können zuvor erlernte Verhaltensweisen, wie Sicherheitsvorkehrungen, verschlechtern.

Frühere Forschung konzentrierte sich weitgehend auf Minderungsstrategien zum Zeitpunkt der nachgelagerten Aufgabe (z. B. Rehearsal, parameter-effizientes Fine-Tuning, Regularisierung oder Modell-Merging), um Fähigkeiten im Nachhinein zu bewahren. Die Autoren argumentieren jedoch, dass diese Methoden unzureichend sind, da sie von spezifischen nachgelagerten Verfahren abhängen. Wenn ein Modell später einem „gewöhnlichen" Fine-Tuning unterzogen wird, das von diesen spezifischen Protokollen abweicht, bleiben zuvor erlernte Verhaltensweisen brüchig. Die Kernherausforderung besteht darin, sicherzustellen, dass die Basis-Policy selbst gegenüber zukünftiger Anpassung robust ist, unabhängig von der nachgelagerten Aufgabe oder dem verwendeten Algorithmus.

Methodik: Fine-tuning Robust Policy Optimization (FRPO)

Die Autoren schlagen Fine-tuning Robust Policy Optimization (FRPO) vor, ein Framework, das darauf ausgelegt ist, die Basis-Policy für Stabilität gegenüber zukünftigen nachgelagerten Verschiebungen zu optimieren.

1. Problemformulierung: Max-Min-Robustheit

Die Autoren formalisieren das Problem als Max-Min-Optimierung. Anstatt die erwartete Belohnung ausschließlich bei der aktuellen Policy $\pi_\theta$ zu maximieren, zielt das Ziel darauf ab, die minimale Belohnung zu maximieren, die von jeder nachgelagerten Policy $Q$ erreichbar ist, die innerhalb einer durch KL-Divergenz begrenzten Nachbarschaft von $\pi_\theta$ liegt.

Das Optimierungsproblem ist definiert als:
$$\max_{\pi_\theta} \left( \inf_{Q: \mathbb{E}[\text{KL}(Q \parallel \pi_\theta)] \le \rho} \mathbb{E}_{x,y \sim Q}[r(x,y)] - \beta \text{KL}(\pi_\theta \parallel \pi_{\text{ref}}) \right)$$

Hierbei repräsentiert $\rho$ den Radius der KL-Nachbarschaft und spiegelt Standardbeschränkungen beim nachgelagerten Fine-Tuning wider (z. B. Vertrauensbereiche, kleine Lernraten oder LoRA-Beschränkungen). Die Beschränkung wird im Erwartungswert über die Prompt-Verteilung auferlegt, was erhebliche Verschiebungen bei prompt-spezifischen Aufgaben erlaubt, während die durchschnittliche Nähe zur Basis-Policy gewahrt bleibt.

2. Theoretische Herleitung

Unter Verwendung der Dualitätstheorie für distributionell robuste Optimierung (speziell unter einer $f$-Divergenz-Beschränkung) leiten die Autoren eine duale Form des inneren Minimierungsproblems ab.

• Entropisches Risiko-Ziel: Die duale Formulierung zeigt, dass das robuste Ziel äquivalent zur Optimierung eines entropischen Risiko-Ziels ist. Die innere Minimierung über $Q$ transformiert das Problem in die Maximierung von:
  $$-\lambda \log \mathbb{E}_{y \sim \pi_\theta} \left[ e^{-r(x,y)/\lambda} \right]$$
  wobei $\lambda$ ein einstellbarer Hyperparameter ist, der die Risikosensitivität steuert.
• Interpretation: Dieses Ziel bestraft Trajektorien mit niedriger Belohnung stärker als die Standardmaximierung der erwarteten Belohnung. Es sucht effektiv nach „flachen" Regionen in der Belohnungslandschaft innerhalb des Policy-Raums, wodurch sichergestellt wird, dass kleine Störungen (nachgelagertes Fine-Tuning) nicht zu starken Einbrüchen der Leistung führen.
• Grenzfall: Wenn $\lambda \to \infty$, gewinnt das Ziel das Standard-GRPO-Ziel (Group Relative Policy Optimization) zurück, das die erwartete Belohnung maximiert.

3. Algorithmische Implementierung

Die Autoren passen den GRPO-Algorithmus an, um dieses neue Ziel zu optimieren, ohne zusätzlichen Rechenaufwand einzuführen (z. B. keinen gelernten Critic oder zusätzliche Forward-Passes).

• Importance Sampling: Der Erwartungswert unter der neuen Policy wird mithilfe von Stichproben aus einer verzögerten Policy $\pi_{\text{old}}$ via Importance Sampling geschätzt.
• Clipping und Baselines: Um Stabilität zu gewährleisten, setzt der Algorithmus PPO-artiges Clipping auf den Importance Ratios ein. Entscheidend ist, dass die Autoren einen Baseline-Term ableiten, um die hohe Varianz im Gradienten der Log-Partition-Funktion zu kompensieren, wenn $\lambda$ groß ist.
• Bias-Reduktion: Die Monte-Carlo-Schätzung der Log-Partition-Funktion führt zu einem $O(1/G)$-Bias (wobei $G$ die Gruppengröße ist). Die Autoren wenden eine Jackknife-Technik (Leave-One-Out-Schätzung) an, um diesen Bias auf $O(1/G^2)$ zu reduzieren und so eine stabile Konvergenz auch bei kleinen Gruppengrößen sicherzustellen.

Hauptbeiträge

1. Robustes RLHF-Framework: Die Arbeit führt eine Max-Min-Formulierung ein, die alle nachgelagerten Verschiebungen innerhalb einer KL-Kugel berücksichtigt. Sie zeigt, dass diese Formulierung mathematisch äquivalent zur Optimierung eines entropischen Risiko-Ziels ist und liefert eine theoretische Grundlage für Belohnungsflachheit im Policy-Raum.
2. FRPO-Algorithmus: Eine praktische, direkte Ersatzlösung für GRPO, die keine zusätzlichen Berechnungen erfordert. Sie lässt sich nahtlos in bestehende RLHF-Pipelines integrieren, enthält eine abgeleitete Baseline für Stabilität und gewinnt das Standard-GRPO zurück, wenn $\lambda \to \infty$.
3. Empirische Validierung: Die Methode wird in verschiedenen Szenarien evaluiert, einschließlich Sicherheitsausrichtung und mathematischem Schlussfolgern, und zeigt eine überlegene Bewahrung vorheriger Fähigkeiten im Vergleich zu Standard-Fine-Tuning und anderen Robustheits-Baselines.

Experimentelle Ergebnisse

Die Autoren evaluieren FRPO in zwei Hauptszenarien: Sicherheitstraining und Mathematisches Schlussfolgern.

1. Sicherheitstraining (Sicherheitsvorkehrungen)

• Setup: Modelle (Mistral-v0.1 und Qwen2.5) wurden mit FRPO unter Verwendung verschiedener $\lambda$-Werte auf einem Sicherheitsdatensatz trainiert und anschließend auf nachgelagerte Aufgaben feinabgestimmt (Alpaca für Instruktionsbefolgung, GSM8K für Mathematik und UltraFeedback für RL).
• Ergebnisse:
  • Mit FRPO trainierte Modelle ($\lambda \approx 0.2$) behielten nach dem Fine-Tuning im Vergleich zu Standard-GRPO signifikant höhere Sicherheits-Ablehnungsraten und niedrigere Schädlichkeitswerte (StrongREJECT) bei.
  • Die Methode bewahrte Sicherheitsvorkehrungen effektiv, selbst unter aggressiven Fine-Tuning-Einstellungen (vollparametrisches SFT, verschiedene Lernraten, LoRA).
  • Trade-off-Analyse: Während kleinere $\lambda$-Werte die Sicherheit besser bewahrten, verbesserten größere $\lambda$-Werte die Hilfsbereitschaft leicht auf Kosten der Sicherheit. Der optimale $\lambda$-Wert erwies sich als aufgabenabhängig, war aber für Sicherheit im Allgemeinen niedrig.
  • Nachgelagerte Leistung: Entscheidend ist, dass FRPO die Leistung bei nachgelagerten Aufgaben (z. B. GSM8K-Genauigkeit, MMLU-Scores) im Vergleich zu GRPO nicht beeinträchtigte, was bestätigt, dass die Bewahrung der Sicherheit nicht auf einem Versagen beim Erlernen der neuen Aufgabe beruhte.

2. Kontinuierliches Lernen (Mathematik zu Code)

• Setup: Modelle wurden zunächst auf mathematisches Schlussfolgern (MATH-Datensatz) trainiert und anschließend auf Code-Generierung (OpenCodeInstruct) feinabgestimmt.
• Ergebnisse:
  • FRPO mit $\lambda = 2.0$ (optimiert für die höhere Varianz in mathematischen Belohnungen) bewahrte die mathematische Genauigkeit signifikant besser als GRPO.
  • Insbesondere behielt FRPO nach dem Code-Fine-Tuning 22 % höhere Genauigkeit auf dem MATH500-Benchmark bei als GRPO, während vergleichbare Gewinne auf Code-Benchmarks (MBPP+) erzielt wurden.

3. Ablationsstudien

• Vergleich mit SAM: Die Autoren verglichen FRPO mit Sharpness-Aware Minimization (SAM). Während SAM (operierend im Parameterraum) eine gewisse Verbesserung gegenüber GRPO bot, war FRPO (operierend im Policy-/KL-Raum) überlegen. Die Autoren stellen fest, dass die L2-Kugel von SAM im Parameterraum nachgelagerte Policy-Verschiebungen nicht vollständig erfasst, während die KL-Kugel Verschiebungen der Ausgabeverteilung direkt einschränkt.
• Rolle von $\lambda$: Experimente bestätigten, dass $\lambda$ den Trade-off zwischen mittlerer Belohnung und Varianz steuert. Der optimale $\lambda$ korreliert mit der Standardabweichung der Belohnungen und dem KL-Budget $\rho$.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass Robustheit gegenüber nachgelagertem Fine-Tuning direkt in die Basis-Policy während des RLHF integriert werden sollte, anstatt sich auf spätere Interventionen zu verlassen. Durch die Optimierung auf Belohnungsstabilität innerhalb einer durch KL begrenzten Nachbarschaft identifiziert FRPO Policies, die sich in „belohnungsflachen" Regionen befinden, was sie inhärent resistent gegen katastrophales Vergessen macht.

Die Autoren positionieren FRPO als allgemeine Lösung, die:

• Über verschiedene Domänen hinweg funktioniert (Sicherheit, Mathematik, Programmieren).
• Im Vergleich zu Standard-GRPO keine zusätzlichen Rechenkosten erfordert.
• Einen prinzipiellen, mathematisch fundierten Ansatz für kontinuierliches Lernen in LLMs bietet, der unabhängig von spezifischen nachgelagerten Fine-Tuning-Protokollen ist.

Die Arbeit legt nahe, dass das Prinzip der Optimierung auf Robustheit zukünftiger Anpassungen über RLHF hinaus auf Pretraining und supervised Fine-Tuning ausgedehnt werden kann und eine neue Richtung für Forschungsarbeiten zur Ausrichtung und zum kontinuierlichen Lernen bietet.