MAD: Manifold Attracted Diffusion

Sintesi Tecnica: Manifold Attracted Diffusion (MAD)

Problema
I modelli di diffusione basati sul punteggio (score-based) si sono affermati come metodi allo stato dell'arte per la sintesi di immagini, superando spesso le GAN e i VAE in termini di qualità e diversità dei campioni. Tuttavia, sorge un limite significativo quando i dati di addestramento sono corrotti da rumore derivante da errori di misurazione, artefatti di compressione o processi di raccolta dati. I modelli di diffusione standard addestrati direttamente su tali dataset rumorosi imparano a riprodurre la corruzione, generando campioni che mantengono le caratteristiche del rumore dei dati di addestramento. Mentre gli approcci esistenti alla gestione di dati rumorosi richiedono spesso la modifica del processo di addestramento stesso (ad esempio, adattando le funzioni di perdita o richiedendo un mix di dati puliti e rumorosi), manca un metodo efficiente che possa generare campioni puliti da una distribuzione rumorosa utilizzando framework di addestramento consolidati e reti pre-addestrate senza necessità di ri-addestramento.

Metodologia
Il metodo proposto, Manifold Attracted Diffusion (MAD), affronta questo problema modificando la procedura di inferenza dei modelli basati sul punteggio piuttosto che il processo di addestramento. L'approccio si fonda sull'ipotesi del manifold, la quale postula che i dati significativi risiedano su un manifold a bassa dimensionalità all'interno di uno spazio ambiente ad alta dimensionalità. Gli autori interpretano il rumore come variazioni di piccola entità in direzioni ortogonali (fuori dal manifold) rispetto a questa struttura, mentre le variazioni dei dati significativi avvengono lungo il manifold (sul manifold).

L'innovazione tecnica centrale è l'introduzione dello score esteso.

• Definizione: Per una misura di probabilità $p$ e un parametro $\gamma > 0$, lo score esteso $H_\gamma p$ è definito come $(1 + \gamma)S(p * g_\gamma) + \gamma \frac{d}{d\gamma}S(p * g_\gamma)$, dove $S$ è l'operatore di score standard ($\nabla \log p$) e $g_\gamma$ è un kernel gaussiano.
• Proprietà: A differenza dello score standard, che è indefinito per distribuzioni di Dirac delta (che rappresentano dati su un manifold a dimensione inferiore con varianza nulla in certe direzioni), lo score esteso è ben definito. Fondamentalmente, per una delta di Dirac, lo score esteso si comporta in modo simile allo score di una Gaussiana con varianza unitaria. Ciò consente al processo di inferenza di trattare le direzioni a bassa varianza (fuori dal manifold) come se avessero una varianza non nulla, applicando efficacemente un effetto di "soglia morbida" (soft thresholding).
• Algoritmo di Inferenza: Gli autori propongono un algoritmo di inferenza iterativo che utilizza un'approssimazione dello score esteso. Dato un network di denoising pre-addestrato $D_\theta$ (che approssima lo score standard), lo score esteso viene approssimato valutando il network a un livello di rumore leggermente perturbato e calcolando una differenza finita per approssimare la derivata rispetto a $\gamma$.
• Implementazione: L'algoritmo introduce gli iperparametri ($a, b, p$) che controllano la forza dell' "attrazione al manifold". Regolando questi parametri, la procedura di inferenza sopprime le piccole variazioni nelle direzioni fuori dal manifold (rumore) preservando al contempo le variazioni maggiori nelle direzioni sul manifold (segnale). Il metodo è compatibile con i framework consolidati, come la formulazione dell'ODE del flusso di probabilità di Karras et al. (2022), e può utilizzare reti pre-addestrate.

Contributi Chiave

1. Definizione Formale e Analisi: Il documento fornisce una definizione matematica rigorosa dello score esteso, analizzando il suo comportamento sia su distribuzioni standard che su misure di Dirac delta, dimostrando la sua capacità di generalizzare l'operatore di score a distribuzioni degenerate.
2. Algoritmo di Inferenza Efficiente: Gli autori presentano un algoritmo di inferenza che riduce il rumore nei campioni generati senza richiedere una procedura di addestramento speciale. Ciò rende il metodo compatibile con le pipeline di addestramento esistenti e con i modelli pre-addestrati.
3. Validazione Empirica: L'efficacia dell'approccio è dimostrata attraverso problemi banali (toy problems), dati sintetici e dataset del mondo reale, inclusi FFHQ, AFHQ, ImageNet e dati di criomicroscopia elettronica (Cryo-EM).

Risultati

• Dati Toy e Sintetici: In esempi illustrativi in 2D, lo score esteso riesce a far collassare le traiettorie delle Gaussiane a bassa varianza alle loro medie (mimando le delta di Dirac) preservando al contempo la dispersione delle componenti ad alta varianza. Su dataset sintetici con rumore gaussiano aggiunto e rimozione di pixel, MAD supera significativamente l'inferenza dello score standard, generando campioni che recuperano le forme originali pulite.
• Dati Immagine del Mondo Reale: Gli esperimenti su FFHQ, AFHQv2 e ImageNet mostrano che l'aumento del parametro di attrazione al manifold ($b$) produce sfondi più puliti e caratteristiche primarie più nitide, effettuando efficacemente un filtraggio del rumore ad alta frequenza e delle variazioni secondarie.
• Metriche Quantitative: Su dataset CIFAR-10 e FFHQ corrotti da rumore gaussiano, shot e rimozione di pixel, MAD ottiene punteggi FID (Fréchet Inception Distance) significativamente più bassi rispetto all'inferenza standard. Notevolmente, i FID dell'inferenza standard corrispondono strettamente ai FID dei dati di addestramento rumorosi, confermando che i modelli standard riproducono la corruzione, mentre MAD deноisifica con successo.
• Applicazione Cryo-EM: Applicato al dataset EMPIAR-11618 (immagini Cryo-EM a singola particella rumorose), MAD ha generato campioni che somigliano fortemente alle forme raffinate ottenute con i metodi tradizionali, nonostante la rete sia stata addestrata esclusivamente su dati rumorosi.

Significatività e Rivendicazioni
Il documento sostiene che MAD offra una soluzione pratica per il "denoising generativo cieco" (blind generative denoising), in cui il modello di degradazione specifico è sconosciuto e sono disponibili solo dati rumorosi. La significatività del metodo risiede nella sua capacità di sfruttare la struttura geometrica dei manifold dei dati per filtrare il rumore durante la fase di generazione, piuttosto che fare affidamento su complesse adattazioni durante il tempo di addestramento.

Gli autori sono modesti riguardo all'ambito delle loro rivendicazioni:

• Non pretendono di risolvere problemi inversi (come il denoising di una singola immagine data), ma piuttosto di apprendere da un dataset rumoroso per produrre nuovi campioni puliti dalla distribuzione sottostante.
• Riconoscono che raggiungere prestazioni allo stato dell'arte in domini specifici (come la Cryo-EM) richiederebbe l'integrazione di conoscenze a priori più specifiche del dominio, il che è al di fuori dell'ambito di questo lavoro.
• Notano che, sebbene il metodo sia efficace, il costo computazionale è al massimo il doppio rispetto all'inferenza standard (a causa della valutazione extra del network), sebbene osservazioni empiriche suggeriscano che utilizzare lo score esteso solo nella seconda metà dei passi di inferenza possa essere sufficiente.

Il lavoro posiziona lo score esteso come un meccanismo geometricamente fondato per migliorare la qualità della generazione, potenzialmente applicabile anche a modelli addestrati su dati puliti, e apre la strada a ricerche future sulla selezione adattiva dei parametri e l'integrazione con modelli di diffusione condizionali.