Sintesi Tecnica: Problemi Inversi con Modelli di Diffusione tramite Perdita di Ricerca del Modo (Mode-Seeking Loss)

Definizione del Problema

Il documento affronta la sfida di risolvere problemi inversi utilizzando modelli di diffusione incondizionati pre-addestrati. Un problema inverso consiste nel stimare un campione di dati originale $x_0$ dato un elemento misurato parzialmente degradato $y$, governato dall'equazione $y = A(x_0) + \eta$, dove $A$ è un operatore di degradazione noto e $\eta$ è rumore gaussiano.

Sebbene i modelli di diffusione pre-addestrati possano risolvere questi problemi in modo zero-shot (plug-and-play), gli approcci esistenti presentano limitazioni significative:

1. Metodi di Campionamento della Posteriore: Mirano a estrarre campioni proporzionali alla probabilità a posteriori $p(x_0|y)$. Tuttavia, spesso si affidano ad approssimazioni di modellazione (ad es., approssimazioni gaussiane dello score condizionale) o metodi Monte-Carlo sequenziali, che possono essere computazionalmente onerosi e potrebbero non aderire strettamente alla vera posteriore senza una validazione basata su dati reali (ground-truth).
2. Metodi di Stima MAP: Mirano a trovare la stima Massima A Posteriori (MAP) (il modo della posteriore). I solver esistenti utilizzano spesso approssimazioni euristiche, si affidano allo splitting delle variabili (ad es., HQS, ADMM) o richiedono modelli di consistenza. Questi approcci implicano spesso costose computazioni del gradiente e assunzioni di modellazione che potrebbero non valere per tutti gli operatori lineari.

Gli autori sostengono che molti compiti di restauro d'immagine diano priorità alla generazione di un'immagine plausibile coerente con le misurazioni piuttosto che alla copertura dell'intera distribuzione posteriore. Ciò si allinea più strettamente alla stima MAP, motivando la necessità di una strategia MAP più rigorosa e priva di approssimazioni.

Metodologia: Variational Mode-Seeking Loss (VML)

Il contributo principale è la Variational Mode-Seeking Loss (VML), una nuova funzione obiettivo progettata per guidare i campioni verso la stima MAP durante il processo di diffusione inversa.

Fondamento Teorico

Gli autori derivano la VML da una prospettiva variazionale minimizzando la divergenza di Kullback-Leibler (KL) tra la posteriore di diffusione $p(x_0|x_t)$ e la posteriore della misurazione $p(x_0|y)$ ad ogni passo di diffusione inversa $t$:
$$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$$

L'intuizione è che man mano che $t \to 0$, la distribuzione $p(x_0|x_t)$ diventa progressivamente concentrata attorno a $x_t$. Trovando un $x_t^*$ ottimale che minimizzi questa divergenza KL ad ogni passo, il processo guida il campione verso i modi della posteriore $p(x_0|y)$.

Derivazione in Forma Chiusa per Operatori Lineari

Un importante traguardo teorico è la derivazione di un'espressione in forma chiusa per la VML nel caso di problemi inversi lineari ($A(x) = Hx$) senza alcuna approssimazione di modellazione. La completa espressione VML include:

1. Un termine di consistenza della misurazione: $\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$.
2. Termini del prior: che coinvolgono il denoiser $D(x_t, t)$, lo schedule del rumore $\sigma_t$ e la covarianza della posteriore $Cov[x_0|x_t]$.

Obiettivi Semplificati e Ripesati

La VML completa contiene termini di ordine superiore che coinvolgono la matrice di covarianza $Cov[x_0|x_t]$, i quali sono computazionalmente costosi da calcolare e differenziare. Gli autori dimostrano che questi termini di ordine superiore convergono uniformemente a zero (o a una costante) quando $t \to 0$. Di conseguenza, propongono una Simplified-VML ($VML_S$) che scarta questi termini, mantenendo solo i termini di consistenza della misurazione e del prior.

Per migliorare ulteriormente la stabilità dell'ottimizzazione, introducono una reweighted simplified-VML ($VML_S^\rho$):
$$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$$
dove $\rho_t$ è uno schedule del peso del prior tale che $\rho_t \to 1$ quando $t \to 0$. Questa formulazione consente un calcolo efficiente del gradiente utilizzando la differenziazione automatica standard.

Algoritmi

Sulla base della minimizzazione di VML, gli autori propongono:

1. VML-MAP: Un algoritmo che utilizza la discesa del gradiente per minimizzare $VML_S$ ad ogni passo di diffusione inversa.
2. VML-MAP$_{pre}$: Una variante precondizionata per operatori lineari mal condizionati. Utilizza la decomposizione in valori singolari (SVD) della matrice di degradazione $H$ per costruire un precondizionatore $P$, accelerando la convergenza per compiti difficili come il deblurring.
3. VML-MAP$^\tau$: Una strategia per il restauro di immagini rumorose. Per evitare artefatti acuti causati dalle sfide di ottimizzazione a bassi livelli di rumore, l'algoritmo minimizza la VML solo fino a una soglia di tempo di diffusione $\tau$, per poi passare al campionamento DDIM standard per generare l'immagine finale.
4. LatentVML-MAP: Un'estensione ai modelli di diffusione latente (LDM), approssimando la VML nello spazio latente.

Contributi Chiave

1. Nuova Funzione di Perdita: Introduzione della VML, che teoricamente guida i campioni verso le stime MAP minimizzando la divergenza KL tra le posteriori di diffusione e di misurazione.
2. Derivazione Senza Approssimazioni: Per i problemi inversi lineari, la VML è derivata in forma chiusa senza approssimazioni di modellazione euristica, fornendo una base teorica formale per la stima MAP con i modelli di diffusione.
3. Analisi di Convergenza: Dimostrazione formale che gli obiettivi VML semplificati e ripesati convergono alla negativa log-posterior (l'obiettivo MAP) quando $t \to 0$ sotto deboli assunzioni.
4. Algoritmi Pratici: Sviluppo di VML-MAP e della sua variante precondizionata, offrendo una soluzione empiricamente efficace per diversi compiti di restauro d'immagine.
5. Validazione Estesa: Dimostrazione di prestazioni superiori rispetto ai baseline esistenti (ad es., DDRM, $\Pi$GDM, MAPGA, DAPS) attraverso molteplici dataset (ImageNet, FFHQ, CelebA) e compiti (inpainting, super-resolution, deblurring).

Risultati Sperimentali

Gli autori hanno valutato i loro metodi su compiti di restauro d'immagine utilizzando modelli pre-addestrati su ImageNet256, FFHQ256 e CelebA.

• Metriche: Il documento enfatizza le metriche percettive (LPIPS e FID) rispetto alle metriche pixel-wise (PSNR, SSIM), sostenendo che queste ultime spesso non riescono a catturare la qualità percettiva delle ricostruzioni plausibili in problemi inversi mal posti.
• Prestazioni:
  • VML-MAP ha superato costantemente i metodi di campionamento della posteriore (DDRM, $\Pi$GDM) e altri stimatori MAP (MAPGA, DAPS) in termini di LPIPS e FID in compiti di inpainting, super-resolution e deblurring.
  • VML-MAP$_{pre}$ ha mostrato miglioramenti significativi rispetto alla versione non precondizionata per compiti mal condizionati (ad es., deblurring), ottenendo punteggi LPIPS e FID inferiori.
  • Efficienza Computazionale: VML-MAP ha ottenuto una migliore qualità percettiva con un tempo di esecuzione inferiore (meno valutazioni di funzioni neurali) rispetto ai metodi che richiedono molti passi di diffusione inversa.
• Restauro con Rumore: La strategia VML-MAP$^\tau$ ha gestito con successo misurazioni rumorose, producendo ricostruzioni di alta qualità dove lo standard VML-MAP faticava a causa di artefatti.
• Modelli Latenti: LatentVML-MAP ha dimostrato efficacia su CelebA, sebbene gli autori notino che l'ottimizzazione rimane più complessa nello spazio latente a causa della non linearità del decoder.

Significato e Rivendicazioni

Il documento afferma che la VML rappresenta la prima modellazione senza approssimazioni della stima MAP per problemi inversi lineari sfruttando esclusivamente un modello di diffusione incondizionato pre-addestrato, supportata da un'analisi teorica formale.

Gli autori sottolineano che il loro approccio differisce fondamentalmente dal campionamento della posteriore (che mira alla copertura della distribuzione) e dai moderni solver MAP (che spesso si affidano a euristiche o modelli di consistenza). Minimizzando la VML, il metodo punta efficacemente ai "modi" della posteriore, comportamento che gli autori ritengono sia l'obiettivo desiderato per generare restauri d'immagine plausibili e di alta qualità.

Il lavoro è modesto riguardo ai propri limiti, riconoscendo che:

• L'efficacia dipende dalla disponibilità di ottimizzatori praticamente efficienti per l'obiettivo VML.
• Le sfide di ottimizzazione persistono per operatori mal condizionati o non lineari, sebbene il precondizionamento aiuti.
• È necessario ulteriore lavoro per sviluppare ottimizzatori di ordine superiore per mitigare ulteriormente i problemi di ottimizzazione senza costi computazionali proibitivi.

Il documento conclude che VML-MAP fornisce un framework principato, teoricamente fondato ed empiricamente efficace per risolvere problemi inversi con i modelli di diffusione, colmando il divario tra la stima MAP teorica e il restauro d'immagine pratico.