Sintesi Tecnica: DiffRGD

Definizione del Problema

I modelli di diffusione sono diventati fondamentali per la generazione di immagini, tuttavia controllarli per compiti specifici (ad esempio, restauro di immagini, generazione condizionata) senza costosi processi di ri-addestramento o fine-tuning rimane una sfida. Gli esistenti metodi di guida durante l'inferenza (ad esempio, DPS, FreeDoM) guidano il processo di campionamento iniettando gradienti da un obiettivo differenziabile direttamente nello spazio latente. Tuttavia, questi metodi operano nello spazio ambiente senza rispettare la struttura gaussiana sottostante del processo di diffusione.

Questa mancanza di consapevolezza della distribuzione porta al drift distributivo (deriva della distribuzione), in cui le variabili latenti deviano dal manifold gaussiano originale indotto dal modello. Man mano che i passi di campionamento si accumulano, questo fenomeno "fuori dal manifold" degrada la qualità del campione pulito finale. I tentativi precedenti per vincolare gli aggiornamenti, come MPGD (che assume la linearità locale) o DSG (che proietta su una ipersfera a raggio fisso), o si affidano a encoder perfetti non fattibili o non riescono a rispettare completamente le specifiche proprietà gaussiane della distribuzione latente, talvolta degradando verso una distribuzione uniforme.

Metodologia: DiffRGD

Gli autori propongono DiffRGD, un framework di guida durante l'inferenza plug-and-play che preserva esplicitamente la struttura gaussiana latente. L'intuizione centrale è che ad ogni timestep di campionamento DDIM $t$, la distribuzione latente è una gaussiana isotropa centrata su una media predetta $\mu_t$.

1. Decomposizione Polare e Costruzione del Manifold

Il metodo sfrutta una proprietà di decomposizione polare delle distribuzioni gaussiane isotrope. In qualsiasi timestep $t$, un campione latente $x_t$ può essere scomposto in una componente radiale $r$ e una componente direzionale $u$ statisticamente indipendenti:
$$x_t = \mu_t + \sigma_t r u$$
dove $r$ segue una distribuzione chi scalata e $u$ è uniformemente distribuito su una sfera unitaria.

Invece di ottimizzare nello spazio ambiente o su una sfera a raggio fisso, DiffRGD costruisce un manifold sferico $\mathcal{S}_{t,r}$ indotto dalla specifica distribuzione latente:
$$\mathcal{S}_{t,r} := \{ x \in \mathbb{R}^n \mid \|x - \mu_t\|_2 = \sigma_t r \}$$
Qui, il raggio $\sigma_t r$ è campionato dalla scala naturale della distribuzione, garantendo che l'insieme ammissibile rispetti le proprietà gaussiane del processo di diffusione.

2. Discesa del Gradiente Riemanniana (RGD)

Il problema di guida ad ogni timestep è formulato come un problema di ottimizzazione vincolata:
$$x_{t-1} = \arg\min_{x \in \mathcal{S}_{t,r}} \mathcal{L}(\hat{x}_0(x, t-1), y)$$
dove $\mathcal{L}$ è la perdita di guida (ad esempio, perdita di ricostruzione per il restauro o allineamento delle feature per la generazione condizionata).

Per risolvere questo, DiffRGD impiega la Discesa del Gradiente Riemanniana sul manifold sferico $\mathcal{S}_{t,r}$:

1. Proiezione nello Spazio Tangente: Il gradiente euclideo della perdita viene proiettato nello spazio tangente della sfera nel punto corrente, rimuovendo la componente radiale per garantire che la direzione dell'aggiornamento sia tangente al manifold.
2. Ritrazione: Dopo aver compiuto un passo lungo lo spazio tangente, un'operazione di ritrazione mappa il punto aggiornato nuovamente sul manifold sferico $\mathcal{S}_{t,r}$, preservando il raggio $\sigma_t r$.

Questo processo viene ripetuto per $K$ iterazioni interne (tipicamente $K=3$) per ogni timestep di diffusione. Gli autori dimostrano che, sotto ipotesi di $L$-smoothness, questa ottimizzazione converge a un punto stazionario.

Contributi Chiave

1. Geometria Consapevole della Distribuzione: Gli autori identificano una proprietà di decomposizione polare che rivela una geometria indotta dalla distribuzione per la guida della diffusione. Propongono un framework riemanniano che rispetta la geometria della distribuzione latente, mitigando il degrado dei campioni causato dal drift fuori dal manifold.
2. Analisi Teorica: Il lavoro fornisce un'analisi della preservazione della distribuzione e dimostra la convergenza dell'algoritmo di ottimizzazione riemanniana verso un punto stazionario. In pratica, l'ottimizzazione proposta, consapevole della geometria, converge più velocemente dei metodi basati su ADMM.
3. Performance Plug-and-Play: DiffRGD è dimostrato essere un metodo plug-and-play che non richiede ri-addestramento o fine-tuning, raggiungendo risultati allo stato dell'arte attraverso diversi compiti.

Risultati Sperimentali

Gli autori hanno valutato DiffRGD rispetto ai baseline (DPS, FreeDoM, MPGD, DSG, ADMMDiff) su una singola GPU NVIDIA GeForce RTX 4090 utilizzando il framework HuggingFace Diffusers.

Restauro di Immagini

Valutato su FFHQ 256×256 e ImageNet 256×256 per quattro compiti: Inpainting, Super-Risoluzione (4×), Deblurring Gaussiano e Deblurring di Movimento.

• Quantitativo: DiffRGD ha ottenuto miglioramenti statisticamente significativi in PSNR, SSIM, LPIPS e FID rispetto a tutti i baseline nella maggior parte dei compiti. Ad esempio, nell'Inpainting, DiffRGD ha raggiunto un PSNR di 34.04 rispetto ai 32.38 del secondo miglior metodo (ADMMDiff).
• Robustezza: DiffRGD ha mantenuto prestazioni elevate anche con meno passi di campionamento (100 passi), mentre metodi come DPS e ADMMDiff hanno mostrato un degrado evidente.
• Qualitativo: I risultati visivi hanno mostrato meno artefatti e una migliore coerenza dello sfondo rispetto a DPS e MPGD, con una qualità comparabile o superiore a ADMMDiff ma con un tempo computazionale inferiore.

Generazione Condizionata

Valutato su CelebA-HQ 256×256 per la guida tramite Mappa di Segmentazione, Sketch e FaceID.

• DiffRGD ha superato costantemente gli esistenti metodi di inferenza (FreeDoM, DSG, ADMMDiff) in termini di metriche di allineamento (mIoU, Sketch-$\ell_2$, FaceID-$\ell_2$) e qualità della generazione (FID, KID).
• I risultati qualitativi hanno indicato che DiffRGD produce risultati più naturali e visivamente accattivanti con un migliore allineamento alle condizioni rispetto a DSG, che spesso soffre di un controllo instabile.

Forza della Guida

Gli esperimenti sulla forza della guida ($\eta_t$) hanno mostrato che DiffRGD raggiunge un migliore compromesso tra allineamento (mIoU) e qualità della generazione (FID) attraverso un intervallo più ampio di intensità rispetto a DSG, che ha performato scarsamente in tutte le impostazioni.

Significato e Rivendicazioni

Il paper sostiene che DiffRGD affronti una lacuna critica nella guida durante l'inferenza: il fallimento degli esistenti metodi nel preservare la struttura gaussiana latente passo dopo passo. Formulando la guida come un'ottimizzazione vincolata su un manifold sferico indotto dalla distribuzione, DiffRGD previene i fenomeni "fuori dal manifold" che degradano la qualità del campione.

Gli autori affermano che la guida geometricamente coerente è un'alternativa valida e superiore all'applicazione diretta dei gradienti nello spazio ambiente o all'uso di manifold restrittivi e non consapevoli della distribuzione. Il metodo offre una soluzione plug-and-play semplice che migliora le metriche dei compiti preservando al contempo l'alta fedeltà del campione inerente ai modelli di diffusione pre-addestrati, senza la necessità di costosi ri-addestramenti.