VIPaint: Image Inpainting with Pre-Trained Diffusion Models via Variational Inference

1. Enunciato del Problema

Il paper affronta la sfida dell'image inpainting (riempimento di regioni mancanti o mascherate di un'immagine) utilizzando Modelli Diffusivi Pre-addestrati (DM), in particolare Modelli Diffusivi Latenti (LDM) come Stable Diffusion.

• La Difficoltà Fondamentale: Sebbene i DM eccellano nel generare immagini nuove dal rumore, condizionarli su immagini corrotte o parzialmente osservate (inpainting) non è banale. L'obiettivo è campionare dalla vera distribuzione a posteriori condizionata $p(x|y)$, dove $y$ è l'immagine osservata (mascherata) e $x$ è l'immagine completa.
• Limiti dei Metodi Esistenti:
  • Metodi Euristic/Iterativi: Approcci come RePaint, CoPaint e DPS (Diffusion Posterior Sampling) spesso si basano su raffinamenti iterativi o euristiche di "viaggio nel tempo". Questi falliscono frequentemente nel catturare la struttura globale quando sono mascherate regioni ampie, portando a risultati incoerenti o irrealistici.
  • Metodi Variazionali (REDdiff): I precedenti tentativi di inferenza variazionale (ad es. REDdiff) ottimizzano spesso una singola immagine a livello privo di rumore ($t=0$) utilizzando una semplice approssimazione gaussiana. Questo non riesce a catturare l'incertezza, è soggetto a instabilità di ottimizzazione (specialmente in regimi a basso rumore) e fatica a generalizzare agli LDM.
  • Requisiti di Addestramento: Molti metodi all'avanguardia richiedono il fine-tuning del modello generativo su dataset o maschere specifici, il che è computazionalmente costoso e impedisce l'applicazione "zero-shot" a nuovi compiti.

2. Metodologia: VIPaint

Gli autori propongono VIPaint, un algoritmo di inferenza variazionale (VI) gerarchica che tratta il modello diffusivo pre-addestrato come un prior. Invece di ottimizzare una singola immagine, VIPaint ottimizza una distribuzione su variabili latenti a più livelli di rumore.

A. Posterior Variazionale Gerarchico

VIPaint definisce un'approssimazione non gaussiana e markoviana del vero posterior diffusivo. Introduce un insieme di $K$ "punti chiave" (step temporali intermedi) $\{h(1), \dots, h(K)\}$ all'interno della traiettoria di diffusione inversa (dove $h(K)$ è ad alto rumore e $h(1)$ è a basso rumore).

La distribuzione variazionale è definita come:
$$q_\lambda(z_{h(1):h(K)}) = q_\lambda(z_{h(K)}) \prod_{i=1}^{K-1} q_\lambda(z_{h(i)} | z_{h(i+1)})$$

• Struttura:
  • Alto Rumore ($h(K)$): Modellato come una Gaussiana fattorizzata.
  • Rumore Intermedio/Basso: Modellato come una Gaussiana dove la media è una combinazione convessa della predizione del prior diffusivo (dal modello pre-addestrato) e di un parametro variazionale appreso specifico per l'immagine.
  • Varianza: La varianza è appresa e specifica per dimensione, permettendo al modello di aumentare dinamicamente l'incertezza nelle regioni mascherate mantenendo la certezza nelle regioni osservate.

B. La Funzione Obiettivo

VIPaint ottimizza un limite inferiore variazionale sulla verosimiglianza marginale dell'osservazione $y$. La funzione di perdita consta di tre termini:

1. Perdita di Ricostruzione: Assicura che i campioni generati al livello di rumore più basso ($h(1)$) siano coerenti con i pixel osservati $y$. Questo è approssimato utilizzando una predizione di denoising in un singolo passo.
2. Perdita Prior: Regolarizza il livello di rumore più alto ($h(K)$) per allinearla con le semantica globali codificate dal prior diffusivo nel regime ad alto rumore.
3. Perdita Gerarchica: Assicura la coerenza tra i punti chiave adiacenti nella gerarchia, catturando dettagli e strutture di livello medio dell'immagine.

C. Ottimizzazione e Campionamento

• Ottimizzazione: I parametri variazionali $\lambda$ (medie, varianze e pesi di miscelazione) sono ottimizzati per immagine utilizzando la discesa del gradiente (Adam). Crucialmente, ciò avviene senza annealing (viene utilizzato il campionamento casuale del tempo per la perdita), il che fornisce un limite non distorto ed evita l'instabilità osservata nei metodi che applicano annealing da $T$ a $0$.
• Campionamento: Dopo l'ottimizzazione, vengono generate ricostruzioni diversificate tramite:
  1. Campionamento Ancestrale: Campionamento dal posterior gerarchico ottimizzato da $h(K)$ fino a $h(1)$.
  2. Raffinamento: Ulteriore raffinamento del campione da $h(1)$ fino a $t=0$ utilizzando la rete di denoising pre-addestrata, guidata dal gradiente della verosimiglianza (simile a DPS ma partendo da un posterior meglio inizializzato).

3. Contributi Chiave

1. Approssimazione Gerarchica Non Gaussiana: A differenza dei precedenti metodi VI che approssimano il posterior in un singolo punto ($t=0$), VIPaint modella il posterior attraverso una gerarchia di livelli di rumore. Questo gli permette di catturare simultaneamente sia le semantica globali (ad alto rumore) che i dettagli locali (a basso rumore).
2. Adattamento Zero-Shot: VIPaint non richiede il fine-tuning del modello diffusivo pre-addestrato. Adatta il modello a nuove maschere e immagini ottimizzando i parametri variazionali per ogni query, rendendolo applicabile a qualsiasi DM pre-addestrato (basato su pixel o latente).
3. Stabilità e Modellazione dell'Incertezza: Evitando il regime a basso rumore durante la fase principale di ottimizzazione (dove le reti di denoising sono non lisce e instabili) e modellando esplicitamente la varianza, VIPaint produce imputazioni diversificate e plausibili che riflettono l'incertezza dei dati mancanti.
4. Generalizzazione agli LDM: Il metodo è esplicitamente progettato per funzionare con Modelli Diffusivi Latenti, abilitando l'inpainting ad alta risoluzione (es. 1024x1024) con un costo computazionale inferiore rispetto agli approcci basati su pixel.

4. Risultati Sperimentali

Gli autori hanno valutato VIPaint su ImageNet64, ImageNet256, LSUN-Churches e Stable Diffusion 3.5 (condizionato al testo).

• Prestazioni Quantitative:

  • Maschere Grandi: VIPaint supera significativamente le linee di base (DPS, RePaint, CoPaint, REDdiff) su regioni mascherate ampie (40-80% dell'immagine) in termini di PSNR, SSIM e LPIPS.
  • Modelli Condizionati al Testo: Su Stable Diffusion 3.5, VIPaint ottiene un PSNR superiore (26.69 contro 25.22 per DPS) e un LPIPS migliore (0.050 contro 0.061), producendo immagini globalmente coerenti che preservano i dettagli fini.
  • Altri Problemi Inversi: Il metodo mostra anche prestazioni solide in deblurring e super-resolution, superando linee di base specializzate come PSLD e ReSample.

• Prestazioni Qualitative:

  • I risultati visivi mostrano che VIPaint genera strutture globalmente coerenti (ad es. forme di oggetti e illuminazione consistenti) dove altri metodi producono artefatti sfocati o incoerenti.
  • Genera con successo molteplici completamenti diversificati e plausibili per la stessa immagine mascherata, dimostrando la sua capacità di modellare la natura multi-modale del posterior.

• Efficienza:

  • VIPaint è computazionalmente efficiente. Per $K=2$, richiede solo circa 150 chiamate alla rete di denoising per l'ottimizzazione (rispetto a 1000+ per i metodi basati su campionamento) ed è più veloce di molte linee di base pur offrendo una qualità superiore.

5. Significato

VIPaint rappresenta un avanzamento significativo nei problemi inversi con modelli diffusivi. Spostandosi dal campionamento euristico o dall'ottimizzazione a punto singolo a un framework di inferenza variazionale gerarchica principiale, risolve la questione critica della incoerenza globale nell'inpainting con maschere ampie.

La sua capacità di funzionare zero-shot con modelli fondazionali all'avanguardia (come Stable Diffusion 3.5) senza re-addestramento lo rende uno strumento altamente pratico per compiti reali di editing, restauro e completamento di immagini. Colma il divario tra il rigore teorico dell'inferenza variazionale e le capacità pratiche dell'IA generativa moderna.