A Wavelet Diffusion GAN for Image Super-Resolution

Riepilogo Tecnico: Un GAN a Diffusione con Ondelette per la Super-Risoluzione delle Immagini

Enunciato del Problema
Sebbene i modelli di diffusione siano emersi come un'alternativa superiore alle Reti Generative Avversariali (GAN) per la generazione di immagini ad alta fedeltà, la loro applicazione in scenari in tempo reale è significativamente ostacolata da velocità di addestramento e inferenza lente. I tradizionali modelli di diffusione nello spazio dei pixel richiedono migliaia di passi temporali per generare campioni di alta qualità, richiedendo spesso diversi minuti per immagine. Sebbene i modelli di diffusione nello spazio latente tentino di mitigare questo problema addestrando un Autoencoder Variazionale (VAE), dipendono ancora da dataset su larga scala e mostrano tempi di esecuzione relativamente elevati. Al contrario, mentre i Diffusion GAN hanno accelerato la generazione incondizionata, permane la necessità di una soluzione che affronti le sfide specifiche della Super-Risoluzione a Immagine Singola (SISR), in particolare il compromesso tra output ad alta fedeltà ed efficienza computazionale.

Metodologia
Gli autori propongono WaDiGAN-SR, un nuovo framework che integra un Diffusion GAN condizionato con la Trasformata Discreta in Ondelette (DWT) per eseguire la super-risoluzione delle immagini. La metodologia opera attraverso i seguenti meccanismi chiave:

1. Diffusione nello Spazio delle Ondelette: A differenza degli approcci standard che eseguono la diffusione nello spazio dei pixel, questo metodo applica il processo di diffusione nel dominio delle ondelette. Le immagini ad alta risoluzione vengono decomposte mediante DWT in quattro sotto-bande (LL, LH, HL, HH), che rappresentano i componenti a bassa e alta frequenza. Questa decomposizione riduce le dimensioni spaziali dell'input (da $H \times W$ a $H/2 \times W/2$ per sotto-banda) preservando al contempo i dettagli ad alta frequenza essenziali per la nitidezza dell'immagine.
2. Architettura del Diffusion GAN Condizionato:
   • Processo Inverso: Il rumore viene aggiunto progressivamente alla decomposizione in ondelette dell'immagine ad alta risoluzione ($x_0$) fino a quando rimane solo rumore Gaussiano puro ($x_t$).
   • Processo Inverso: Un generatore $G(x_t, x_{lr}, t)$ prende l'immagine in ondelette rumorosa e la decomposizione in ondelette dell'input a bassa risoluzione ($x_{lr}$) come condizione. Invece di prevedere direttamente il rumore, il generatore prevede l'immagine pulita $x_0$.
   • Riduzione della Dimensione del Passo: Sfruttando un discriminatore GAN per distinguere tra coppie reali e false di $(x_{t-1}, x_t)$, il modello consente dimensioni del passo più grandi nel processo di diffusione inverso. Ciò riduce i passi temporali richiesti da migliaia a un numero predefinito (ad esempio, $T \leq 8$), accelerando significativamente l'inferenza.
3. Ottimizzazione: L'obiettivo di addestramento combina perdite avversariali (per garantire che la distribuzione generata corrisponda alla distribuzione dei dati reali) con una perdita di ricostruzione ($L_1$) tra l'immagine pulita generata e il ground truth. Questo approccio duale previene la perdita di informazioni di frequenza e mantiene la coerenza tra le sotto-bande delle ondelette.

Contributi Chiave

• Primo Diffusion GAN Condizionato Basato su Ondelette per SISR: Il documento introduce il primo approccio per applicare un Diffusion GAN condizionato basato su ondelette specificamente alla Super-Risoluzione delle Immagini. Sfrutta la riduzione delle dimensioni spaziali e il recupero migliorato dei dettagli attraverso le sotto-bande di ondelette ad alta frequenza.
• Bilanciamento tra Efficienza e Qualità: Il metodo proposto riduce con successo sia i tempi di addestramento che quelli di inferenza, superando le metodologie più avanzate (state-of-the-art) nella qualità dell'immagine. Funge da trampolino di lancio verso applicazioni in tempo reale dei modelli di diffusione nella super-risoluzione.
• Innovazione Architettonica: L'integrazione di blocchi consapevoli della frequenza e connessioni residue all'interno di un generatore simile a U-Net consente al modello di essere altamente sensibile ai componenti ad alta frequenza, producendo output più nitidi senza il carico computazionale della diffusione standard nello spazio dei pixel.

Risultati Sperimentali
Gli autori hanno valutato il modello sul dataset CelebA-HQ per un compito di ingrandimento 8x ($16 \times 16 \to 128 \times 128$) e hanno condotto uno studio di ablazione sul dataset Shipspotting.

• Prestazioni Quantitative: Nonostante sia stato addestrato solo per 25.000 iterazioni (limitato dai vincoli di risorse), il modello ha raggiunto metriche superiori rispetto a ESRGAN, SR3 e DiWa. Nello specifico, ha ottenuto un PSNR di 23.38 (contro 21.13 di ESRGAN), un SSIM di 0.68 e un LPIPS significativamente più basso di 0.061.
• Efficienza: Il modello ha dimostrato una drastica riduzione del tempo di inferenza, completando una generazione in 0.12 secondi rispetto ai 60.3 secondi di SR3 e ai 34.7 secondi di DiWa. Il modello è anche leggero, contenente 57M di parametri, che è meno rispetto alle baseline di diffusione (98M e 92M).
• Analisi Qualitativa: I risultati visivi hanno indicato una significativa riduzione degli artefatti (come pattern a griglia e spostamenti di colore) comuni in altre baseline. Le immagini generate hanno preservato dettagli più fini e hanno mostrato una maggiore fedeltà visiva.

Significato e Affermazioni
Il documento afferma che WaDiGAN-SR supera efficacemente gli svantaggi intrinseci dei modelli di diffusione nelle applicazioni sensibili al tempo. Combinando la stabilità e la flessibilità dei modelli di diffusione con la velocità delle GAN e la riduzione della dimensionalità delle trasformate in ondelette, l'approccio offre una qualità all'avanguardia con efficienza delle risorse. Gli autori notano che, sebbene i loro esperimenti siano stati vincolati da limitazioni hardware (addestramento su dimensioni di immagine più piccole e su un numero inferiore di iterazioni), i risultati validano l'efficienza del modello e la sua idoneità per la super-risoluzione in tempo reale. Il lavoro si posiziona come una soluzione praticabile per i ricercatori che operano sotto vincoli hardware e che richiedono output ad alta fedeltà senza i costi computazionali proibitivi dei modelli di diffusione tradizionali.