🤖 machine learning

$f$ -Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data

Dieses Papier stellt den $f$ -Trajectory-Balance vor, eine Familie von Verlustfunktionen, die den Ansatz des mittleren quadratischen Fehlers auf alle $f$ -Divergenzen erweitert und damit das Training generativer Modelle und von LLMs mit Off-Policy-Daten ermöglicht, wobei die spezifischen Optimierungseigenschaften (wie das Abdecken von Modi) der entsprechenden On-Policy- $f$ -Divergenz-Gradienten erhalten bleiben.

Ursprüngliche Autoren: Jake Fawkes, Jason Hartford

Veröffentlicht 2026-05-18

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Jake Fawkes, Jason Hartford

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie unterrichten einen Roboter, wie er ein riesiges, dunkles Höhlensystem erkundet. Die Höhle verfügt über viele verborgene Kammern (die sogenannten „Modi"), von denen einige Gold enthalten (hohe Belohnungen), während andere nur leer oder gefährlich sind. Ihr Ziel ist es, dem Roboter beizubringen, so viele verschiedene goldene Kammern wie möglich zu finden, anstatt nur eine große Goldmine zu entdecken und alles andere zu ignorieren.

Dieser Artikel stellt eine neue Reihe von „Unterrichtswerkzeugen" (Verlustfunktionen) vor, um diese Roboter zu trainieren, unabhängig davon, ob sie Moleküle für neue Medikamente generieren, Kunst erschaffen oder Code schreiben.

Hier ist die Aufschlüsselung der Ideen des Artikels unter Verwendung einfacher Analogien:

1. Das Problem: Der „Goldrausch" vs. der „Entdecker"

In der Vergangenheit war die Standardmethode, diese Roboter zu unterrichten, wie ein Goldrausch. Der Roboter würde eine Stelle mit etwas Gold finden, sich begeistern und dann seine gesamte Zeit damit verbringen, dort zu graben. Er würde alle anderen Goldminen in der Höhle ignorieren.

In technischen Begriffen: Dies wird als „Modus-Suche" (mode-seeking) bezeichnet. Das Modell kollabiert auf einige wenige Antworten mit hoher Wahrscheinlichkeit und ignoriert den Rest der Vielfalt.
Das Ziel des Artikels: Wir wollen eine „Modus-Abdeckung" (mode-covering). Wir wollen, dass sich der Roboter ausbreitet und alle Goldminen erkundet, auch die kleineren, um eine vollständige Karte der Höhle zu erhalten.

2. Das alte Werkzeug: Der „quadratische KL"-Verlust

Kürzlich haben Forscher einen klugen Trick entdeckt, um Roboter mit einer Methode des „quadratischen Fehlers" zu unterrichten (die Messung des Abstands zwischen dem, was der Roboter denkt, und dem, was er tun sollte).

Die Analogie: Stellen Sie sich einen Lehrer vor, der einen Schüler benotet. Wenn der Schüler die Antwort falsch hat, sagt der Lehrer nicht nur „Falsch". Sie berechnen das Quadrat des Fehlers.
Der Vorteil: Diese Methode ist sehr stabil. Sie funktioniert auch dann, wenn der Schüler aus alten Notizen lernt (Off-Policy-Daten) und nicht aus Live-Unterricht.
Der Mangel: Obwohl sie stabil ist, verhält sich diese spezifische „quadratische" Methode immer noch wie ein Goldrausch. Sie drängt den Roboter natürlich dazu, sich nur auf einige wenige Antworten zu konzentrieren und die Vielfalt der Höhle zu verpassen.

3. Die neue Lösung: Die Familie der „f-Trajektorien-Balance"

Die Autoren erkannten, dass die „quadratische" Methode nur ein spezifischer Geschmack einer viel größeren Familie von Unterrichtswerkzeugen war. Sie nennen diese Familie f-Trajektorien-Balance.

Stellen Sie sich diese Familie als einen Drehregler oder einen Lautstärkeregler an einem Radio vor.

Den Regler in die eine Richtung drehen (Niedrige Zahlen): Sie erhalten Werkzeuge, die wie ein Super-Entdecker wirken. Diese Werkzeuge zwingen den Roboter, sich auszubreiten und jede mögliche Goldmine zu finden, auch die schwer erreichbaren. Dies ist hervorragend für die Wirkstoffentdeckung, wo Sie jedes Molekül finden möchten, das funktionieren könnte, und nicht nur das offensichtlichste.
Den Regler in die andere Richtung drehen (Hohe Zahlen): Sie erhalten Werkzeuge, die wie ein Goldgräber wirken. Diese Werkzeuge sagen dem Roboter, er solle die kleinen Minen ignorieren und sich intensiv auf den größten, offensichtlichsten Goldhaufen konzentrieren. Dies ist nützlich, wenn Sie nur die einzelne beste Antwort wollen.
Der Mittelweg: Sie können den Regler irgendwo dazwischen einstellen, um eine Mischung aus Erkundung und Fokus zu erhalten.

4. Warum dies eine große Sache ist

Der Artikel beweist zwei Hauptpunkte:

Der „magische Schalter": Sie zeigten, dass Sie das Standardwerkzeug „quadratisch" gegen eines dieser neuen „Drehregler"-Werkzeuge austauschen können, und die Mathematik funktioniert immer noch perfekt. Der Roboter lernt genauso gut, aber mit einer anderen Persönlichkeit (mehr erkundend oder mehr fokussiert).
Stabilität: Genau wie das alte Werkzeug funktionieren diese neuen Werkzeuge auch dann, wenn der Roboter aus alten Daten lernt (Off-Policy). Dies ist für reale Anwendungen wie Large Language Models (LLMs) entscheidend, bei denen das Training oft asynchron stattfindet (der Roboter lernt aus Daten, die vor einiger Zeit generiert wurden).

5. Tests in der realen Welt (Die Höhlenkarten)

Die Autoren testeten diese Werkzeuge in drei verschiedenen „Höhlen":

Das Gitter-Spiel: Ein einfaches Computerdabyrinth mit vier Ecken voller Gold. Das Standardwerkzeug fand nur eine Ecke. Die neuen „Entdecker"-Werkzeuge fanden alle vier Ecken schnell.
Molekülforschung (SynFlowNet): Sie versuchten, neue chemische Moleküle zu generieren. Indem sie den Regler auf „Entdecker" stellten, generierten sie eine viel breitere Vielfalt einzigartiger Moleküle, die potenziell Medikamente sein könnten, anstatt nur Variationen derselben wenigen Chemikalien.
Sprachmodelle (LLMs): Sie passten KI-Modelle an, um mathematische Probleme zu lösen. Die neuen Werkzeuge ermöglichten es der KI, verschiedene Wege zur Lösung von Problemen zu erkunden, ohne in einer Schleife festzustecken, die dieselbe Antwort wiederholt, selbst wenn die Trainingsdaten verzögert waren (asynchron).

Zusammenfassung

Dieser Artikel erfindet keinen neuen Roboter-Typ. Stattdessen erfindet er eine neue Reihe von Anweisungen, wie man sie trainiert.

Alter Weg: „Finde die größte Goldmine und grabe dort." (Stabil, aber langweilig).
Neuer Weg: „Hier ist ein Drehregler. Sie können wählen, ob Sie ein Goldgräber, ein Super-Entdecker oder etwas dazwischen sein möchten. Und egal wofür Sie sich entscheiden, das Training wird stabil sein und mit alten Daten funktionieren."

Dies gibt Ingenieuren einen einfachen Regler, um zu steuern, wie kreativ oder fokussiert ihre KI-Modelle sein sollen, ohne den Trainingsprozess zu stören.

Technische Zusammenfassung: f-Trajektorien-Balance

Problemstellung

Die Ausrichtung und Feinabstimmung generativer Modelle, einschließlich Large Language Models (LLMs) und Generativer Flussnetzwerke (GFlowNets), stützt sich typischerweise auf die Minimierung der Diskrepanz zwischen der Verteilung eines Modells und einer Zielverteilung. Beim Reinforcement-Learning (RL) im Rahmen der Feinabstimmung wird dies häufig als Maximierung der erwarteten Belohnung unter einer Kullback-Leibler-(KL)-Divergenz-Strafkomponente formuliert.

Eine erhebliche Herausforderung in diesem Bereich ist der Kompromiss zwischen modensuchendem (mode-seeking) und modenabdeckendem (mode-covering) Verhalten. Standardziele, wie die reverse KL-Divergenz (oft implementiert über Trajektorien-Balance oder quadratischen KL-Verlust), neigen dazu, modensuchend zu sein, wodurch Modelle auf wenige hochwahrscheinliche Modi kollabieren. Dies ist für Aufgaben suboptimal, die Exploration erfordern, wie etwa die Wirkstoffentdeckung oder die Generierung vielfältiger Kandidaten, wo „modenabdeckendes" Verhalten (Erfassung der gesamten Vielfalt der Zielverteilung) bevorzugt wird.

Darüber hinaus kämpfen bestehende Methoden oft mit Off-Policy-Gültigkeit. Während Schätzer wie REINFORCE oder PPO üblich sind, leiden sie unter hoher Varianz. Neuere Ansätze, die den quadratischen KL-Verlust ( $KL_{sq}$ ) verwenden, bieten on-policy Gradienten mit niedriger Varianz und Off-Policy-Gültigkeit, sind jedoch inhärent an die Eigenschaften der reverse KL-Divergenz gebunden, was die Fähigkeit zur Kontrolle des Trade-offs zwischen Exploration und Exploitation einschränkt.

Methodik

Die Autoren schlagen f-Trajektorien-Balance vor, eine Familie von Surrogat-Verlustfunktionen, die den quadratischen KL-Verlust auf die gesamte Familie der f-Divergenzen verallgemeinert.

Theoretische Grundlage

Der zentrale theoretische Beitrag besteht in der Herleitung einer Eins-zu-Eins-Entsprechung zwischen translationsinvarianten Verlustfunktionen auf Log-Wahrscheinlichkeiten und f-Divergenzen.

Von Divergenz zu Verlust: Für jede konvexe Funktion $f$ , die eine f-Divergenz $D_f(p_\theta \| p^*)$ definiert, leiten die Autoren eine Verlustfunktion $L_f$ ab, die auf der Differenz der Log-Wahrscheinlichkeiten $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$ definiert ist:
$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$
Dieser Verlust besitzt zwei kritische Eigenschaften:
- On-Policy-Gradient-Äquivalenz: Wenn Stichproben aus der Modellpolitik gezogen werden ( $\mu = p_\theta$ ), stimmen die erwarteten automatisch differenzierten Gradienten von $L_f$ mit den Gradienten der entsprechenden f-Divergenz $D_f$ überein.
- Off-Policy-Gültigkeit: Wenn auf Off-Policy-Daten ausgewertet wird (gestichprobt aus einer Verhaltensverteilung $\mu$ ), bleibt der Verlust ein gültiges Ziel mit demselben globalen Minimierer ( $p_\theta = p^*$ ), sofern der Träger der Zielverteilung im Träger der Verhaltenspolitik enthalten ist.
Von Verlust zu Divergenz: Umgekehrt entspricht jede strikt konvexe, translationsinvariante Verlustfunktion $\ell$ auf Log-Wahrscheinlichkeiten der Minimierung einer spezifischen f-Divergenz. Die Arbeit beweist, dass diese Abbildungen zueinander invers sind.

Praktische Implementierungen

Um praktische Herausforderungen beim Training zu adressieren, führen die Autoren mehrere Erweiterungen ein:

DevGrad-Verlust (Batch-weise Normalisierung): Um mit unlösbaren Partitionierungsfunktionen umzugehen (häufig bei GFlowNets und nicht-normalisierten Zielen), verallgemeinern die Autoren den VarGrad-Schätzer. Anstatt die Batch-Varianz zu verwenden, nutzen sie die batch-weise generalisierte Abweichung der Log-Wahrscheinlichkeitsdifferenzen. Dies beinhaltet die Schätzung der Normalisierungskonstante $\log Z$ durch Minimierung des Verlusts über den Batch, was effektiv die Koeffizienten der Score-Funktion zentriert und die Varianz reduziert.
Temperierter Verlust: Um numerisches Überlaufen bei kleinen Temperaturparametern $\beta$ zu verhindern (wobei die Zielverteilung $p^* \propto \pi_{ref} \exp(r/\beta)$ ist), schlagen die Autoren einen temperierten Verlust vor. Dieser skaliert den Verlust mit $1/\beta$ und wendet ihn auf die temperierte Verteilung an, wodurch sichergestellt wird, dass die Gradienten beschränkt und stabil bleiben.
f-Trajektorien-Balance für GFlowNets: Das Framework wird auf GFlowNets angewendet, indem Trajektorien-Balance-Constraints unter Verwendung des $L_f$ -Verlusts erzwungen werden. Dies verallgemeinert die Standard-Trajektorien-Balance (die der reverse KL entspricht) auf beliebige f-Divergenzen.

Hauptbeiträge

Verallgemeinerte Verlustfamilie: Herleitung einer allgemeinen Familie translationsinvarianter Surrogat-Verluste ( $L_f$ ), deren On-Policy-Gradienten mit denen der entsprechenden f-Divergenz übereinstimmen, während die Off-Policy-Gültigkeit erhalten bleibt.
Theoretische Äquivalenz: Beweis der inversen Beziehung zwischen translationsinvarianten Log-Wahrscheinlichkeits-Verlusten und f-Divergenzen, was die systematische Entwicklung neuer Verlustfunktionen mit gewünschten Eigenschaften (z. B. modenabdeckend) ermöglicht.
Erweiterung auf nicht-normalisierte Ziele: Verallgemeinerung der VarGrad-Technik auf beliebige f-Divergenzen durch den DevGrad-Verlust, wodurch das Problem der unlösbaren Partitionierungsfunktion beim RL-Feintuning adressiert wird.
Kontrolle der Exploration: Demonstration, dass die Wahl der f-Divergenz (parametrisiert durch $\alpha$ in $\alpha$ -Divergenzen) das modensuchende versus modenabdeckende Verhalten direkt steuert. Niedrigere $\alpha$ -Werte (z. B. Forward KL, Hellinger) fördern modenabdeckendes Verhalten, während höhere Werte (z. B. Pearson $\chi^2$ ) modensuchendes Verhalten fördern.

Experimentelle Ergebnisse

Die Autoren bewerten die Verlustfamilie über vier verschiedene Aufgaben hinweg:

Synthetisches Raster (GFlowNets): Bei einer 2D-Raster-Aufgabe, die entwickelt wurde, um die Modenabdeckung zu testen, entdeckten Verluste mit niedrigerem $\alpha$ (Forward KL, Hellinger) alle vier Modi der Zielverteilung deutlich schneller und zuverlässiger als die Standard-Trajektorien-Balance (Reverse KL). Umgekehrt zeigte Pearson $\chi^2$ ( $\alpha=2$ ) stark modensuchendes Verhalten und blieb im ersten gefundenen Modus stecken.
Molekülenentdeckung (SynFlowNet): Angewendet auf das SynFlowNet-Framework zur Generierung synthetisch zugänglicher Moleküle. Die Autoren fanden heraus, dass die Feinabstimmung von $\alpha$ die Vielfalt der generierten Moleküle steuern konnte. Ein Abkühlungsplan (Start mit niedrigerem $\alpha$ für Exploration und Erhöhung auf $\alpha=1$ für Exploitation) ergab den besten Kompromiss und generierte Moleküle mit höherer Belohnung und größerer Vielfalt im Vergleich zur Standard-Trajektorien-Balance.
Diffusionsmodelle: Beim bedingten Sampling von MNIST-Ziffern (ungerade vs. gerade) versagte die Standard-Trajektorien-Balance beim gleichmäßigen Sampling über die Zielmodi hinweg (Übersampling spezifischer Ziffern wie 0 und 6). f-Trajektorien-Balance-Verluste, insbesondere solche mit niedrigerem $\alpha$ , erreichten ein gleichmäßigeres Sampling über die Zielmodi hinweg.
Asynchrones LLM-Feintuning: Die Verluste wurden angewendet, um LLMs (Qwen- und OLMo-Familien) bei mathematischen Problemen (GSM8k, MATH) mit einer asynchronen Verzögerung von 50 Schritten zu feinabstimmen. Die f-Divergenz-Verluste zeigten stabiles Training und konsistente Entropie-Belohnungs-Trade-offs, während Standard-PPO-Implementierungen unter denselben asynchronen Bedingungen Instabilität zeigten.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass die Wirksamkeit des quadratischen KL-Verlusts kein isoliertes Phänomen ist, sondern ein spezifisches Beispiel einer breiteren Familie von Verlusten. Durch die Erweiterung der Eigenschaft des „On-Policy-Gradient-Matchings" und der „Off-Policy-Gültigkeit" auf die gesamte Familie der f-Divergenzen bieten die Autoren ein einheitliches Framework für die Feinabstimmung generativer Modelle.

Die primäre Bedeutung liegt in der Fähigkeit, den Exploration-Exploitation-Trade-off explizit zu steuern, einfach durch die Auswahl einer anderen Surrogat-Verlustfunktion (oder das Einstellen eines $\alpha$ -Parameters), ohne den zugrunde liegenden Trainingsalgorithmus zu ändern. Dies ermöglicht Praktikern, zwischen modensuchendem (zur Maximierung der Belohnung) und modenabdeckendem (für Vielfalt und Exploration) Verhalten zu wählen, während die Stabilität und Off-Policy-Kompatibilität erhalten bleiben, die für das moderne, großskalige asynchrone Training von LLMs und GFlowNets erforderlich sind. Die Arbeit legt nahe, dass dieses Framework besonders wertvoll für Anwendungen wie die Wirkstoffentdeckung und explorative Agenten ist, bei denen die Abdeckung des Raums hochbelohnter Kandidaten entscheidend ist.