An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

Problematica
I modelli di diffusione (DM) hanno dimostrato una capacità eccezionale nel risolvere problemi inversi di imaging computazionale modellando complessi prior d'immagine. Tuttavia, la loro applicazione pratica è severamente limitata dalla dipendenza da dataset su larga scala, di alta qualità e "puliti" per il pre-addestramento. In molti scenari reali, in particolare nell'imaging scientifico e biomedico (ad esempio, la TC a raggi X, la criomicroscopia elettronica), l'acquisizione di tali dati puliti è impraticabile o proibitivamente costosa. Al contrario, le osservazioni corrotte (rumorose, sfocate o mascherate) con modelli forward differenziabili sono abbondanti. Ciò crea un dilemma "uovo e gallina": addestrare un modello di diffusione accurato richiede immagini pulite, ma ricostruire immagini pulite da osservazioni corrotte richiede un modello di diffusione ben addestrato. I metodi esistenti per apprendere da dati corrotti, come AmbientDiffusion o SURE-Score, spesso faticano con specifici tipi di corruzione (ad esempio, inpainting su larga scala o rumore significativo) o non riescono a generalizzare attraverso diversi problemi inversi.

Metodologia: EMDiffusion
Gli autori propongono EMDiffusion, un nuovo framework di tipo Expectation-Maximization (EM) progettato per risolvere congiuntamente i problemi inversi di imaging e addestrare un modello di diffusione direttamente da osservazioni corrotte. Il metodo opera attraverso un ciclo iterativo di due fasi:

1. Inizializzazione: Il processo inizia con un prior di diffusione "vago" addestrato su un set molto piccolo di immagini pulite (anche solo 10–50 immagini). Questa inizializzazione è cruciale per evitare che il modello collassi in una distribuzione definita esclusivamente dagli input corrotti. Gli autori notano che queste immagini pulite iniziali non devono necessariamente appartenere alla stessa distribuzione del task target (dati fuori distribuzione possono servire come efficaci inizializzatori).

2. E-Step (Expectation): Utilizzando l'attuale modello di diffusione come prior, l'algoritmo esegue un Adaptive Diffusion Posterior Sampling per ricostruire immagini pulite dalle osservazioni corrotte.

   • Il campionamento posteriore standard spesso fallisce nelle prime iterazioni perché il prior è debole, portando al "collasso del modo" (mode collapse) dove i campioni ignorano l'osservazione e ritornano alla distribuzione di addestramento.
   • Per mitigare questo problema, EMDiffusion introduce un iperparametro $\lambda$ che riscala il termine della verosimiglianza (likelihood) nella SDE (equazione differenziale stocastica) a tempo inverso.
   • $\lambda$ è sintonizzato in modo adattivo in ogni E-step per minimizzare la perdita dei dati ($\|y - A(\hat{x}_0)\|_2^2$), bilanciando l'influenza del prior imperfetto rispetto ai dati osservati. Ciò consente la generazione di ricostruzioni più pulite anche quando il prior è inizialmente scarso.

3. M-Step (Maximization): L'algoritmo aggiorna i pesi del modello di diffusione utilizzando le immagini ricostruite (campioni posteriori) generate nell'E-step.

   • Questo passaggio assomiglia al classico denoising score matching, dove il modello impara ad approssimare la funzione di score dei dati ricostruiti "puliti".
   • Per accelerare l'addestramento ed evitare l'overfitting sulle prime ricostruzioni rumorose, il metodo impiega una strategia di addestramento a due fasi: inizialmente, effettua il fine-tuning dei pesi ereditati dall'iterazione precedente; una volta che la qualità della ricostruzione migliora (indicata dal fatto che il $\lambda$ ottimale scende sotto 1), il modello viene riinizializzato e riaddestrato da zero sull'intero set di campioni posteriori.

Contributi Chiave

• Framework EM per Dati Corrotti: Il paper introduce il primo framework generale per addestrare modelli di diffusione basati sullo score da osservazioni arbitrariamente corrotte utilizzando un approccio EM, rompendo efficacemente la dipendenza da grandi dataset puliti.
• Adaptive Posterior Sampling: La proposta di un fattore di scala adattivo ($\lambda$) per il termine di verosimiglianza nella SDE inversa affronta il problema del collasso del modo inerente alle iterazioni EM iniziali, consentendo una ricostruzione stabile anche con prior deboli.
• Raffinamento Iterativo: Il metodo dimostra che una piccola quantità di dati puliti può agire come un regularizzatore implicito, permettendo al modello di convergere progressivamente verso la vera distribuzione dei dati puliti attraverso passi alternati di ricostruzione e addestramento.
• Generalizzabilità: A differenza dei lavori precedenti focalizzati su corruzioni specifiche (come il masking), EMDiffusion fornisce un approccio unificato applicabile all'inpainting casuale, al denoising e al deblurring.

Risultati Sperimentali
Gli autori validano EMDiffusion su dataset CIFAR-10 e CelebA attraverso tre task: inpainting casuale (mascheramento di pixel al 60%), denoising gaussiano e deblurring gaussiano.

• Inpainting: EMDiffusion supera significativamente i baseline come AmbientDiffusion e SURE-Score, raggiungendo una qualità di ricostruzione (PSNG/LPIPS) comparabile a quella di Diffusion Posterior Sampling (DPS) che utilizza un prior addestrato sull'intero dataset pulito.
• Denoising & Deblurring: Il metodo supera SURE-Score e i baseline auto-supervisionati (Noise2Self), recuperando dettagli fini (ad esempio, ruote di auto, tratti del viso) che altri metodi tendono a smussare.
• Qualità Generativa: I modelli di diffusione appresi raggiungono punteggi FID allo stato dell'arte rispetto ad altri metodi addestrati su dati corrotti. Notevolmente, un modello appreso da 50.000 immagini corrotte tramite EM performa meglio di un modello addestrato su circa 15.000 immagini pulite.
• Efficienza: Il metodo ottiene questi risultati utilizzando un'architettura DDPM standard (35.7M di parametri), che è significativamente più piccola delle architetture utilizzate dai baseline concorrenti (ad esempio, i 113M+ parametri di AmbientDiffusion).

Significatività e Rivendicazioni
Il paper sostiene che EMDiffusion dimostra che è informazionalmente possibile apprendere la distribuzione sottostante dei dati puliti principalmente da misurazioni corrotte. La significatività risiede in:

1. Superamento della Scarsità di Dati: Offre una soluzione pratica per domini in cui i dati puliti non sono disponibili, sfruttando l'abbondanza di osservazioni corrotte.
2. Inizializzazione Robusta: Stabilisce che un set minimo di immagini pulite (anche fuori distribuzione) è sufficiente per inizializzare il manifold, impedendo al modello di collassare in una distribuzione distorta dei dati corrotti.
3. Prestazioni allo Stato dell'Arte: Il metodo raggiunge nuove prestazioni allo stato dell'arte nei compiti di imaging inverso senza richiedere modelli pre-addestrati su grandi dataset puliti, colmando il divario tra i prior teorici e i vincoli pratici dei dati.

Gli autori mantengono la modestia riguardo all'ambito, riconoscendo che il metodo si basa sull'assunto che il modello forward sia noto e differenziabile, e che il lavoro futuro potrebbe estendersi all'apprendimento di modelli forward o statistiche di rumore sconosciute.