Stochastic Generative Plug-and-Play Priors

Titel: Stochastic Generative Plug-and-Play Priors (SGPnP)

Auteurs: Chicago Y. Park, Edward P. Chandler, Yuyang Hu, Michael T. McCann, Cristina Garcia-Cardona, Brendt Wohlberg, Ulugbek S. Kamilov.

1. Het Probleem

Beeldherstelproblemen (zoals ontvlechting, super-resolutie, inpainting en MRI-reconstructie) zijn vaak "ill-posed" (onderbepaald), wat betekent dat er meerdere mogelijke oplossingen zijn voor een enkele set metingen.

• Plug-and-Play (PnP) priors: Traditionele PnP-methoden lossen deze problemen op door een geleerde denoiser (ruisverwijderaar) te integreren in optimalisatie-algoritmen. Echter, deze methoden hebben moeite met ernstig onderbepaalde problemen.
• De beperkingen:
  1. Distributiemismatch: Tussenliggende iteraties in PnP bevatten gestructureerde artefacten, terwijl de denoisers meestal zijn getraind op additief Gaussisch ruis. Hierdoor wordt de denoiser buiten zijn trainingsdomein gebruikt.
  2. Beperkte ruisdekking: De meeste PnP-denoisers zijn geoptimaliseerd voor lage ruisniveaus en falen bij zwaar beschadigde beelden waar sterke prioren nodig zijn.
• Score-based Diffusion Models (SBDMs): Hoewel SBDMs uitstekende generatieve prestaties leveren door training over een breed scala aan ruisniveaus, is het niet duidelijk hoe ze systematisch als prioren in PnP-algoritmen kunnen worden gebruikt zonder de hele reverse diffusion-sampling te herhalen bij elke stap. Bestaande stochastische PnP-methoden (zoals SNORE) lossen de mismatch slechts gedeeltelijk op en blijven beperkt in sterk ill-posed scenario's.

2. Methodologie

De auteurs introduceren een nieuw raamwerk dat PnP en score-based diffusion modellen combineert via een Score-based Interpretatie en een Stochastische Generatieve PnP (SGPnP) strategie.

A. Score-based Interpretatie van PnP

De auteurs leggen een wiskundig verband tussen klassieke PnP-iteraties en score-based denoising.

• Ze gebruiken de formule van Tweedie om een link te maken tussen de MMSE-denoiser en de scorefunctie (de gradiënt van de log-dichtheid van ruisige data).
• Ze ontwikkelen een methode om vooraf getrainde SBDM's (zowel Variance-Exploding als Variance-Preserving) direct te mappen naar PnP-denoisers.
• Parameter Matching: Omdat SBDM's vaak werken met discrete tijdstappen ($t$) en PnP met continue ruisniveaus ($\sigma$), wordt een interpolatiestrategie ontwikkeld om de juiste tijdstap te vinden die overeenkomt met het gewenste ruisniveau van de PnP-iteratie.

B. Het SGPnP Framework

Het kernidee is het introduceren van een stochastische ruisinjectie vóór de denoising-stap.

• Deterministische PnP (SDPnP): Gebruikt de SBDM direct als denoiser op de tussenliggende iteratie.
• Stochastische Generatieve PnP (SGPnP): Voegt expliciet Gaussische ruis toe aan de invoer van de denoiser voordat deze wordt gedenoist.
  • Formule: $x_{k+1} = D_\theta(z_k + \sigma_{inject} n; \sigma_{cond})$, waarbij $n \sim \mathcal{N}(0, I)$.
  • Decoupling: Een cruciaal verschil met eerdere methoden is dat het injectie-niveau ($\sigma_{inject}$) en het conditioneringsniveau ($\sigma_{cond}$) van elkaar kunnen verschillen. Dit houdt rekening met de extra corruptie (artefacten) die wordt geïntroduceerd door de data-consistentiestap, niet alleen door de injecteerde ruis.

3. Belangrijkste Bijdragen

1. Score-based Interpretatie: Een wiskundige onderbouwing die rechtstreeks het gebruik van vooraf getrainde SBDM's binnen PnP-algoritmen rechtvaardigt zonder reverse diffusion sampling.
2. SGPnP Framework: Een nieuw algoritme dat ruisinjectie gebruikt om:
   • De invoer van de denoiser beter af te stemmen op de trainingsdistributie van het SBDM.
   • Stochasticiteit in te voeren die helpt om te ontsnappen aan strikte zadelpunten (strict saddle points) in de optimalisatie-landschap.
3. Theoretische Garanties:
   • Ontsnapping aan zadelpunten: Bewezen dat onder specifieke voorwaarden de injecteerde ruis zorgt voor voldoende drift langs richtingen met negatieve kromming, waardoor het algoritme niet vastloopt in suboptimale zadelpunten.
   • Convergentie: Bewezen dat bij een afkoelend ruis-schema (annealed noise schedule), de iteraties convergeren naar een stationair punt van het oorspronkelijke, niet-gesmoorde doelwit (MAP-objectief).
4. Empirische Validatie: Uitgebreide experimenten op natuurlijke RGB-afbeeldingen (FFHQ) en multi-coil MRI-data (fastMRI).

4. Resultaten

De methode werd getest op uitdagende taken zoals inpainting met grote maskers (box inpainting), super-resolutie, ontvlechting en versnelde MRI-reconstructie.

• Prestatieverbetering: SGPnP overtreft consequent traditionele PnP-methoden (zoals DPIR met DRUNet) en bestaande stochastische methoden (SNORE).
• Vergelijking met Diffusie-oplossers: SGPnP bereikt prestaties die concurrerend zijn met pure diffusie-based oplossers (zoals DPS en DiffPIR), maar behoudt de flexibiliteit van een optimalisatie-gebaseerd raamwerk.
• Kwalitatieve resultaten: Bij "box inpainting" (grote ontbrekende gebieden) falen deterministische PnP-methoden vaak en produceren onvolledige reconstructies. SGPnP slaagt erin realistische semantische inhoud te "hallucineren" en complete, geloofwaardige beelden te genereren.
• Ablatiestudies:
  • Het gebruik van SBDM's in plaats van CNN-denoisers (DRUNet) levert al verbetering op door de bredere ruisdekking.
  • De stochastische ruisinjectie is essentieel; zonder deze injectie (deterministische variant met SBDM) presteert het algoritme aanzienlijk slechter in ernstig ill-posed scenario's.
  • Het decouplen van injectie- en conditioneringsruis levert verdere stabiliteitswinst op.

5. Significance (Betekenis)

Dit werk vormt een brug tussen klassieke optimalisatie-methoden (PnP) en moderne generatieve modellen (Diffusion Models).

• Het lost het probleem op van hoe je krachtige generatieve prioren kunt gebruiken binnen een efficiënt optimalisatie-raamwerk zonder de rekenkosten van volledige sampling.
• Het biedt een theoretische onderbouwing voor het gebruik van ruisinjectie in PnP, niet alleen als een heuristiek, maar als een mechanisme om de optimalisatie te sturen weg van lokale minima en zadelpunten.
• De methode maakt robuuste reconstructie mogelijk in scenario's die voorheen als te moeilijk werden beschouwd voor PnP-methoden, zoals het invullen van zeer grote delen van een afbeelding of het herstellen van zwaar onderbemonsterde MRI-scans.

Kortom, SGPnP combineert de sterktes van score-based generatieve modellen met de stabiliteit van optimalisatie-algoritmen, resulterend in een krachtige tool voor complexe beeldherstelproblemen.