Technische Samenvatting: Leren Steproeven van Diffusiemodellen via Inverse Reinforcement Learning

Probleemstelling

Diffusiemodellen zijn een standaard geworden voor generatie van hoogdimensionele data, maar het stepprocedere blijft een kritieke bottleneck. Hoewel het denoisings-neuraal netwerk doorgaans vaststaat na de training, vereist het steleuralgoritme zelf zorgvuldige afstemming van hyperparameters (bijv. ruisschema's, guidance-schalen, stochasticiteitsprofielen). Momenteel wordt deze afstemming uitgevoerd via kostbare empirische grid-searches. Voor complexe modellen zoals die getraind op ImageNet-64, kan de computationele kosten van een uitputtende hyperparameterzoektocht de kosten van het trainen van de denoiser zelf benaderen of zelfs overstijgen. Bovendien vereisen bestaande reinforcement learning (RL) of controle-benaderingen voor diffusiemodellen vaak expliciete beloningsfuncties of het finetunen van de denoiser, wat in strijd is met het doel om de sampler te optimaliseren zonder de vooraf getrainde generator aan te passen.

Methodologie

De auteurs stellen voor om diffusie-sampling te formuleren als een sequentieel beslissingsprobleem dat oplosbaar is via Inverse Reinforcement Learning (IRL). De kerninnovatie is het leren van een sampling-policy die doel-toestandsdistributies benadert zonder toegang tot expert-trajecten, expliciete beloningen of het hertrainen van de denoiser.

1. MDP-formulering

Het diffusie-samplingproces wordt geformaliseerd als een Markov Beslissingsproces (MDP) met een eindig horizontont:

• Toestand ($s_t$): Gedefinieerd als $(x_t, \sigma_t)$, waarbij $x_t$ de huidige sample is en $\sigma_t$ het ruisniveau.
• Actie ($a_t$): Discrete controlebeslissingen toegepast op de sampling-dynamiek, zoals het selecteren van guidance-schalen, niveaus van ruisinjectie, of het besluiten om opnieuw te ruisen (herstarten) naar een eerder ruisniveau.
• Dynamiek: De transitie-kernel wordt geïnduceerd door de gediscretiseerde diffusie-dynamiek en de geselecteerde controle-actie. Het proces eindigt wanneer het ruisniveau $\Sigma_0$ bereikt (een absorberende toestand).

2. State-Only Inverse Reinforcement Learning

In tegen tegenstelling tot traditionele IRL, die vaak uitgaat van expert-acties of trajecten, opereert deze methode onder een state-only regime. De "expert" wordt uitsluitend gedefinieerd door de marginale distributies van dataprojecten bij geselecteerde ruisniveaus.

• Doel: Het doel is om een policy $\pi_\theta$ te vinden waarvan de state occupancy measure $\mu_\theta$ (de marginale distributie van bezochte toestanden) de expert occupancy measure $\mu_E$ benadert.
• Optimalisatie: In plaats van een beloningsfunctie te leren, minimaliseren de auteurs direct de $f$-divergentie tussen de expert en de policy occupancy measures:
  $$\min_\theta D_f(\mu_E \parallel \mu_\theta)$$
  Het paper onderzoekt specifiek de Kullback-Leibler (KL) en Reverse-KL (rKL) divergenties, waarbij wordt opgemerkt dat de keuze van de $f$-divergentie de afweging bepaalt tussen mode-covering (diversiteit) en mode-seeking (getrouwheid/fidelity).

3. Policy Gradient en Optimalisatie

De auteurs leiden een policy gradient theorem af voor het direct minimaliseren van $f$-divergenties. De gradiënt wordt uitgedrukt als:
$$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$$
waarbij $h_f$ een functie is afgeleid van de generator van de $f$-divergentie.

• Density Ratio Estimation: Omdat $\mu_E$ en $\mu_\theta$ onbekend zijn, wordt de dichtheidsratio geschat met behulp van een discriminator die getraind is om onderscheid te maken tussen expert-toestanden (geruisde data) en door de policy gegenereerde toestanden.
• Stabilisatie: Om de hoge variantie in gradiënt-schattingen te beheersen, gebruikt de methode een PPO-achtige clipped objective, importance sampling correcties en een globale batch mean baseline.

4. Geëvalueerde Samplingstrategieën

Het framework wordt toegepast op drie specifieke interventies tijdens de inferentie:

1. Adaptieve Classifier-Free Guidance: Het leren van een toestand-afhankelijke guidance-schaal $\omega(x, \sigma)$.
2. Stochasticiteit-injectie: Het leren van een toestand-afhankelijk niveau van ruisinjectie $\gamma_{EDM}(x, \sigma)$.
3. Renoising (Restart Sampling): Het leren wanneer het proces moet worden teruggezet naar een eerder ruisniveau om slechte lokale generaties te vermijden.

Belangrijkste Bijdragen

1. MDP-formulering voor Sampling: Het paper beschouwt diffusie-sampling als een MDP om samplers tijdens de inferentie aan te passen zonder de denoiser te hertrainen.
2. Directe State-Marginal Matching: Het leidt een policy gradient objectief af dat $f$-divergenties tussen state occupancy measures minimaliseert, waarmee de noodzaak voor expliciete beloningsmodellering of expert-actiegegevens wordt omzeild.
3. Interpreteerbare Afwegingen: Het framework legt een controleerbare afweging bloot tussen sample-diversiteit (Recall) en getrouwheid (Precision) via de keuze van de $f$-divergentie (KL vs. rKL).
4. Toestand-afhankelijke Controle: De geleerde policies fungeren als controlemechanismen tijdens de inferentie die selectief interveniëren (bijv. het toepassen van hoge guidance alleen bij specifieke ruisniveaus of voor specifieke typen samples) in plaats van vaste globale schema's toe te passen.

Experimentele Resultaten

De methode is geëvalueerd op CIFAR-10, FFHQ en ImageNet-64 met behulp van vooraf getrainde EDM-modellen.

• Prestaties: De geleerde policies evenaren of overtreffen de prestaties van de best handmatig afgestemde baselines (gevonden via grid search) op basis van FID, Precision en Recall metrieken.
  • Adaptieve Guidance: Verbeterde de kwaliteit-diversiteit afweging vergeleken met vaste guidance. KL-divergentie verbeterde de Recall (diversiteit), terwijl rKL de Precision (getrouwheid) verbeterde.
  • Stochasticiteit-injectie: De geleerde toestand-afhankelijke profielen herstelden de kwalitatieve patronen van optimale EDM-schema's (lage stochasticiteit aan het einde van het traject) en presteerden beter dan vaste baselines op complexe datasets zoals ImageNet.
  • Renoising: Verminderde het aantal functie-evaluaties (NFE) aanzienlijk op ImageNet, terwijl de FID werd behouden of verbeterd vergeleken met standaard restart-strategieën.
• Computationele Efficiëntie:
  • Training: Een enkele training run vervangt een uitputtende grid search. Op ImageNet-64 reduceert dit het optimalisatiebudget met factor 9$\times$ (van ~1437 EFLOPs voor een 3-hyperparameter grid search naar ~161 EFLOPs voor de IRL-training).
  • Inference: De policy voegt een bescheiden overhead toe van 16% tot 40% in FLOPs vergeleken met de denoiser alleen, wat als verwaarloosbaar wordt beschouwd gezien het elimineren van handmatige afstemming.
• Ablatie-bevindingen:
  • Conditionering: Policies die geconditioneerd zijn op de volledige toestand $(x, \sigma)$ presteren significant beter dan die enkel geconditioneerd zijn op het ruisniveau $\sigma$, en presteren bijna even goed als die met expliciete klasse-labels $(x, \sigma, c)$, wat suggereert dat de policy impliciet klasse-onderscheidende kenmerken leert van de beeldinhoud.
  • Keuze van Divergentie: Er is aangetoond dat KL en rKL het meest effectief zijn voor het controleren van de precision-recall afweging, terwijl symmetrische divergenties (TV, Hellinger) minder interpreteerbare resultaten produceerden.

Betekenis en Claims

Het paper claimt dat dit framework een systematisch, geautomatiseerd alternatief biedt voor handmatige hyperparameter-afstemming voor diffusie-sampling. Door sampling te behandelen als een IRL-probleem, maakt het de optimalisatie van complexe, toestand-afhankelijke samplingstrategieën mogelijk die moeilijk handmatig te ontwerpen zijn.

De auteurs benadrukken dat de methode bijzonder waardevol is voor grootschalige modellen waar de kosten van evaluatie per trial in een grid search prohibitief zijn. De aanpak wijzigt het generatieve model zelf niet, maar leert een optimale "controle-policy" voor het inferentie-traject, wat een praktische weg biedt om de kwaliteit en diversiteit van samples te verbeteren zonder de computationele last van hertraining of uitputtende zoektochten. Het paper blijft bescheiden en merkt op dat de voordelen dataset-afhankelijk zijn en dat het huidige framework steunt op discrete actieruimtes, wat suggereert dat toekomstig werk zich moet richten op continue actieruimtes en robuustere initialisatie.