Technischer Zusammenfassung: POETS (Policy Ensembles for Thompson Sampling)

Problemstellung

Die Balance zwischen Exploration und Exploitation bleibt eine Kernherausforderung in der sequenziellen Entscheidungsfindung und der Black-Box-Optimierung, insbesondere bei der Anwendung auf Large Language Models (LLMs). Eine effektive Exploration erfordert eine prinzipielle Quantifizierung von Unsicherheit. Während das Ensembling unabhängig trainierter Modelle eine bewährte, modellagnostische Methode zur Schätzung epistemischer Unsicherheit darstellt, ist die direkte Anwendung naiver Ensembling-Verfahren auf LLMs aufgrund von Speicher- und Rechenbeschränkungen rechnerisch prohibitiv. Darüber hinaus verlassen sich bestehende Explorationsmethoden für LLMs häufig auf explizite Belohnungsmodelle (was verschachtelte Optimierungskomplexitäten einführt) oder sind stark von den Instruktionsfolgefähigkeiten des Modells abhängig.

Methodik

Die Arbeit stellt POETS (Policy Ensembles for Thompson Sampling) vor, ein Framework, das Unsicherheitsquantifizierung und Policy-Optimierung verbindet, ohne ein explizites probabilistisches Belohnungsmodell zu erfordern.

Kernidee: Implizite Belohnungskodierung

POETS basiert auf der Dualität zwischen Belohnungsfunktionen und optimalen Policies unter regularisierten Zielen. Konkret kodieren Policies, die mit Kullback-Leibler (KL)-Regularisierung trainiert wurden, implizit eine zugrunde liegende Belohnungsfunktion. Anstatt ein Belohnungsmodell zu lernen und dann eine Policy daran anzupassen (ein zweistufiger Prozess), passt POETS die impliziten Belohnungsfunktionen, die von einem Policy-Ensemble kodiert werden, direkt an beobachtete Online-Daten an.

Der Algorithmus

1. Policy-Ensemble: POETS verwaltet ein Ensemble von $n$ unabhängigen Policies, $\Pi = \{\pi_1, \dots, \pi_n\}$.
2. Poisson-Bootstrapping: Um epistemische Unsicherheit zu erfassen, werden die Ensemble-Mitglieder durch Poisson-Bootstrapping diversifiziert. Für jeden eingehenden Datenbatch erhält jedes Ensemble-Mitglied ein unabhängiges ganzzahliges Gewicht, das aus einer Poisson(1)-Verteilung gezogen wird. Dies ermöglicht dem Ensemble, eine approximative Bayes'sche Posterior-Verteilung über Belohnungsfunktionen zu lernen, ohne einen statischen Datensatz zu benötigen.
3. Verlustfunktion: Die Methode optimiert eine Verlustfunktion, die die implizite Belohnung einer einzelnen Policy mit der wahren Belohnung abgleicht. Eine theoretische Herleitung zeigt, dass die Optimierung dieser Verlustfunktion mathematisch äquivalent zu einem Policy-Gradient-Update (speziell eine Variante von GRPO) ist, skaliert mit den Poisson-Gewichten.
4. Thompson Sampling: Durch das Anpassen eines impliziten Belohnungsensembles an wahre Belohnungen unter Verwendung unabhängiger Bootstraps lernt POETS eine Posterior-Verteilung über Belohnungsfunktionen. Die Ausführung von KL-regularisiertem Thompson Sampling wird erreicht, indem einfach eine Policy gleichmäßig aus dem Ensemble $\Pi$ gezogen und entsprechend gehandelt wird.

Effiziente Architektur: „Trunk & Branch"

Um die prohibitiven Kosten der Aufrechterhaltung von $n$ vollständigen LLM-Kopien zu überwinden, nutzt POETS eine parameter-effiziente Architektur:

• Gemeinsamer Trunk: Ein vortrainierter Backbone (der Großteil der Gewichte) wird über alle Ensemble-Mitglieder hinweg geteilt und nur einmal pro Vorwärts-/Rückwärtsdurchlauf evaluiert.
• Unabhängige Branches: Die Ensemble-Diversität wird vollständig in unabhängigen, low-rank Adaptation (LoRA)-Modulen eingefangen, die als letzte Ausgabeschichten fungieren.
• Effizienz: Dieses Design stellt sicher, dass das Generieren von Antworten und das Akkumulieren von Gradienten für $n$ Policies praktisch die gleiche Wanduhrzeit erfordert wie für eine einzelne Policy, wobei der Speicheraufwand nur mit der Größe der LoRA-Branches skaliert.

Hauptbeiträge

1. Neues Framework: Einführung von POETS, das epistemische Unsicherheit über Policy-Ensembles erfasst und implizit KL-regularisiertes Thompson Sampling ohne explizites Belohnungsmodell durchführt.
2. Theoretische Garantien: Beweis, dass POETS die starken kumulativen weichen Regret-Schranken von $O(\sqrt{T\gamma_T})$ erbt, die mit KL-regularisiertem Thompson Sampling verbunden sind, wobei $\gamma_T$ der maximale Informationsgewinn ist.
3. Rechnerische Effizienz: Demonstration, dass eine „Trunk & Branch"-LoRA-Architektur tiefes Ensembling mit minimalem Overhead ermöglicht (z. B. ~7,7 % Laufzeitsteigerung für $n=16$ auf einer H200-GPU).
4. Empirische Leistung: State-of-the-Art-Probeneffizienz über diverse kombinatorische Black-Box-Optimierungsbereiche hinweg (FAQ-Verfeinerung, Proteinsuche, Quantenschaltungsdesign) und verbesserte Robustheit in Reinforcement-Learning-Szenarien, insbesondere in Off-Policy-Szenarien mit Experience Replay.

Experimentelle Ergebnisse

Die Autoren bewerteten POETS gegen Baselines, darunter GRPO, TOSFIT (das ein explizites Belohnungsmodell verwendet), FIBO (In-Context-Thompson-Sampling) und EVOSEARCH (evolutionäre Suche).

• Black-Box-Optimierung: POETS erreichte State-of-the-Art-Probeneffizienz bei Proteinsuche, Quantenschaltungsdesign und FAQ-Verfeinerung. Bemerkenswerterweise übertraf es die Baselines sogar ohne Replay-Puffer, während Standard-GRPO versagte, die Diversität aufrechtzuerhalten und sich an frühe Kandidaten anpasste (Overfitting).
• Reinforcement Learning (RL): Im AIME-Datensatz-Setting (kontextueller Bandit mit überprüfbaren Belohnungen) verbesserte POETS die Optimierungspfade sowohl für On-Policy- als auch für Off-Policy-Settings (Experience Replay). Es gelang ihm erfolgreich, die bei GRPO beobachteten Overfitting-Probleme bei wiederholtem Kontakt mit Kontexten zu mildern.
• Diversität: Das Ensemble bewahrte während des Trainings die Generierungsdiversität (gemessen durch die Jensen-Shannon-Divergenz) und verhinderte den Entropiekollaps, der bei Standard-Policy-Gradient-Methoden üblich ist.
• Ablationsstudien: Experimente bestätigten, dass Poisson-Bootstrapping für die Diversität unerlässlich ist, dass kleine Ensembles ($n=8$) für hohe Leistung ausreichen und dass die „Trunk & Branch"-Architektur Unsicherheit mit minimalem Rechenaufwand effektiv erfasst.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass POETS einen bedeutenden Fortschritt darstellt, indem es prinzipielle, unsicherheitsbewusste Exploration für LLMs in der wissenschaftlichen Entdeckung und im RL ermöglicht, ohne die rechnerische Last expliziter Belohnungsmodellierung oder die Fragilität von In-Context-Instruktionen.

• Theoretisch: Es wird festgestellt, dass Policy-Ensembles mit Bootstrapping ausreichen, um KL-regularisiertes Thompson Sampling zu implementieren, und dabei optimale Regret-Schranken erben.
• Praktisch: Es bietet eine skalierbare Lösung, die tiefes Ensembling für LLMs machbar macht und die effektive Nutzung von Experience Replay sowie robuste Exploration in Regimen mit hoher Unsicherheit ermöglicht.
• Umfang: Die Arbeit wird als grundlegender algorithmischer Beitrag für das Post-Training von LLMs und die wissenschaftliche Entdeckung positioniert, speziell mit Zielrichtung auf kombinatorische Optimierungsprobleme, bei denen Bewertungen kostspielig sind.

Die Autoren weisen darauf hin, dass POETS zwar die Probeneffizienz und Robustheit verbessert, aber keine Misspezifikation von Belohnungen, Fairness oder Alignment adressiert, die weiterhin Verantwortlichkeiten der Systemdesigner bleiben. Als zukünftige Arbeit wird vorgeschlagen, das Framework von kontextuellen Banditen auf Mehrschritt-Szenarien mit feinkörniger Kreditvergabe zu erweitern.