Adaptive Moments are Surprisingly Effective for Plug-and-Play Diffusion Sampling

Probleemstelling

Diffusiemodellen zijn uitgegroeid tot een toonaangevende techniek voor generatieve modellering. Een belangrijke toepassing is Plug-and-Play (PnP) conditionele generatie, waarbij een model dat is getraind op de marginale verdeling $p(x)$ wordt gebruikt om te bemonsteren uit een conditionele verdeling $p(x|y)$ zonder specifieke training voor de taak.

De kernuitdaging bij PnP-sampling ligt in het benaderen van de likelihood-score ($\nabla_{x_t} \log p(y|x_t)$). Volgens Bayes' regel is de posterior-score de som van de prior-score (geleverd door het diffusiemodel) en de likelihood-score. Omdat de likelihood-score vaak niet direct berekenbaar is (intractable), moeten bestaande methoden zoals Diffusion Posterior Sampling (DPS) en Classifier Guidance (CG) deze benaderen.

Deze benaderingen introduceren echter aanzienlijke ruis in de sampling-dynamica. De gradiënten die worden gebruikt om de sampling te sturen, zijn onstabiel en kunnen leiden tot artefacten, wazige resultaten of een mislukte alignering met de gewenste conditie (bijvoorbeeld bij zware super-resolutie of inpainting). Bestaande geavanceerde methoden proberen deze ruis te verminderen door complexere algoritmen te ontwikkelen (zoals Monte Carlo benaderingen of iteratieve verfijning), maar deze zijn vaak computationeel duur en verliezen hun effectiviteit naarmate de taak moeilijker wordt (bijv. minder conditionele informatie).

Methodologie: Adaptive Moment Estimation

De auteurs stellen een eenvoudige maar krachtige oplossing voor: het toepassen van Adaptive Moment Estimation (bekend van de Adam-optimizer in stochastische optimalisatie) om de ruis in de likelihood-gradiënten te stabiliseren tijdens het sampling-proces.

In plaats van complexere benaderingsmethoden te ontwikkelen, behouden ze exponentiële voortschrijdende gemiddelden van de eerste en tweede momenten van de likelihood-score schattingen over de sampling-stappen heen:

1. Eerste moment (Momentum): Accumuleert gradiëntinformatie over stappen om de optimalisatietraject te gladstrijken.
   $$m_k = \beta_1 m_{k-1} + (1 - \beta_1) g_t$$
2. Tweede moment (Adaptieve schaling): Past de updates aan op basis van de historische variantie van elke gradiëntcomponent.
   $$v_k = \beta_2 v_{k-1} + (1 - \beta_2) g_t^2$$

De gestabiliseerde likelihood-score $\hat{g}_t$ wordt vervolgens berekend als:
$$\hat{g}_t = \frac{\hat{m}_k}{\sqrt{\hat{v}_k} + \delta}$$
waarbij $\hat{m}_k$ en $\hat{v}_k$ de bias-correcteerde momenten zijn en $\delta$ een kleine constante voor numerieke stabiliteit.

De auteurs passen deze techniek toe op twee bestaande methoden:

• AdamDPS: Adaptive Moments toegepast op Diffusion Posterior Sampling.
• AdamCG: Adaptive Moments toegepast op Classifier Guidance.

Dit resulteert in een "plug-and-play" modificatie die slechts enkele regels code vereist in bestaande implementaties.

Belangrijkste Bijdragen

1. Prestatieverbetering met minimale complexiteit: Het paper demonstreert dat adaptive moment estimation de prestaties van PnP-methoden substantieel verbetert, wat leidt tot state-of-the-art resultaten op zowel reconstructie- als class-conditional generatietaken, zonder significante extra rekenkosten.
2. Empirische analyse van ruisreductie: De auteurs leveren uitgebreide analyses (zowel synthetisch als op real data) die aantonen dat het mitigeren van gradiëntruis via adaptive moments de alignering met de conditie verbetert. Ze tonen aan dat DPS-gradiënten vaak in tegenstrijdige richtingen wijzen (negatieve cosine-similariteit tussen opeenvolgende stappen), terwijl AdamDPS consistente, coherente updates biedt.
3. Inzicht in taakmoeilijkheid: Het paper onthult dat taakmoeilijkheid (bijv. mate van degradatie van de input) de relatieve prestaties van methoden sterk beïnvloedt. Complexe methoden zoals Training-Free Guidance (TFG) presteren goed op makkelijke taken, maar degraderen snel bij moeilijke taken, terwijl AdamDPS robuust blijft.

Resultaten

De methoden zijn getest op diverse datasets (ImageNet, CIFAR-10, Cats vs. Dogs) en taken:

• Reconstructie-taken (Super-resolutie, Deblurring, Inpainting):
  • AdamDPS overtreft alle geëvalueerde methoden (waaronder DPS, TFG, UGD, MPGD) op alle reconstructie-taken, inclusief extreme scenario's zoals 16x super-resolutie en 90% inpainting.
  • Kwalitatief produceert AdamDPS scherpere details en minder artefacten dan DPS of TFG.
  • Quantitatief (gemeten in LPIPS en FID) behaalt AdamDPS de beste scores.
• Class-conditional Sampling:
  • Op CIFAR-10 verbetert AdamDPS de classificatie-accuraatte met 9,86% ten opzichte van de volgende beste methode (DPS).
  • Op ImageNet behaalt AdamDPS 10,49% top-10 nauwkeurigheid, terwijl andere methoden dicht bij willekeurige kans (1%) blijven.
  • Voor time-aware classifiers (AdamCG) wordt de nauwkeurigheid met meer dan 19 procentpunten verbeterd ten opzichte van standaard CG.
• Robuustheid bij toenemende moeilijkheid:
  • Naarmate de taak moeilijker wordt (minder conditionele informatie), presteren geavanceerde methoden zoals TFG slechter dan de simpele DPS-baseline. AdamDPS blijft daarentegen consistent beter presteren dan DPS, ongeacht de moeilijkheidsgraad.
• Efficiëntie:
  • De methode voegt verwaarloosbare overhead toe aan de wall-clock tijd in vergelijking met DPS. TFG is daarentegen aanzienlijk trager vanwege de herhaalde gradiëntberekeningen.

Significantie en Conclusie

Deze paper biedt een fundamenteel nieuw perspectief op het verbeteren van Plug-and-Play diffusie-sampling. In plaats van te focussen op het verfijnen van de likelihood-score-benadering zelf (wat vaak leidt tot steeds complexere algoritmen), focust de auteurs op het stabiliseren van de updates die voortvloeien uit deze benaderingen.

De conclusie is dat ruismitigatie via adaptive moments een eenvoudige, maar uiterst effectieve strategie is. Dit is vooral waardevol voor wetenschappelijke domeinen waar inverse problemen vaak onderbepaald zijn (onderdetermined) en de conditionele informatie schaars en ruisig is. De methode biedt een praktische, goedkope en robuuste alternatief voor de trend van steeds complexere likelihood-approximaties, en kan breed worden toegepast op bestaande PnP-frameworks.