Evaluating GFlowNet from partial episodes for stable and flexible policy-based training

Diese Arbeit verbindet die wertbasierte und die politikbasierte Sichtweise von GFlowNets, indem sie zeigt, dass Flussgleichgewicht auch einen fundierten Politikbewerter liefert, und schlägt ein Evaluierungsgleichgewichts-Objektiv über Teil-Episoden vor, das das Training stabiler macht und die Integration parametrisierter Rückwärtsrichtlinien sowie Offline-Daten ermöglicht.

Das Wesentliche

Titel: Wie man einen perfekten Schatzsucher baut – Eine einfache Erklärung von GFlowNets und dem neuen „Sub-EB"-Trick

Stell dir vor, du bist ein Schatzsucher in einem riesigen, verworrenen Labyrinth (dem „combinatorial space"). Dein Ziel ist es, alle möglichen Schätze zu finden, aber nicht nur irgendeinen – du willst die besten Schätze finden, und zwar in einer Häufigkeit, die ihrer Qualität entspricht. Ein sehr wertvoller Diamant soll öfter gefunden werden als ein kleiner Kieselstein.

Das Problem? Das Labyrinth ist so groß, dass du nie alle Wege ablaufen kannst, um zu wissen, wo die besten Schätze liegen. Du musst also einen Weg finden, der dich intelligent durch das Labyrinth führt, ohne jedes einzelne Steinchen zu prüfen.

Hier kommt GFlowNet ins Spiel.

1. Das Grundproblem: Der verwirrte Wanderer

GFlowNets sind wie ein cleverer Wanderer, der durch das Labyrinth läuft. Er trifft Entscheidungen an jedem Abzweig (einem „Zustand"), um zum nächsten Punkt zu kommen.

• Die alte Methode (Werte-basiert): Der Wanderer hat eine Landkarte, auf der für jeden Pfad ein „Fluss-Wert" steht. Er versucht, sicherzustellen, dass die Menge an Wasser (Fluss), die durch einen Pfad fließt, genau der Menge entspricht, die am Ende ankommt. Das funktioniert gut, ist aber wie das Berechnen des Wasserdrucks in jedem einzelnen Rohr – sehr rechenintensiv und starr.
• Die neue Methode (Policy-basiert): Der Wanderer hat einen inneren Kompass (eine „Bewertungsfunktion" oder V). Dieser Kompass sagt ihm: „Hey, von hier aus ist die Wahrscheinlichkeit, einen tollen Schatz zu finden, eigentlich noch nicht so gut wie sie sein könnte." Der Wanderer passt dann seine Schritte an, um diesen Fehler zu korrigieren.

Das Problem mit dem Kompass: Bis jetzt war es sehr schwierig, diesen Kompass genau zu kalibrieren. Wenn der Kompass falsch ist, läuft der Wanderer in die falsche Richtung oder bleibt stecken. Die Schätzung des „Fehlers" war oft unzuverlässig.

2. Die Lösung: Der „Sub-EB"-Trick (Teil-Episode-Balance)

Die Autoren dieses Papers haben eine geniale Verbindung hergestellt. Sie haben erkannt, dass die alte „Fluss-Methode" und die neue „Kompass-Methode" eigentlich zwei Seiten derselben Medaille sind.

Stell dir vor, der Wanderer läuft nicht nur den ganzen Weg bis zum Ende ab, um zu sehen, ob er richtig lag. Stattdessen schaut er sich Teilstrecken an.

• Die Analogie: Stell dir vor, du lernst ein neues Instrument.
  • Die alte Methode sagt: „Spiele das ganze Stück durch. Wenn am Ende alles klingt, war es gut." (Das ist schwer zu bewerten, wenn das Stück sehr lang ist).
  • Die neue Methode (Sub-EB) sagt: „Schau dir nur die nächsten drei Takte an. War die Harmonie zwischen dem, was du gerade gespielt hast, und dem, was kommen sollte, im Gleichgewicht?"

Der neue Trick heißt Sub-EB (Subtrajectory Evaluation Balance). Er besagt:

„Wenn du an einem Punkt im Labyrinth stehst, muss die Differenz zwischen deinem aktuellen Gefühl (dem Kompass) und dem Ziel genau der Differenz entsprechen, die du über den nächsten kleinen Abschnitt hinweg erwarten würdest."

Das ist wie ein Gleichgewichtsspiel: Der Wanderer prüft ständig, ob sein inneres Gefühl für die Zukunft mit der Realität der nächsten Schritte übereinstimmt. Wenn ja, ist der Kompass perfekt kalibriert.

3. Warum ist das so toll? (Die Vorteile)

1. Stabilität (Der ruhige Wanderer)
Bisher war der Kompass manchmal verrückt. Mal zeigte er nach Norden, mal nach Süden, obwohl der Wanderer sich kaum bewegt hatte. Mit Sub-EB wird der Kompass viel stabiler. Der Wanderer lernt zuverlässiger und schneller, ohne hin und her zu schwanken. Das ist wie ein GPS, das nicht mehr ständig „Neu berechnen" sagt, sondern einen klaren Weg weist.

2. Flexibilität (Der flexible Wanderer)
Früher musste der Wanderer immer genau so laufen, wie er es gerade gelernt hat (On-Policy). Er durfte keine alten Daten von einem anderen Wanderer nutzen.
Mit Sub-EB kann der Wanderer jetzt auch alte Aufzeichnungen (Offline-Daten) nutzen. Er kann sagen: „Schau mal, dieser andere Wanderer war hier schon mal und hat einen tollen Weg gefunden. Ich lerne daraus, auch wenn ich selbst noch nicht dort war." Das macht das Training viel effizienter.

3. Der Rückwärts-Trick (Der Rückwärts-Planer)
Normalerweise plant man einen Weg nur vorwärts. Sub-EB erlaubt es dem System, auch einen Rückwärts-Planer zu haben, der dynamisch lernt. Stell dir vor, du kannst nicht nur planen, wie du zum Schatz kommst, sondern auch lernen, wie man vom Schatz zurück zum Start geht, um den Weg besser zu verstehen. Das System passt beide Richtungen gleichzeitig an, was viel schneller zum Ziel führt.

4. Wo wurde es getestet?

Die Autoren haben ihren neuen Kompass in verschiedenen Welten getestet:

• Künstliche Gitter: Wie ein riesiges Schachbrett, auf dem man die besten Züge finden muss.
• Biologie & Chemie: Hier hilft der Wanderer, neue Medikamente oder DNA-Sequenzen zu entwerfen. Er findet Moleküle, die wie perfekte Schlüssel für Krankheiten sind.
• Netzwerk-Strukturen: Er hilft dabei, die beste Struktur für ein soziales Netzwerk oder ein biologisches System zu finden.

In allen Fällen hat der Wanderer mit dem neuen Sub-EB-Kompass schneller gelernt, bessere Ergebnisse erzielt und war robuster gegen Fehler als die alten Methoden.

Fazit

Dieses Paper ist wie die Einführung eines super-intelligenten Navigators für KI-Systeme, die komplexe Dinge erschaffen sollen. Anstatt nur zu raten, ob ein ganzer Weg gut ist, prüft der Navigator ständig kleine Abschnitte auf ihr Gleichgewicht. Das macht das Lernen stabiler, schneller und erlaubt es der KI, aus vergangenen Erfahrungen (auch von anderen) zu lernen, ohne jedes Rad neu zu erfinden.

Kurz gesagt: Sub-EB ist der Schlüssel, um KI-Schatzsucher von chaotischen Anfängern zu perfekten Meistern zu machen.

Technische Zusammenfassung

1. Problemstellung

Generative Flow Networks (GFlowNets) sind generative Modelle, die darauf abzielen, kombinatorische Objekte (z. B. Graphen, Sequenzen) so zu sampeln, dass ihre Wahrscheinlichkeit proportional zu einer Belohnungsfunktion $R(x)$ ist. Das Training von GFlowNets erfolgt typischerweise über zwei Hauptansätze:

• Wertbasierte Methoden (Value-Based): Diese nutzen Fluss-Balance-Bedingungen (z. B. Sub-Trajectory Balance, Sub-TB), um den Fluss durch den Zustandsraum zu modellieren. Sie sind robust, erfordern jedoch oft Off-Policy-Datensammlung und können bei der Schätzung von Flüssen in großen Räumen instabil sein.
• Policy-basierte Methoden (Policy-Based): Diese nutzen einen Actor-Critic-Ansatz, ähnlich wie im Reinforcement Learning (RL). Ein Critic (eine Bewertungsfunktion $V(s)$) schätzt die Divergenz zwischen der aktuellen Vorwärts-Policy und der Zielverteilung, um die Policy zu aktualisieren.

Das zentrale Problem: Bei policy-basierten Methoden ist die zuverlässige Schätzung der Bewertungsfunktion $V(s)$ (die die Kullback-Leibler-Divergenz approximiert) schwierig. Bisherige Ansätze (wie $\lambda$-TD) leiden unter hoher Varianz oder Bias, wenn sie nur auf aktuellen Daten basieren, und erfordern oft starre Annahmen (z. B. eine feste Rückwärts-Policy $\pi_B$). Zudem ist die Integration von Offline-Daten oder parametrisierten Rückwärts-Policies in diesen Rahmenwerken komplex.

2. Methodik: Sub-Trajectory Evaluation Balance (Sub-EB)

Die Autoren schlagen eine neue Methode vor, die die Lücke zwischen wertbasierten und policy-basierten Ansätzen schließt, indem sie die Fluss-Balance nutzen, um eine Evaluation-Balance für die Bewertungsfunktion $V(s)$ abzuleiten.

• Theoretische Verbindung: Die Arbeit zeigt, dass für eine feste Vorwärts-Policy $\pi_F$ die Lösung der Fluss-Balance-Bedingung für den Zustandsfluss $F(s)$ exakt der wahren Bewertungsfunktion $V(s)$ (definiert als KL-Divergenz) entspricht.
• Sub-EB-Bedingung: Es wird eine neue Bedingung, die Sub-Trajectory Evaluation Balance (Sub-EB), eingeführt. Diese besagt, dass für jede Teiltrajektorie (Partial Episode) zwischen zwei Zuständen $s_i$ und $s_j$ die erwartete Differenz der gelernten Divergenzen ($V(s_i) - V(s_j)$) der wahren Divergenz der Trajektorien zwischen diesen Zuständen entsprechen muss.
  • Mathematisch: $E_{P_F}[\log(P_F(\tau_{i:j}|s_i) \exp V(s_i))] = E_{P_F}[\log(P_B(\tau_{i:j}|s_j) \exp V(s_j))]$.
• Sub-EB-Objektivfunktion: Basierend auf dieser Bedingung wird eine neue Verlustfunktion $L_V$ definiert, die den quadratischen Fehler der Sub-EB-Bedingung über alle Teiltrajektorien minimiert. Im Gegensatz zu $\lambda$-TD, das nur auf Kantenfehlern und Startpunkten basiert, nutzt Sub-EB Informationen von Teiltrajektorien, die vor und nach dem aktuellen Zustand enden/starten. Dies führt zu einem stabileren und ausgewogeneren Lernen von $V(s)$.
• Erweiterungen:
  • Parametrisierte Rückwärts-Policy ($\pi_B$): Sub-EB erlaubt es, $\pi_B$ gemeinsam mit $V$ zu optimieren, ohne separate Trainingsphasen oder zusätzliche Ziele zu benötigen.
  • Offline-Training: Durch die Einführung einer inversen Bewertungsfunktion $W(s)$ für die Rückwärtsrichtung wird ein Workflow ermöglicht, der Offline-Datensammlung ($\pi_D \neq \pi_F$) unterstützt, was für die Exploration in großen Suchräumen entscheidend ist.

3. Wichtige Beiträge

1. Theoretische Brücke: Etablierung einer formalen Verbindung zwischen dem Zustandsfluss $F(s)$ und der Bewertungsfunktion $V(s)$, die zeigt, dass Fluss-Balance eine hinreichende Bedingung für die korrekte Evaluation von Policies darstellt.
2. Neues Lernziel (Sub-EB): Entwicklung der Sub-EB-Objektivfunktion, die Teiltrajektorien als Grundeinheit nutzt und so eine zuverlässigere Schätzung von $V(s)$ ermöglicht als bestehende $\lambda$-TD-Methoden.
3. Flexibilität: Demonstration, dass Sub-EB parametrisierte Rückwärts-Policies und Offline-Datensammlung nahtlos in den policy-basierten Trainingsprozess integriert, was die Flexibilität und Effizienz erheblich steigert.
4. Umfassende Evaluation: Validierung der Methode auf synthetischen (Hypergrids) und realen Datensätzen (biologische Sequenzen, molekulare Graphen, Bayesian Network Struktur-Lernen).

4. Ergebnisse

Die Experimente wurden auf verschiedenen Aufgaben durchgeführt und verglichen mit State-of-the-Art-Methoden wie Sub-TB (wertbasiert), RL (policy-basiert mit $\lambda$-TD), CV (Control Variates) und Q-Much.

• Hypergrids (Synthetisch):
  • Sub-EB zeigte eine signifikant verbesserte Stabilität und Konvergenzgeschwindigkeit im Vergleich zum RL-Ansatz mit $\lambda$-TD, insbesondere in hochdimensionalen Räumen (z. B. $256 \times 256$).
  • Die Verwendung einer parametrisierten Rückwärts-Policy mit Sub-EB (Sub-EB-P) führte zu den besten Ergebnissen in Bezug auf Stabilität und Endleistung.
• Sequenzdesign (Biologie/Moleküle):
  • Sub-EB ermöglichte die effektive Integration von Offline-Sampling-Techniken (Sub-EB-B), was zu einer schnelleren Entdeckung von High-Reward-Moden führte, ohne die Verteilungsmodellierung zu stark zu beeinträchtigen.
  • In der Verteilungsmodellierung (gemessen an Total Variation und JSD) schnitt Sub-EB besser oder gleichauf mit RL ab, übertraf aber Sub-TB deutlich.
• Bayesian Network (BN) Struktur-Lernen:
  • In großen kombinatorischen Räumen (bis zu $10^{35}$ mögliche Graphen) erreichte Sub-EB die höchste durchschnittliche Belohnung und konvergierte schneller als Sub-TB und Q-Much.
  • Sub-EB zeigte eine hohe Diversität der gefundenen Lösungen, was auf eine gute Balance zwischen Exploration und Exploitation hindeutet.
• Molekulare Graph-Design:
  • Sub-EB erzielte die höchste durchschnittliche Belohnung bei der JNK3-Aufgabe und die schnellste Konvergenz, während die Diversität mit anderen Top-Methoden vergleichbar blieb.

5. Bedeutung und Fazit

Diese Arbeit stellt einen bedeutenden Fortschritt im Training von GFlowNets dar, indem sie die Vorteile wertbasierter Fluss-Balance-Methoden in den policy-basierten Rahmen überträgt.

• Stabilität: Sub-EB löst das Problem der instabilen Schätzung der Bewertungsfunktion $V(s)$, was ein Hauptlimitierungsfaktor bei policy-basierten GFlowNet-Methoden war.
• Flexibilität: Die Methode macht policy-basiertes Training robuster gegenüber der Wahl der Datensammlungspolitik und erlaubt die Nutzung von parametrisierten Rückwärts-Policies und Offline-Daten.
• Skalierbarkeit: Die Ergebnisse zeigen, dass Sub-EB effektiv auf sehr große kombinatorische Suchräume skaliert und sowohl für die Suche nach optimalen Lösungen (Mode Seeking) als auch für die genaue Nachbildung von Verteilungen geeignet ist.

Zusammenfassend bietet Sub-EB einen principleden und effizienten Weg, um GFlowNets zu trainieren, der die Zuverlässigkeit von wertbasierten Ansätzen mit der Flexibilität von policy-basierten Ansätzen vereint.