Technisches Resümee: AEGPO – Adaptive Entropie-gesteuerte Policy-Optimierung für Diffusionsmodelle

1. Problemstellung

Reinforcement Learning from Human Feedback (RLHF), insbesondere Group Relative Policy Optimization (GRPO), hat sich zu einem Standardparadigma für die Ausrichtung von Diffusions- und Flow-Matching-Modellen an menschlichen Präferenzen entwickelt. Bestehende GRPO-Varianten leiden jedoch unter zwei grundlegenden Ineffizienzen aufgrund ihrer Abhängigkeit von statischen Heuristiken:

1. Ineffiziente Stichprobenallokation: Standardmethoden weisen jeder Prompt eine gleichmäßige Anzahl an Rollouts zu und ignorieren dabei die signifikante Variation des „Lernwerts“ zwischen den einzelnen Stichproben. Einige Prompts bewirken minimale Änderungen der Policy, während andere substanzielle Verbesserungen vorantreiben.
2. Statische Zeitschritt-Exploration: Explorationsstrategien (z. B. Verzweigungen in baumbasierten Rollouts) folgen typischerweise fest vorgegebenen Zeitplänen. Diese versäumen es, sich an die dynamische Natur des Denoising-Prozesses anzupassen, bei dem kritische Momente für die Exploration zwischen verschiedenen Prompts und Trainingsphasen variieren.

Diese Einschränkungen führen zu verschwendeten Rechenressourcen bei niederwertigen Stichproben und einer suboptimalen Exploration an nicht-kritischen Zeitschritten, was die Konvergenz verlangsamt und die endgültige Alignment-Leistung begrenzt.

2. Methodik: AEGPO

Die Autoren schlagen die Adaptive Entropy-Guided Policy Optimization (AEGPO) vor, ein Plug-and-Play-Framework, das intrinsische Attention-Signale nutzt, um sowohl die Stichprobenauswahl als auch die Exploration dynamisch zu steuern. Die zentrale Erkenntnis ist, dass die Attention-Entropie als ein leistungsstarkes duales Signal-Proxy dient.

2.1 Kernsignale

Die Methode definiert zwei intrinsische Metriken, die aus den Attention-Maps des Modells abgeleitet werden:

• $\Delta$Entropy (Stichproben-Ebene Signal): Die relative Änderung der Attention-Dispersion zwischen der aktuellen Policy und der Basis-Policy.
  • Definition: $\Delta\text{Entropy} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropy}_{\theta}(t) - \text{Entropy}_{\theta_{base}}(t)|$.
  • Funktion: Es fungiert als robuster Proxy für den Lernwert einer Stichprobe. Eine hohe $\Delta$Entropy deutet darauf hin, dass ein Prompt das Modell dazu zwingt, seine Attention-Strategie signifikant zu reorganisieren, was auf ein hohes Lernpotenzial hinweist.
• Entropy($t$) (Zeitschritt-Ebene Signal): Die absolute Attention-Dispersion bei einem spezifischen Denoising-Schritt $t$.
  • Definition: Die mittlere Shannon-Entropie der Attention-Gewichte über die Text-Token für die Bild-Features bei Schritt $t$.
  • Funktion: Peaks in Entropy($t$) identifizieren kritische Explorations-Zeitschritte. Diese Peaks entsprechen Momenten hoher Unsicherheit oder breiter semantischer Integration (oft in frühen Strukturierungs- und späten Verfeinerungsphasen auftretend), in denen Exploration die höchste Reward-Varianz und Diversität liefert.

2.2 Dual-Level Adaptive Strategie

AEGPO nutzt diese Signale, um die Berechnung neu zu allokieren, ohne das gesamte Rollout-Budget zu erhöhen:

1. Globale Adaptive Allokation:

   • Der Batch an Prompts wird basierend auf dem Median der $\Delta$Entropy in „hochwertige“ und „niederwertige“ Tiers unterteilt.
   • Hochwertige Prompts (hohe $\Delta$Entropy) erhalten ein größeres Rollout-Budget ($r_{high}$), während niederwertige Prompts weniger Rollouts erhalten ($r_{low}$).
   • Constraint: $r_{high} + r_{low} = 2r_{avg}$, wodurch die Gesamtrechenleistung konstant bleibt.

2. Lokale Adaptive Exploration:

   • Anstatt fester Branching-Schedules identifiziert die Methode dynamisch die Top-$K$ Zeitschritte mit den höchsten Entropy($t$)-Peaks für jeden Rollout.
   • Die Verzweigung (Exploration) erfolgt ausschließlich an diesen identifizierten High-Dispersion-Zeitschritten.
   • Dies stellt sicher, dass die Exploration auf die informativsten Momente der Denoising-Trajektorie konzentriert wird.

3. Zentrale Beiträge

• Entdeckung des Dual-Signals: Die Arbeit etabliert die Attention-Entropie als duales intrinsisches Signal-Proxy, wobei die relative Änderung ($\Delta$Entropy) den Lernwert der Stichprobe angibt und absolute Peaks (Entropy($t$)) kritische Explorations-Zeitschritte identifizieren.
• AEGPO-Framework: Einführung eines dual-level adaptiven Frameworks, das die Stichprobenselektion und die Zeitschritt-Exploration mittels leichtgewichtiger, entropie-gesteuerter Mechanismen verbessert.
• Generalisierbarkeit: Nachweis, dass AEGPO modell- und framework-agnostisch ist und verschiedene Policy-Optimierungsmethoden (DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT) sowie Basismodelle (FLUX.1-dev, SD3.5-M) erfolgreich verbessert.

4. Experimentelle Ergebnisse

Die Experimente wurden auf den Datensätzen HPDv2.1 und Pick-a-Pic unter Verwendung von FLUX.1-dev und SD3.5-M als Basismodellen durchgeführt.

• Konvergenzgeschwindigkeit: AEGPO beschleunigt die Policy-Optimierung signifikant. Es erreicht eine 2-mal schnellere Konvergenz bei DanceGRPO und eine 5-mal schnellere Konvergenz bei DiffusionNFT im Vergleich zu Standardvarianten und erzielt in weniger Schritten überlegene finale Rewards.
• Leistungsmetriken:
  • Bei FLUX.1-dev erreichte BranchGRPO + AEGPO einen neuen State-of-the-Art HPS-v2.1 Score von 0,374 (gegenüber 0,363) und einen Pick Score von 0,232 (gegenüber 0,229).
  • Bei SD3.5-M verbesserte DiffusionNFT + AEGPO den HPS-v2.1 auf 0,343 und den ImageReward auf 1,57.
• Reward-KL-Trade-off: AEGPO-Modelle dominieren die Baseline auf der Reward-KL-Pareto-Front konsequent und erzielen signifikant höhere Alignment-Rewards für jedes gegebene Niveau der Policy-Abweichung (KL-Divergenz).
• Ablationsstudien: Sowohl die globale Allokation (G) als auch die lokale Exploration (L) bieten unabhängige Vorteile, wobei der kombinierte Ansatz (G+L) die höchste Leistung liefert.
• Rechenaufwand: Die Methode führt einen moderaten Overhead ein (ca. 11,1 % Steigerung der Schrittzeit und vernachlässigbare VRAM-Zunahme) durch die Extraktion intermediärer Attention-Maps, was durch die dramatische Reduktion der benötigten Gesamttrainingsschritte zur Erreichung der Spitzenleistung gerechtfertigt ist.

5. Bedeutung und Ansprüche

Das Paper behauptet, dass AEGPO die Kernineffizienzen der statischen Stichprobenentnahme bei Diffusionsmodellen adressiert. Durch den Wechsel von heuristischer Allokation zu intrinsisch signalbasierter Adaptation stellt die Methode sicher, dass Rechenressourcen auf die informativsten Stichproben und die kritischsten Momente des Generierungsprozesses konzentriert werden.

Die Autoren betonen, dass AEGPO ein generalisierbares, Plug-and-Play-Modul ist, das keine neuen Verlustfunktionen oder komplexen Reward-Modellierungen erfordert. Stattdessen nutzt es die internen Attention-Dynamiken des Modells selbst, um die Optimierung zu leiten, und bietet so einen effizienteren Weg zu hochwertigem Alignment über diverse Diffusionsarchitekturen und RLHF-Frameworks hinweg. Die Arbeit legt nahe, dass intrinsische Signale wie die Attention-Entropie untergenutzte Ressourcen zur Verbesserung der Effizienz der Policy-Optimierung in generativen Modellen sind.