Technische Samenvatting: Uitlijnen van Few-Step Generatieve Modellen door Amortisatie van Op Steekproeven Gebaseerde Variatie-Inferentie

Probleemstelling

Het uitlijnen van few-step generatieve modellen (bijv. GANs, consistentiemodellen, flow maps, driftmodellen) met specifieke doelen, zoals beloningsfuncties of menselijke voorkeuren, vormt een aanzienlijke uitdaging. Bestaande uitlijningskaders vertrouwen doorgaans op restrictieve structurele aannames die few-step-generators vaak schenden:

1. Berekenbare Likelihoods: Veel methoden vereisen expliciete likelihood-evaluatie over denoising-trajecten, wat niet beschikbaar is voor impliciete modellen zoals GANs of one-step-mappings.
2. Specifieke Oplossers: Benaderingen gebaseerd op Stochastische Optimale Controle (SOC) of Versterkingsleer (RL) gaan vaak uit van specifieke oplossers voor Gewone Differentiaalvergelijkingen (ODE) of Stochastische Differentiaalvergelijkingen (SDE) en sequentiële denoising-processen.
3. Modelspecificiteit: Formuleringen zijn vaak op maat gemaakt voor specifieke modelfamilies (bijv. diffusiemodellen), waardoor ze niet generaliseren naar diverse few-step-architecturen zoals VAEs of flow maps.

Bijgevolg ontbreekt er een algemeen uitlijningskader dat diverse few-step-generators kan fine-tunen terwijl hun kernvoordeel behouden blijft: snelle, single-step-inferentie.

Methodologie: FAV

De auteurs stellen FAV (Few-step Generative Models Alignment via Sample-based Variational Inference) voor, een algemeen kader dat uitlijning ontkoppelt van de generatorfamilie en de steekproefdynamiek. FAV formuleert uitlijning als het steekproeven uit een beloningsgeoriënteerde verdeling verankerd in een referentieverdeling:
$$q^*(x) \propto p_{\text{ref}}(x) \exp(\beta \cdot r(x))$$
waarbij $p_{\text{ref}}$ een referentieverdeling is (bijv. een vooraf getrainde generator of empirische data) en $r(x)$ de doelfunctie is.

De methode verloopt in drie belangrijke fasen:

1. Op Steekproeven Gebaseerde Benadering via SVGD

FAV benadert de doelverdeling $q^*$ met Stein Variational Gradient Descent (SVGD). In plaats van dichtheids-evaluatie te vereisen, stelt SVGD de verdeling voor met deeltjes en transporteert deze iteratief naar de doelverdeling. Het optimale Stein-velocityveld $\phi^*$ wordt afgeleid om de afname van de KL-divergentie te maximaliseren:
$$\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \left[ k(x', x) \nabla_{x'} \log p_{\text{ref}}(x') + \beta \cdot k(x', x) \nabla_{x'} r(x') + \nabla_{x'} k(x', x) \right]$$
Dit veld bestaat uit drie componenten:

• Prior-uitlijning: Drijft steekproeven naar de referentieverdeling $p_{\text{ref}}$.
• Beloningsgeleiding: Duwt steekproeven naar gebieden met hoge beloning.
• Diversiteitsafdwinging: Voorkomt mode-collapse via de kernel-gradiëntterm.

2. Non-parametrische Score-schatting (KDE)

Een kritieke uitdaging is dat de score $\nabla \log p_{\text{ref}}$ vaak onberekenbaar is voor empirische data of impliciete generators. FAV lost dit op door de score non-parametrisch te schatten met Kernel Density Estimation (KDE). Door de referentieverdeling te convolueren met een Gaussische kernel, wordt de score benaderd als een gewogen gemiddelde van verschillen tussen referentiesteekproeven en huidige deeltjes, waardoor de noodzaak voor expliciete likelihoods effectief wordt omzeild.

3. Amortisatie via Vastpunt-regressie

Om de snelle inferentiesnelheid van few-step-modellen te behouden, voert FAV geen iteratieve SVGD-bijwerkingen uit tijdens de inferentie. In plaats daarvan amortiseert het het transport in de generatorparameters $\theta$ via vastpunt-regressie. De generator wordt getraind om de afstand tussen zijn output en het getransporteerde doel te minimaliseren:
$$\mathcal{L}(\theta) = \mathbb{E}_{\epsilon} \left[ \| \text{stopgrad}(f_\theta(\epsilon) + \hat{\phi}^*_{q_\theta, q^*_\sigma}(f_\theta(\epsilon))) - f_\theta(\epsilon) \|_2^2 \right]$$
Hier fungeert het getransporteerde steekproef als een stabiel regressiedoel. Dit stelt de generator in staat om in één forward-pass uitgelijnde steekproeven te produceren.

Extensie naar Black-Box-beloningen: Voor niet-differentieerbare beloningen maakt FAV gebruik van nulde-orde gradiënt-schatting (Gaussische smoothing) om $\nabla r(x)$ te benaderen, waardoor uitlijning mogelijk is zelfs wanneer de beloningsfunctie een black box is.

Belangrijkste Bijdragen

1. Algemeen Uitlijningskader: FAV is het eerste kader dat diverse few-step-generators (inclusief VAEs, GANs, consistentiemodellen en flow maps) uitlijnt zonder berekenbare likelihoods, specifieke ODE/SDE-oplossers of multi-step-trajecten te vereisen.
2. Op Steekproeven Gebaseerde VI-formulering: Door uitlijning te gieten als het steekproeven uit een beloningsgeoriënteerde verdeling met SVGD en KDE, ontkoppelt FAV het uitlijningsproces van de interne dynamiek van de generator.
3. Behoud van Efficiëntie: Door amortisatie via vastpunt-regressie behoudt FAV de single-step-inferentie-efficiëntie van de basisgenerator, waardoor de rekenkundige overhead van iteratief steekproeven of distillatieprocessen wordt vermeden.
4. Compatibiliteit met Black Boxes: De integratie van nulde-orde gradiënt-schatting stelt FAV in staat om willekeurige black-box-beloningen te optimaliseren.

Experimentele Resultaten

De auteurs hebben FAV geëvalueerd op twee primaire domeinen: robotica-manipulatie en uitlijning van beeldgeneratoren.

1. Robotica-manipulatie (Offline en Offline-naar-Online RL)

• Benchmarks: 56 offline RL-taken (OGBench en D4RL) en 30 offline-naar-online RL-taken.
• Prestaties: FAV behaalde state-of-the-art prestaties over alle taken, presterend beter dan 10 baselines, waaronder gereparameteriseerde policy gradients, Adjoint Matching en gedistilleerde generatieve policies (bijv. FQL, IFQL).
• Belangrijkste Bevinding: FAV demonstreerde superieure policy-extractiecapaciteiten en overgangsprestaties van offline naar online leren, zelfs zonder aanvullende generatieve modellen of multi-step-denoising.

2. Uitlijning van Beeldgeneratoren

• Schaal: Experimenten varieerden van ImageNet-256 (256x256) tot high-resolution text-to-image-synthese (1024x1024).
• Modellen: Getest op StyleGAN-XL, driftmodellen, iMeanFlow en SANA-Sprint.
• Maten: Geëvalueerd tegen esthetische scores, menselijke voorkeurscores (HPS) en veiligheidsdetectoren (NSFW).
• Prestaties:
  • FAV behaalde consistent de beste Pareto-frontier tussen beloningsoptimalisatie en steekproefdiversiteit/kwaliteit, en beperkte beloningsoveroptimalisatie beter dan baselines zoals DRaFT en Adjoint Matching.
  • Inferentie-efficiëntie: FAV genereerde uitgelijnde steekproeven ongeveer 180x tot 280x sneller dan inferentie-tijd uitlijningsmethoden (bijv. Best-of-N, ReNO) terwijl het concurrerende beeldkwaliteit behield.
  • Schalbaarheid: De methode schaalde succesvol van 256x256 naar 1024x1024 resoluties en generaliseerde over verschillende modelarchitecturen.

Betekenis en Claims

Het artikel stelt dat FAV een universele, model-agnostische oplossing biedt voor het uitlijnen van few-step generatieve modellen. De betekenis hiervan ligt in:

• Verwijdering van Structurele Barrières: Het elimineert de noodzaak voor berekenbare likelihoods of specifieke steekproeftrajecten, waardoor uitlijning toegankelijk wordt voor een bredere klasse impliciete generatieve modellen.
• Balans tussen Efficiëntie en Kwaliteit: Het verzoent succesvol de afweging tussen generatie met hoge beloning en de rekenkundige efficiëntie van few-step-inferentie, een uitdaging waarbij eerdere methoden vaak prestaties degradeerden of complexe distillatie vereisten.
• Robuustheid: Het kader demonstreert robuustheid over diverse taken, van low-dimensional controle tot high-resolution beeldsynthese, en verwerkt zowel differentieerbare als black-box beloningsfuncties.

De auteurs concluderen dat FAV een sterke, algemene aanpak vertegenwoordigt voor policy-extractie in RL en beloningsuitlijning in beeldgeneratie, en een praktische alternatief biedt voor complexere, aannames-zware methoden.