Sintesi Tecnica: Restauro di Immagini Basato su Modelli di Diffusione tramite l'Apprendimento Interattivo di Determinismo e Stocasticità

Definizione del Problema

Il restauro di immagini (IR) mira a ricostruire informazioni visive latenti da osservazioni degradate, un compito complicato dalle significative discrepanze distribuzionali tra diversi tipi di degradazione (ad es., rumore, pioggia, ombre, scarsa illuminazione) nello spazio delle caratteristiche. Gli esistenti metodi di deep learning, inclusi CNN, Transformer e GAN, spesso faticano a generalizzare attraverso diversi compiti di IR senza riprogettazioni specifiche per il task. Sebbene i modelli di diffusione offrano alta fedeltà e interpretabilità, gli approcci attuali affrontano due limitazioni primarie:

1. Inefficienza nel IR: I modelli di diffusione standard iniziano il processo inverso partendo dal rumore puro, il che è inefficiente per i compiti di IR in cui l'input contiene già informazioni strutturali.
2. Modellazione Statica di Determinismo e Stocasticità: I precedenti tentativi di unificare la diffusione deterministica (basata sui residui) e quella stocastica (basata sul rumore), come il Residual Denoising Diffusion Model (RDDM), impiegano U-Net indipendenti per ciascun percorso. Questa separazione statica non riesce a catturare le distinte distribuzioni di probabilità dei diversi tipi di degradazione e può introdurre una diversità eccessiva che ostacola la ricostruzione precisa.

Metodologia

Gli autori propongono un nuovo framework di diffusione che impiega l'Apprendimento Interattivo di Determinismo e Stocasticità (ILDS) per affrontare queste limitazioni. L'architettura centrale utilizza una rete a doppio percorso basata su due semplici U-Net.

1. Costruzione del Processo di Diffusione Forward

A differenza della diffusione standard che aggiunge solo rumore, il processo forward proposto incorpora sia i residui (determinismo) che il rumore (stocasticità). La transizione è definita come:
$$x_t = x_{t-1} + \alpha_t x_{res} + \beta_t \epsilon_{t-1}$$
dove $x_{res} = x_{in} - x_0$ rappresenta il residuo di degradazione, e $\epsilon$ è il rumore gaussiano. Ciò consente alla traiettoria forward di preservare le strutture legate alla degradazione, permettendo al processo inverso di partire da stati più vicini al target rispetto al rumore puro, riducendo così i passi di campionamento richiesti.

2. Meccanismo di Apprendimento Interattivo

Per superare i limiti della modellazione indipendente di RDDM, gli autori introducono una strategia di apprendimento interattivo all'interno dell'architettura U-Net:

• Predizione a Doppio Percorso: Due U-Net vengono utilizzate rispettivamente per predire il rumore ($\epsilon_\theta$) e il residuo ($x_{res}^\theta$).
• Fusione delle Caratteristiche: Invece di operare in isolamento, gli encoder delle due U-Net condividono le caratteristiche. Nello specifico, dopo il primo blocco di downsampling, le caratteristiche di entrambi i percorsi vengono concatenate e processate attraverso convoluzioni $1 \times 1$ ($\tilde{h}_{res}$ e $\tilde{h}_\epsilon$).
• Modellazione Dinamica Dipendente dal Tempo: Gli input del timestep per i due percorsi sono modulati dinamicamente da parametri che controllano l'accumulo temporale ($\bar{\alpha}_t$) e l'intensità del rumore ($\bar{\beta}_t$). Ciò consente al modello di regolare adattivamente il contributo delle componenti deterministiche e stocastiche in base alla specifica distribuzione di probabilità delle caratteristiche del tipo di degradazione.
• Contesto Condiviso: Le caratteristiche interagite forniscono informazioni contestuali cross-path condivise durante la fase di decoding, garantendo che la guida deterministica e la diversità stocastica si influenzino a vicenda.

3. Addestramento e Generazione

Il modello è addestrato per minimizzare la divergenza di Kullback-Leibler (KL) tra la vera posterior e la transizione inversa appresa. La funzione di perdita finale è la somma degli errori quadratici medi per le predizioni del rumore e del residuo:
$$L = L_{res}(\theta) + L_{\epsilon}(\theta)$$
In particolare, questi componenti non sono pesati, poiché il meccanismo interattivo facilita un equilibrio auto-adattivo durante l'addestramento. Durante l'inferenza, il modello genera immagini di alta qualità campionando iterativamente dalla distribuzione condizionale, con il processo che diventa completamente deterministico quando il fattore di scala del rumore $\eta = 0$.

Contributi Chiave

1. Nuovo Framework di Diffusione: Il design di un framework che modella interattivamente gli aspetti deterministici e stocastici della degradazione dell'immagine, permettendo al processo di diffusione di adattarsi a compiti che richiedono o alta fedeltà (precisione) o diversità.
2. Modellazione Dinamica Dipendente dal Tempo: L'incorporazione di parametri di accumulo temporale e intensità del rumore per modulare dinamicamente gli input del timestep per i percorsi di doppia diffusione, migliorando la generalizzazione attraverso diversi scenari di IR.
3. Architettura a Doppio Percorso Efficiente: L'implementazione della doppia diffusione con l'uso di sole due semplici U-Net con caratteristiche dell'encoder condivise. Questo design raggiunge prestazioni competitive con un overhead di parametri minimo e un'alta scalabilità.

Risultati Sperimentali

Il metodo proposto è stato valutato su quattro compiti rappresentativi di IR: Denoising di Immagini, Rimozione di Ombre, Miglioramento della Scarsa Illuminazione e Derain (Rimozione Pioggia).

• Prestazioni: Il metodo ha ottenuto prestazioni superiori o competitive in tutti i compiti.
  • Denoising (SIDD): Ha raggiunto il miglior PSNR (41.67 dB) e SSIM (0.9385), superando RDDM e altri baseline.
  • Rimozione di Ombre (ISTD): Ha ottenuto il miglior PSNR (29.29 dB) e MAE (8.776), con il secondo miglior SSIM.
  • Derain (Raindrop): Ha raggiunto il miglior PSNR (31.30 dB) e SSIM (0.9460), rimuovendo con successo le striature di pioggia dove altri metodi lasciavano artefatti.
  • Miglioramento della Scarsa Illuminazione (LOL): Ha ottenuto la seconda migliore prestazione (PSNR 24.04, SSIM 0.9118), seguendo leggermente LLFormer nelle metriche, ma mantenendo una forte preservazione strutturale.
• Efficienza: L'apprendimento interattivo ha migliorato significativamente l'efficienza del campionamento. Il modello ha richiesto solo 1–4 passi di campionamento per ottenere risultati di alta qualità, rispetto alle centinaia o migliaia richiesti da altri metodi basati sulla diffusione (ad es., DiffusionVI richiedeva 1000 passi).
• Studi di Ablazione:
  • La rimozione del modulo di interazione ha causato cali di prestazioni in tutti i compiti, confermando il beneficio della fusione delle caratteristiche.
  • L'uso della modellazione dinamica dipendente dal tempo ($\bar{\alpha}_t, \bar{\beta}_t$) ha superato la modellazione statica del timestep ($t$).
  • I passi di campionamento ottimali variavano per task (1 per la scarsa illuminazione, 2 per il denoising, 3 per la rimozione di ombre, 4 per il derain).

Significato e Rivendicazioni

Il paper sostiene che il framework proposto offre una soluzione unificata per il restauro di immagini multi-task che bilancia accuratezza di ricostruzione e diversità. Integrando la modellazione deterministica e stocastica attraverso l'apprendimento interattivo, l'approccio fornisce una maggiore interpretabilità rispetto ai metodi basati su GAN e una migliore generalizzazione rispetto a CNN o Transformer specifici per il task.

Gli autori riconoscono alcune limitazioni, in particolare nei compiti di miglioramento della scarsa illuminazione con grandi gradienti di illuminazione, dove il ripristino della luminosità potrebbe essere insufficiente rispetto a baseline specializzate. Notano inoltre l'attuale dipendenza da dati accoppiati (puliti-degradati) per l'addestramento supervisionato. Il lavoro futuro è diretto verso l'ottimizzazione delle equazioni di diffusione per compiti sensibili all'illuminazione e lo sviluppo di strategie di addestramento non supervisionate per ridurre la dipendenza dai dati.

Nel complesso, lo studio dimosta che un'architettura di doppia diffusione semplice e interattiva può gestire efficacemente l'eterogeneità dei processi di degradazione delle immagini, offrendo un'alternativa scalabile ed efficiente per diverse sfide di restauro di immagini.