Manifold-Optimal Guidance: A Unified Riemannian Control View of Diffusion Guidance

1. Il Problema: Deriva Fuori dalla Varietà (Off-Manifold Drift)

I modelli di diffusione condizionali, in particolare quelli che utilizzano la Classifier-Free Guidance (CFG), sono lo standard per la generazione di immagini di alta qualità. La CFG funziona estrapolando linearmente tra i punteggi (score) incondizionati e condizionati in uno spazio latente euclideo, controllata da un singolo parametro di scala ($w$).

Tuttavia, il paper identifica un fallimento geometrico fondamentale:

• Geometria Euclidea vs. Varietà dei Dati: La CFG tratta lo spazio latente come isotropo e euclideo. Quando la scala di guida ($w$) è alta (necessaria per un forte allineamento al prompt), l'aggiunta del vettore di guida spinge le traiettorie di campionamento fuori dalla varietà ad alta densità dei dati reali ($M$).
• Conseguenze: Questa "deriva fuori dalla varietà" (off-manifold drift) causa artefatti visivi gravi, tra cui:
  • Sovrasaturazione dei colori.
  • Texture innaturali e "plastiche".
  • Collasso strutturale dell'immagine.
• Limiti delle Soluzioni Esistenti: I metodi attuali (es. Dynamic Thresholding, CFG++, APG) agiscono come correzioni post-hoc o euristiche specifiche (es. proiezioni radiali o filtraggio in frequenza) senza fornire un obiettivo unificato che determini sia la direzione ottimale che la forza di aggiornamento in modo coerente con la geometria dei dati.

2. Metodologia: Manifold-Optimal Guidance (MOG)

Gli autori riformulano la guida come un problema di controllo ottimo locale su una varietà riemanniana. L'obiettivo è minimizzare l'energia condizionale (migliorare l'allineamento) restando il più vicino possibile alla varietà dei dati.

A. Formulazione Teorica

Invece di seguire il gradiente euclideo, MOG cerca la direzione di discesa più ripida rispetto a una metrica riemanniana $M_t(x_t)$ che penalizza i movimenti fuori dalla varietà.
Il problema di ottimizzazione è definito come:
$$\min_{u_t} \frac{1}{2} u_t^\top M_t u_t + \beta(t) \langle \nabla_x E, u_t \rangle$$
Dove:

• $u_t$ è l'aggiornamento del punteggio.
• $M_t$ è la metrica riemanniana (che codifica la geometria locale).
• $\nabla_x E$ è il gradiente dell'energia condizionale (approssimato dalla differenza tra score condizionato e incondizionato, $\Delta s$).
• $\beta(t)$ controlla la forza della guida.

La soluzione chiusa per l'aggiornamento ottimale è:
$$s_{MOG} = s_0 + \beta(t) M_t^{-1} \Delta s$$
Questa equazione rappresenta un aggiornamento del gradiente naturale riemanniano. La matrice inversa $M_t^{-1}$ "precondiziona" l'aggiornamento, attenuando le componenti che portano fuori dalla varietà e amplificando quelle tangenti ad essa.

B. Implementazioni Pratiche

Il paper introduce due varianti pratiche che non richiedono riaddestramento:

1. MOG-Score (Metrica Anisotropa):

   • Utilizza lo score incondizionato $s_0$ per stimare la normale alla varietà dei dati rumorosi.
   • Definisce una metrica di rango 1: $M_t = \lambda_\top I + (\lambda_\perp - \lambda_\top) n_t n_t^\top$, dove $n_t$ è la normale normalizzata.
   • Questo penalizza fortemente il movimento lungo la normale (fuori varietà) permettendo movimenti tangenziali.
   • L'inverso della matrice è calcolato efficientemente tramite la formula di Sherman-Morrison con complessità $O(d)$.

2. Auto-MOG (Bilanciamento Energetico Adattivo):

   • Risolve il problema della calibrazione manuale di $\beta(t)$ (la scala di guida).
   • Deriva $\beta(t)$ dinamicamente per bilanciare l'energia del punteggio a priori ($s_0$) e l'energia della guida naturale ($M_t^{-1}\Delta s$).
   • La formula adattiva è: $\beta_{auto}(t) = \gamma \cdot \frac{\|s_0\|_{M_t}}{\|M_t^{-1}\Delta s\|_{M_t} + \epsilon}$.
   • Questo permette alla guida di essere forte quando il segnale condizionale è debole e di attenuarsi quando è forte, prevenendo la sovrasaturazione senza tuning manuale.

3. Contributi Chiave

1. Quadro Teorico Unificato: La prima formulazione della guida come problema di controllo ottimo riemanniano, fornendo una giustificazione geometrica rigorosa per le correzioni necessarie.
2. Soluzione a Chiusura (Closed-Form): Derivazione di un aggiornamento analitico che corregge la deriva fuori dalla varietà senza modificare i pesi del modello.
3. Auto-MOG: Un meccanismo di schedulazione dinamica che elimina la necessità di tuning manuale degli iperparametri, adattandosi automaticamente al livello di rumore e alla difficoltà del prompt.
4. Efficienza Computazionale: L'implementazione richiede solo prodotti scalari e aggiunte vettoriali ($O(d)$), aggiungendo un overhead computazionale trascurabile rispetto all'inferenza della rete neurale.

4. Risultati Sperimentali

Il metodo è stato valutato su sei architetture diverse (UNet, DiT, MMDiT, Rectified Flow) e su diversi dataset (ImageNet, MS-COCO).

• Qualità Visiva (FID e Artefatti): MOG supera sistematicamente le baseline (CFG, CFG++, APG, LF-CFG).
  • Riduce drasticamente la sovrasaturazione e gli artefatti strutturali, specialmente ad alte scale di guida (es. $w=15$ su SD-XL).
  • Su DiT-XL/2, Auto-MOG ottiene un FID di 8.78 contro 12.45 di CFG++.
  • Su FLUX.1, ottiene un FID di 17.84 contro 18.15 di APG.
• Allineamento e Preferenza Umana:
  • Migliora i punteggi di allineamento testo-immagine (CLIP Score e HPSv2).
  • In uno studio con utenti umani (37 partecipanti, 2960 confronti), Auto-MOG è stato preferito nel 62% dei casi per qualità complessiva, con vantaggi significativi in fedeltà cromatica e realismo.
• Robustezza: Funziona bene su diverse scale di guida e architetture, mantenendo un equilibrio ottimale tra diversità e fedeltà (Pareto frontier).

5. Significato e Impatto

Questo lavoro rappresenta un cambio di paradigma nella guida dei modelli di diffusione:

• Dall'Euclideo al Riemanniano: Sposta la prospettiva da una semplice interpolazione lineare in uno spazio piatto a un'ottimizzazione geometrica consapevole della struttura intrinseca dei dati.
• Universalità: Fornisce una soluzione unificata che risolve problemi noti (sovrasaturazione, artefatti) senza richiedere riaddestramento costoso o modifiche architetturali complesse.
• Automazione: Con Auto-MOG, il processo di generazione diventa più robusto e facile da usare, eliminando la necessità per gli utenti di cercare manualmente la "scala di guida perfetta" per ogni modello o prompt.

In sintesi, MOG dimostra che comprendere e rispettare la geometria della varietà dei dati è fondamentale per generare immagini ad alta fedeltà, offrendo un metodo pratico e teoricamente solido per migliorare gli stati dell'arte attuali nella generazione di immagini.