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

Sintesi Tecnica: Colmare il divario tra i priori basati sui dati tramite la funzione di score per il campionamento della distribuzione a posteriori

Definizione del problema
I compiti di restauro d'immagine (IR), come il denoising, il deblurring, l'inpainting e la super-risoluzione, sono tipicamente formulati come problemi inversi lineari ($y = Ax + n$). All'interno di un quadro bayesiano, la risoluzione di questi problemi mal posti richiede la definizione di una distribuzione a priori $p(x)$ per regolarizzare la soluzione. Sebbene i priori tradizionali basati su modelli (ad esempio, la Variazione Totale) siano ben compresi, essi spesso non riescono a catturare le strutture sofisticate delle immagini naturali. Al contrario, i priori basati sui dati appresi tramite il deep learning (ad esempio, Generative Adversarial Networks, Variational Autoencoders, Normalizing Flows) offrono una capacità espressiva superiore, ma pongono sfide significative per il campionamento della distribuzione a posteriori.

Una sfida primaria nel Bayesian IR è l'alta dimensionalità dello spazio delle immagini e la non regolarità delle funzioni di regolarizzazione, che ostacolano l'efficienza degli algoritmi Markov chain Monte Carlo (MCMC). I metodi esistenti spesso si affidano all'Unadjusted Langevin Algorithm (ULA) per campionare dalla distribuzione a posteriori. Tuttavia, l'ULA può incontrare difficoltà in problemi severamente mal condizionati a causa dei vincoli sulla dimensione del passo (step-size) legati alla costante di Lipschitz della verosimiglianza. Inoltre, molti avanzati priori basati sui dati (come quelli basati su denoiser o modelli generativi basati sullo score) sono definiti implicitamente, rendendo difficile il campionamento diretto senza un quadro teorico unificato.

Metodologia
Il lavoro propone un quadro unificato per campionare dalle distribuzioni a posteriori definite da diversi priori basati sui dati, a condizione che possano essere espressi tramite la loro funzione di score ($s(x) = \nabla_x \log p(x)$). La metodologia centrale si basa sul Langevin-within-Split Gibbs Sampler (LwSGS).

1. Asymptotically Exact Data Augmentation (AXDA): L'algoritmo introduce una variabile ausiliaria $z$ per decomporre la distribuzione a posteriori in una forma aumentata:
   $$\pi_\rho(x, z) \propto \exp\left[ -f(x, y) - g(z) - \frac{1}{2\rho^2}\|x - z\|^2 \right]$$
   dove $f(x,y)$ è il termine di fedeltà ai dati e $g(z)$ è il potenziale di regolarizzazione. Quando il parametro di accoppiamento $\rho \to 0$, la distribuzione marginale di $x$ converge alla distribuzione a posteriori target.

2. Split Gibbs Sampling (SGS): L'algoritmo alterna il campionamento di $x$ e $z$:

   • Campionamento di $x$: Dato $z$, la distribuzione condizionata $p(x|y, z)$ è gaussiana (per problemi inversi gaussiani lineari), consentendo un campionamento efficiente in forma chiusa. Fondamentalmente, questo passaggio scollega la dimensione del passo dalla costante di Lipschitz della verosimiglianza, affrontando una limitazione maggiore dell'ULA standard.
   • Campionamento di $z$: Dato $x$, la distribuzione condizionata $p(z|x)$ viene campionata utilizzando un passaggio di Langevin Monte Carlo (LMC). Questo passaggio utilizza esplicitamente la funzione di score del priori:
     $$z^{(t+1)} = z^{(t)} + \gamma \left( s(z^{(t)}) + \frac{1}{\rho^2}(x^{(t)} - z^{(t)}) \right) + \sqrt{2\gamma}\epsilon^{(t)}$$
     dove $s(z) = \nabla \log p(z)$ è la funzione di score del priori.

3. Rappresentazione Unificata dei Priori: Il documento dimostra che quattro distinte classi di priori basati sui dati possono essere unificate sotto questa prospettiva della funzione di score:

   • Regularization-by-Denoising (RED): La funzione di score è derivata direttamente dal residuo di denoising ($s(x) \propto -(x - D(x))$).
   • Score-Based Generative Models (SGM): La funzione di score è appresa esplicitamente tramite lo score matching (o modelli probabilistici di diffusione del denoising).
   • Normalizing Flows (NF): La funzione di score è calcolata tramite differenziazione automatica della log-densità derivata dalla formula del cambiamento di variabili.
   • Convex-Ridge Regularizers (CRR): La funzione di score corrisponde all'output di una singola CNN bias-free.

Contributi Chiave

• Unificazione dei Priori: Il documento stabilisce che i RED, SGM, NF e CRR, nonostante le diverse origini e motivazioni, possono essere tutti integrati in un unico framework di campionamento tramite le loro funzioni di score.
• Versatilità Algoritmica: Dimostra che LwSGS è un algoritmo generico capace di gestire qualsiasi log-prior differenziabile con una funzione di score esplicita o facilmente computabile, superando i limiti dei semplici schemi ULA.
• Disaccoppiamento dei Vincoli: Dividendo il problema, il metodo permette al passaggio di campionamento per il priori di essere indipendente dal condizionamento del termine di fedeltà ai dati, migliorando potenzialmente le proprietà di mixing in problemi mal posti.

Risultati Sperimentali
Gli autori hanno valutato il framework LwSS rispetto ai corrispettivi basati su ULA (PnP-ULA, NF-ULA e i nuovi SGM-ULA e CRR-ULA) su due compiti: inpainting di immagini e super-risoluzione a singola immagine.

• Dataset: Gli esperimenti sono stati condotti su 30 immagini RGB dai dataset FFHQ e ImageNet.
• Metriche di Prestazione: Lo studio ha utilizzato PSNR, SSIM, LPIPS (similarità percettiva) e l'Integrated Autocorrelation Time (IAT) per misurare la qualità della ricostruzione e l'efficienza del campionamento.
• Risultati:
  • Qualità della Ricostruzione: LwSGS ha prodotto costantemente ricostruzioni di alta qualità che corrispondono da vicino ai dettagli della verità fondamentale (ground-truth). Nel compito di inpainting, ha riempito efficacementmente le regioni mancanti; nella super-risoluzione, ha preservato dettagli intricati.
  • Confronto con ULA: LwSGS ha generalmente superato i metodi basati su ULA, in particolare quando utilizzava i priori RED e SGM, producendo dettagli più nitidi e una maggiore fedeltà. Anche con i priori NF e CRR, LwSGS ha mostrato prestazioni competitive o superiori.
  • Efficienza: L'algoritmo LwSGS ha dimostrato una convergenza più rapida, indicata da valori di IAT inferiori e funzioni di autocorrelazione che decadono più velocemente rispetto a ULA.
  • Quantificazione dell'Incertezza: Come metodo di campionamento, LwSGS ha permesso la stima dell'incertezza pixel-wise (deviazione standard), mostrando un'incertezza maggiore nelle regioni inpainting come previsto.
  • Applicazione nel Mondo Reale: Il framework è stato applicato alla restauro di campioni geologici reali (deblurring di immagini fuori fuoco). Utilizzando un priori RED, il metodo ha recuperato con successo dettagli nitidi da immagini sfocate, dimostrando utilità pratica nelle geoscienze.

Significato e Rivendicazioni
Il documento sostiene che il framework LwSGS fornisce un'alternativa versatile e sistematica all'ULA per problemi inversi bayesiani che coinvolgono priori basati sulla funzione di score. La sua importanza risiede in:

1. Ampia Applicabilità: Non è limitato a specifici modelli generativi ma può gestire qualsiasi priori in cui la funzione di score sia accessibile, colmando il divario tra vari paradigmi di regolarizzazione basati sui dati.
2. Efficienza Pratica: La strategia di divisione (splitting) mitiga efficacementamente i problemi di condizionamento inerenti all'ULA, portando a un migliore mixing e a una convergenza più rapida in spazi ad alta dimensionalità.
3. Utilità nel Mondo Reale: Il successo nel restauro di campioni geologici dimostra che il metodo è abbastanza robusto per pipeline operative dove i protocolli di acquisizione possono essere subottimali.

Gli autori concludono che, sebbene il presente lavoro si concentri su specifici priori, il framework unificato è sufficientemente generale da poter gestire qualsiasi distribuzione a posteriori definita da una vasta classe di priori basati sulla funzione di score, aprendo la strada alla futura integrazione di diversi modelli di deep learning nell'inferenza bayesiana.