Technische Samenvatting: Expliciete Critic-Geleiding voor het Aligneren van Diffusiemodellen

Probleemstelling

Het post-trainen van diffusiemodellen om te aligneren met niet-differentieerbare doelstellingen (bijvoorbeeld menselijke voorkeur, compositiële correctheid, tekstweergave) wordt steeds belangrijker, aangezien pretraining alleen vaak suboptimale resultaten oplevert op gebruikersgerichte eigenschappen. Een centrale uitdaging op dit gebied is krediettoewijzing. Diffusiemodellen genereren samples via multi-stap denoising-trajecten, terwijl beloningen doorgaans alleen worden gedefinieerd op de uiteindelijk gedecodeerde afbeelding. Bestaande oplossingen ondervinden aanzienlijke beperkingen:

• Groepsrelatieve Methoden (bijv. GRPO): Deze omzeilen expliciete critics door samples binnen een groep met elkaar te vergelijken. Hoewel ze het trainen van een waardenfunctie vermijden, is het resulterende beloningssignaal grof en op trajectniveau, zonder fijne krediettoewijzing voor intermediaire denoising-beslissingen.
• Critics in Pixelspace: Bestaande expliciete actor-critic-methoden definiëren critics vaak in pixelspace, waarbij gereconstrueerde "schone" afbeeldingen worden geëvalueerd vanuit intermediaire ruwe latents. Dit creëert een toestandsmismatch: het beleid opereert op ruwe latent-toestanden, terwijl de critic proxy-toestanden evalueert die vaak buiten de verdeling liggen voor standaard visuele encoders (bijv. CLIP, BLIP). Deze mismatch leidt tot breekbare waardschatting, instabiliteit in optimalisatie en aanzienlijke decoding- overhead.
• Reward Hacking: Het optimaliseren voor specifieke verifieerbare beloningen leidt vaak tot shortcut-oplossingen (bijv. vereenvoudigde achtergronden, repetitieve patronen) die de doelmetriek maximaliseren terwijl ze niet-gemeten aspecten van de beeldkwaliteit verslechteren.

Methodologie

De auteurs stellen een State-Aligned Latent Actor-Critic Framework voor dat deze problemen aanpakt door de critic direct te integreren in het diffusiegeneratieproces.

1. State-Aligned Latent Critic

In plaats van te opereren in pixelspace, leert het framework een waardenfunctie $V_\phi(z_t, t, y)$ direct op de ruwe latent-toestanden $z_t$ die door het beleid worden bezocht.

• Architectuur: De critic wordt geïnitieerd vanuit de voorgeprende diffusiebackbone (UNet of DiT) en gemodificeerd met lichtgewicht architecturale wijzigingen om een scalair waarde te voorspellen.
  • UNet: De feature map wordt gepatchified, en een leerbaar samenvattingstoken aggregateert informatie via een attention-head en een MLP.
  • DiT: Lichtgewicht critic-branches worden ingevoegd op intermediaire lagen (bijv. lagen 7, 15, 23) met behulp van leerbare query-tokens om visuele informatie te aggregeren.
• State Alignment: Door te opereren in dezelfde latentruimte en de generatieve prior te delen, is de critic van nature bewust van de tijdstap en het ruisniveau, waardoor distributieverplaatsingsproblemen van critics in pixelspace worden vermeden.

2. Stabiliserende Trainingstechnieken

Om stabiele actor-critic-optimalisatie te waarborgen, worden twee specifieke technieken geïntroduceerd:

• Tijdstap-bewuste Conditionering: De value-head wordt expliciet geconditioneerd op de tijdstap $t$ (via AdaLN-stijl modulatie), wat directe toegang biedt tot het ruisniveau om de waardevoorspelling over het traject te verbeteren.
• Value Pretraining: Voordat de gezamenlijke actor-critic-optimalisatie plaatsvindt, wordt een korte warm-up-fase uitgevoerd waarbij alleen de critic wordt bijgewerkt met regressieverlies tegen Monte Carlo-return-doelen. Dit levert betere initiële waardschattingen op, wat de daaropvolgende PPO-updates stabiliseert.

3. Inference-Time Steering

Omdat de critic is gedefinieerd op dezelfde latent-toestanden als het beleid, kan deze tijdens de inferentie worden hergebruikt als een geleidingsmodule. De gradiënt van de waardenfunctie, $\nabla_{z_t} V_\phi$, wijst naar latent-regio's waarvan wordt voorspeld dat ze hogere terminale beloningen opleveren. Deze gradiënt wordt geïnjecteerd in de sampler-update (modificatie van ruisvoorspelling voor DDIM of snelheid voor flow-matching), waardoor het denoising-traject effectief wordt gestuurd zonder extra training.

4. Multi-Reward Optimalisatie

Om reward hacking te mitigeren, wordt het framework uitgebreid tot multi-reward optimalisatie. Een gedeelde latent-backbone wordt gebruikt met aparte, lichtgewicht value-heads voor elk reward-type. Dit ontwerp stelt heterogene beloningen (bijv. compositieel versus esthetisch) in staat om elkaar te regulariseren terwijl ze dure feature-extractie delen, waardoor het instorten in shortcut-oplossingen die gebruikelijk zijn bij single-reward training, wordt voorkomen.

Belangrijkste Bijdragen

1. State-Aligned Latent Critic: Een framework dat een voorgeprent diffusiemodel omzet in zijn eigen op tijdstap geconditioneerde waardenetwerk, waardoor waardevoorspelling op ruwe latents en inference-time steering mogelijk wordt.
2. Stabilisatiestrategieën: Identificatie van tijdstap-bewuste conditionering en value pretraining als essentiële technieken voor het stabiliseren van expliciete actor-critic diffusietraining.
3. Multi-Reward Uitbreiding: Een mechanisme dat aparte value-heads op een gedeelde backbone gebruikt om complementaire beloningen gezamenlijk te optimaliseren, waardoor proxy-overoptimalisatie en reward hacking worden mitigeerd.
4. Empirische Superioriteit: Consistente prestatieverbetering ten opzichte van groepsrelatieve (GRPO) en eerdere actor-critic baselines op zowel single-reward als multi-reward benchmarks, voor zowel UNet- als DiT-backbones.

Experimentele Resultaten

De methode werd geëvalueerd op Stable Diffusion 1.5 (UNet) en Stable Diffusion 3.5M (DiT) over diverse reward-types:

• Single-Reward Benchmarks: De methode behaalde state-of-the-art resultaten op CLIP, HPSv2.1, PickScore, GenEval en OCR-metrieken. Opmerkelijk is dat het aanzienlijke winsten liet zien op verifieerbare beloningen (GenEval, OCR), waar schaarse en discrete signalen het meest profiteren van per-tijdstap waarde-schattingen. Het toonde ook stabiliteit waar DDPO faalde om te convergeren op SD3.5M.
• Multi-Reward Benchmarks: Bij gezamenlijke optimalisatietaken (bijv. CLIP + HPS + GenEval) presteerde de methode beter dan GRPO, met name op GenEval, terwijl het vergelijkbare prestaties behield op voorkeursmetrieken.
• Inference-Time Steering: Het toepassen van de geleerde critic als geleidingssignaal tijdens het sampling verbeterde de generatiekwaliteit verder. Op SD1.5 overtrof critic-geleiding alleen Best-of-N (BoN) sampling; op SD3.5M leverde het combineren van geleiding met BoN de beste resultaten op.
• Efficiëntie: De latent-critic was aanzienlijk sneller dan critics in pixelspace (CLIP/BLIP/HPS) bij zowel waardeberekening (10s versus 28–35s) als trainingsstappen (20s versus 52–86s).
• Ablaties:
  • Het verwijderen van value pretraining verslechterde de prestaties, vooral op schaarse beloningen.
  • Afzonderlijke value-heads voor multi-reward-taken presteerden beter dan gedeelde-head-ontwerpen.
  • Training zonder Classifier-Free Guidance (CFG) leverde de meest stabiele optimalisatie en de hoogste uiteindelijke beloningen op, waardoor de instabiliteit en artefacten die geassocieerd zijn met hoge CFG-schalen, werden vermeden.

Betekenis en Claims

Het artikel claimt een praktisch recept te hebben vastgesteld voor het toepassen van reinforcement learning op moderne diffusiemodellen. De betekenis hiervan ligt in:

1. Oplossen van de Toestandsmismatch: Door de critic af te stemmen op de latent-toestandsruimte van het beleid, worden de breekbaarheid en overhead van evaluatie in pixelspace geëlimineerd.
2. Mogelijk maken van Stabiele Actor-Critic Training: De combinatie van latent-alignement, tijdstap-conditionering en value pretraining maakt expliciete actor-critic training haalbaar en stabiel voor diffusiemodellen, en biedt een meer sample-efficiënt alternatief voor groepsrelatieve methoden.
3. Dubbel Nut: De geleerde critic dient een dubbel doel: het optimaliseren van het beleid tijdens training en het sturen van het generatieproces tijdens inferentie, wat prestatiewinst biedt die verder gaat dan de trainingsfase.
4. Mitigeren van Reward Hacking: De multi-reward uitbreiding biedt een praktisch mechanisme om concurrerende doelstellingen in evenwicht te brengen, waardoor de prevalentie van shortcut-oplossingen in diffusie-RL wordt verminderd.

De auteurs erkennen beperkingen, met name de huidige focus op beelddiffusie (niet video), mogelijke efficiëntievermindering in multi-reward settings vergeleken met single-reward training, en discretisatiefouten die inherent zijn aan SDE-sampling.