Technische Samenvatting: AEGPO – Adaptive Entropy-Guided Policy Optimization voor Diffusiemodellen

1. Probleemstelling

Reinforcement Learning from Human Feedback (RLHF), in het bijzonder Group Relative Policy Optimization (GRPO), is een standaardparadigma geworden voor het afstemmen (alignment) van diffusie- en flow-matchingmodellen op menselijke voorkeuren. Bestaande GRPO-varianten lijden echter aan twee fundamentele inefficiënties door hun afhankelijkheid van statische heuristieken:

1. Inefficiënte Steekproefallocatie: Standaardmethoden wijzen een uniform aantal rollouts toe aan elke prompt, waarbij de aanzienlijke variatie in "leerwaarde" tussen verschillende steekproeven wordt genegeerd. Sommige prompts induceren minimale beleidswijzigingen, terwijl andere substantiële verbeteringen aansturen.
2. Statische Timestep-exploratie: Exploratiestrategieën (bijv. branching in boomstructuur-gebaseerde rollouts) volgen doorgaans vaste, vooraf gedefinieerde schema's. Deze slagen er niet in om te adapteren aan de dynamische aard van het denoising-proces, waarbij kritieke momenten voor exploratie sterk variëren tussen prompts en trainingsfasen.

Deze beperkingen resulteren in verspilde computationele middelen op laagwaardige steekproeven en suboptimale exploratie op niet-kritieke tijdstappen, wat de convergentie vertraagt en de uiteindelijke alignment-prestaties beperkt.

2. Methodologie: AEGPO

De auteurs stellen Adaptive Entropy-Guided Policy Optimization (AEGPO) voor, een plug-and-play framework dat gebruikmaakt van intrinsieke aandachtssignalen om zowel de steekproefselectie als de exploratie dynamisch te sturen. De kerninzicht is dat Attention Entropy dient als een krachtige duale signaalproxy.

2.1 Kernsignalen

De methode definieert twee intrinsieke metrieken afgeleid van de aandachtskaarten (attention maps) van het model:

• $\Delta$Entropy (Steekproefniveau-signaal): De relatieve verandering in aandachtverspreiding tussen het huidige beleid en het basisbeleid.
  • Definitie: $\Delta\text{Entropy} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropy}_{\theta}(t) - \text{Entropy}_{\theta_{base}}(t)|$.
  • Functie: Het fungeert als een robuuste proxy voor de leerwaarde van een steekproef. Een hoge $\Delta$Entropy geeft aan dat een prompt het model dwingt om zijn aandachtstrategie significant te reorganiseren, wat duidt op een hoog leerpotentieel.
• Entropy($t$) (Timestep-niveau-signaal): De absolute aandachtverspreiding op een specifieke denoising-stap $t$.
  • Definitie: De gemiddelde Shannon-entropie van de aandachtgewichten over teksttokens voor beeldkenmerken op stap $t$.
  • Functie: Pieken in Entropy($t$) identificeren kritieke exploratie-tijdstappen. Deze pieken komen overeen met momenten van hoge onzekerheid of brede semantische integratie (vaak voorkomend in vroege structurele en late verfijningsstadia), waar exploratie de hoogste beloningsvariantie en diversiteit oplevert.

2.2 Duaal Adaptief Strategie

AEGPO gebruikt deze signalen om de berekening te herverdelen zonder het totale rollout-budget te verhogen:

1. Globale Adaptieve Allocatie:

   • De batch met prompts wordt gesplitst in "hoogwaardige" en "laagwaardige" segmenten op basis van de mediaan van de $\Delta$Entropy.
   • Hoogwaardige prompts (hoge $\Delta$Entropy) krijgen een groter rollout-budget ($r_{high}$), terwijl laagwaardige prompts minder rollouts ontvangen ($r_{low}$).
   • Restrictie: $r_{high} + r_{low} = 2r_{avg}$, waardoor de totale rekenkracht constant blijft.

2. Lokale Adaptieve Exploratie:

   • In plaats van vaste branching-schema's, identificeert de methode dynamisch de top-$K$ tijdstappen met de hoogste Entropy($t$)-pieken voor elke rollout.
   • Branching (exploratie) vindt uitsluitend plaats op deze geïdentificeerde tijdstappen met hoge dispersie.
   • Dit zorgt ervoor dat exploratie geconcentreerd is op de meest informatieve momenten van het denoising-traject.

3. Belangrijkste Bijdragen

• Ontdekking van Duale Signalen: Het paper vestigt Attention Entropy als een duale intrinsieke proxy-signaal, waarbij relatieve verandering ($\Delta$Entropy) de leerwaarde van de steekproef aangeeft en absolute pieken (Entropy($t$)) de kritieke exploratie-tijdstappen identificeren.
• AEGPO Framework: Introductie van een duaal adaptief framework dat de steekproefselectie en timestep-exploratie verbetert met behulp van lichtgewicht, entropie-gestuurde mechanismen.
• Generaliseerbaarheid: Demonstratie dat AEGPO model-agnostisch en framework-agnostisch is, en succesvol diverse policy-optimalisatiemethoden (DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT) en basismodellen (FLUX.1-dev, SD3.5-M) verbetert.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op de HPDv2.1 en Pick-a-Pic datasets met behulp van FLUX.1-dev en SD3.5-M als basismodellen.

• Convergentiesnelheid: AEGPO versnelt de policy-optimalisatie aanzienlijk. Het bereikt een 2$\times$ snellere convergentie op DanceGRPO en een 5$\times$ snellere convergentie op DiffusionNFT vergeleken met standaardvarianten, waarbij het in minder stappen een superieure uiteindelijke beloning bereikt.
• Prestatie-metrieken:
  • Op FLUX.1-dev bereikte BranchGRPO + AEGPO een nieuwe state-of-the-art HPS-v2.1 score van 0.374 (omhoog van 0.363) en een Pick Score van 0.232 (omhoog van 0.229).
  • Op SD3.5-M verbeterde DiffusionNFT + AEGPO de HPS-v2.1 naar 0.343 en de ImageReward naar 1.57.
• Reward-KL Trade-off: AEGPO-modellen domineren consistent de baseline op de Reward-KL Pareto-front, waarbij ze aanzienlijk hogere alignment-beloningen behalen voor elk gegeven niveau van beleidsafwijking (KL-divergentie).
• Ablatie-studies: Zowel Globale Allocatie (G) als Lokale Exploratie (L) bieden onafhankelijke voordelen, waarbij de gecombineerde aanpak (G+L) de hoogste prestaties levert.
• Computationele Overhead: De methode introduceert een bescheiden overhead (ongeveer 11,1% toename in staptijd en een verwaarloosbare toename in VRAM) door het extraheren van tussenliggende aandachtskaarten, wat gerechtvaardigd wordt door de drastische vermindering van het totaal aantal benodigde trainingsstappen om piekprestaties te bereiken.

5. Betekenis en Claims

Het paper claimt dat AEGPO de kerninefficiënties van statische sampling in diffusiemodellen aanpakt. Door over te stappen van heuristiek-gebaseerde allocatie naar intrinsiek signaal-gebaseerde adaptatie, zorgt de methode ervoor dat computationele middelen worden geconcentreerd op de meest informatieve steekproeven en de meest kritieke momenten van het generatieproces.

De auteurs benadrukken dat AEGPO een generaliseerbare, plug-and-play module is die geen nieuwe verliesfuncties of complexe beloningsmodellering vereist. In plaats daarvan maakt het gebruik van de eigen interne aandachtsdynamiek van het model om de optimalisatie te sturen, wat een efficiënter pad biedt naar hoogwaardige alignment voor diverse diffusiearchitecturen en RLHF-frameworks. Het werk suggereert dat intrinsieke signalen zoals attention entropy onderbenutte bronnen zijn voor het verbeteren van de efficiëntie van policy-optimalisatie in generatieve modellen.