Bridging data-driven priors via the score function for posterior sampling -- Comparative review and experimental study

Technische Samenvatting: Het overbruggen van datagestuurde priors via de score-functie voor posterior sampling

Probleemstelling
Beeldrestauratietaken (IR), zoals denoising, deblurring, inpainting en super-resolutie, worden doorgaans geformuleerd als lineaire inverse problemen ($y = Ax + n$). Binnen een Bayesiaans kader vereist het oplossen van deze slecht gestelde problemen het definiëren van een prior-distributie $p(x)$ om de oplossing te regulariseren. Hoewel traditionele modelgebaseerde priors (bijv. Total Variation) goed begrepen zijn, schieten ze vaak tekort in het vastleggen van de complexe structuren van natuurlijke beelden. Daarentegen bieden datagestuurde priors die zijn geleerd via deep learning (bijv. Generative Adversarial Networks, Variational Autoencoders, Normalizing Flows) een superieure expressiviteit, maar vormen zij aanzienlijke uitdagingen voor posterior sampling.

Een primaire uitdaging in Bayesiaanse IR is de hoge dimensionaliteit van de beeldruimte en de niet-gladde aard van regularisatiefuncties, wat de efficiëntie van Markov chain Monte Carlo (MCMC) algoritmen belemmert. Bestaande methoden vertrouwen vaak op het Unadjusted Langevin Algorithm (ULA) om te samplen uit de posterior. Echter, ULA kan moeite hebben met ernstig slecht gestelde problemen vanwege stapgrootte-beperkingen die gekoppeld zijn aan de Lipschitz-constante van de likelihood-functie. Bovendien zijn veel geavanceerde datagestuurde priors (zoals die gebaseerd op denoisers of score-based generatieve modellen) impliciet gedefinieerd, wat directe sampling moeilijk maakt zonder een verenigd theoretisch kader.

Methodologie
Het artikel stelt een verenigd kader voor om te samplen uit posterior-distributies die gedefinieerd worden door diverse datagestuurde priors, mits deze uitgedrukt kunnen worden via hun score-functie ($s(x) = \nabla_x \log p(x)$). De kernmethodologie berust op de Langevin-within-Split Gibbs Sampler (LwSGS).

1. Asymptotisch exacte data-augmentatie (AXDA): Het algoritme introduceert een hulpvariabele $z$ om de posterior-distributie te deconstrueren naar een geaugmenteerde vorm:
   $$\pi_\rho(x, z) \propto \exp\left[ -f(x, y) - g(z) - \frac{1}{2\rho^2}\|x - z\|^2 \right]$$
   waarbij $f(x,y)$ de data-fidelity term is en $g(z)$ de regularisatie-potentiaal. Wanneer de koppelingsparameter $\rho \to 0$ gaat, convergeert de marginale distributie van $x$ naar de doel-posterior.

2. Split Gibbs Sampling (SGS): Het algoritme wisselt af tussen het samplen van $x$ en $z$:

   • Samplen van $x$: Gegeven $z$, is de conditionele distributie $p(x|y, z)$ Gaussisch (voor lineaire Gaussische inverse problemen), wat efficiënte, gesloten vorm-sampling mogelijk maakt. Cruciaal is dat deze stap de stapgrootte ontkoppelt van de Lipschitz-constante van de likelihood, wat een belangrijke beperking van standaard ULA aanpakt.
   • Samplen van $z$: Gegeven $x$, wordt de conditionele distributie $p(z|x)$ gesampled met behulp van een Langevin Monte Carlo (LMC) stap. Deze stap maakt expliciet gebruik van de prior score-functie:
     $$z^{(t+1)} = z^{(t)} + \gamma \left( s(z^{(t)}) + \frac{1}{\rho^2}(x^{(t)} - z^{(t)}) \right) + \sqrt{2\gamma}\epsilon^{(t)}$$
     waarbij $s(z) = \nabla \log p(z)$ de score-functie van de prior is.

3. Verenigde Prior-representatie: Het artikel demonstreert dat vier verschillende klassen van datagestuurde priors kunnen worden verenigd onder dit score-functie perspectief:

   • Regularization-by-Denoising (RED): De score-functie wordt direct afgeleid van de denoising residu ($s(x) \propto -(x - D(x))$).
   • Score-Based Generative Models (SGM): De score-functie wordt expliciet geleerd via score matching (of denoising diffusion probabilistic models).
   • Normalizing Flows (NF): De score-functie wordt berekend via automatische differentiatie van de log-densiteit afgeleid van de change of variables formule.
   • Convex-Ridge Regularizers (CRR): De score-functie komt overeen met de bias-vrije output van een single-layer CNN.

Belangrijkste Bijdragen

• Vereniging van Priors: Het artikel stelt vast dat RED, SGM, NF en CRR priors, ondanks hun verschillende oorsprong en motivaties, allemaal geïntegreerd kunnen worden in één enkel sampling framework via hun score-functies.
• Algoritmische Veelzijdigheid: Het toont aan dat LwSGS een generiek algoritme is dat elke differentiabele log-prior kan afhandelen met een expliciete of efficiënt berekenbare score-functie, waarmee de beperkingen van eenvoudige ULA-schema's worden overwonnen.
• Ontkoppeling van Beperkingen: Door het probleem te splitsen, staat de methode toe dat de sampling-stap voor de prior onafhankelijk is van de conditionering van de data-fidelity term, wat potentieel de mixing-eigenschappen in slecht gestelde problemen verbetert.

Experimentele Resultaten
De auteurs hebben het LwSGS-framework geëvalueerd tegenover ULA-tegenhangers (PnP-ULA, NF-ULA, en nieuw ontworpen SGM-ULA en CRR-ULA) op twee taken: beeld inpainting en single-image super-resolutie.

• Datasets: Experimenten werden uitgevoerd op 30 RGB-beelden van de FFHQ en ImageNet datasets.
• Prestatie-metrieken: De studie maakte gebruik van PSNR, SSIM, LPIPS (perceptuele gelijkenis) en Integrated Autocorrelation Time (IAT) om zowel de reconstructiekwaliteit als de sampling-efficiëntie te meten.
• Bevindingen:
  • Reconstructiekwaliteit: LwSGS produceerde consistent hoogwaardige reconstructies die de grondwaarheid (ground-truth) details nauwgezet volgden. Bij inpainting-taken vulde het effectief ontbrekende regio's in; bij super-resolutie behield het complexe details.
  • Vergelijking met ULA: LwSGS presteerde over het algemeen beter dan ULA-gebaseerde methoden, vooral bij het gebruik van RED en SGM priors, wat resulteerde in scherpere details en hogere getrouwheid. Zelfs met NF en CRR priors vertoonde LwSGS competitieve of superieure prestaties.
  • Efficiëntie: Het LwSGS-algoritme demonstreerde snellere convergentie, aangegeven door lagere IAT-waarden en sneller afnemende autocorrelatie-functies vergeleken met ULA.
  • Onzekerheidskwantificatie: Als sampling-methode maakte LwSGS de schatting van pixel-gewijze onzekerheid (standaarddeviatie) mogelijk, waarbij hogere onzekerheid in inpainted regio's werd aangetoond, zoals verwacht.
  • Real-World Toepassing: Het framework werd toegepast op het herstellen van echte geologische tailings-monsters (deblurring van out-of-focus beelden). Met een RED-prior slaagde de methode erin om scherpe details te herstellen uit onscherpe beelden, wat de praktische bruikbaarheid in de geowetenschappen aantoont.

Betekenis en Claims
Het artikel claimt dat het LwSGS-framework een veelzijdige en systematische alternatief biedt voor ULA voor Bayesiaanse inverse problemen met score-functie-gebaseerde priors. De betekenis ligt in:

1. Brede Toepasbaarheid: Het is niet beperkt tot specifieke generatieve modellen, maar kan elke prior afhandelen waar de score-functie toegankelijk is, waardoor de kloof tussen diverse datagestuurde regularisatieparadigma's wordt overbrugd.
2. Praktische Efficiëntie: De splitting-strategie mitigeert effectief de conditioneringsproblemen die inherent zijn aan ULA, wat leidt tot betere mixing en snellere convergentie in hoog-dimensionale ruimtes.
3. Real-World Nut: De succesvolle restauratie van geologische monsters toont aan dat de methode robuust genoeg is voor operationele pipelines waar acquisitieprotocollen suboptimaal kunnen zijn.

De auteurs concluderen dat hoewel het huidige werk zich richt op specifieke priors, het verenigde kader voldoende algemeen is om elke posterior-distributie te behandelen die gedefinieerd wordt door een brede klasse van score-functie-gebaseerde priors, wat de weg vrijmaakt voor toekomstige integratie van diverse deep learning modellen in Bayesiaanse inferentie.