Sintesi Tecnica: Apprendimento sulla Varietà con il Riemannian Flow Matching

1. Definizione del Problema

I recenti progressi nella modellazione generativa si sono spostati verso l'utilizzo di encoder di rappresentazione semantica ad alta dimensionalità congelati (es. DiNOv2, SigLIP) per migliorare la fedeltà della sintesi e la velocità di convergenza. Tuttavia, i normali framework Diffusion Transformer (DiT) e Flow Matching (FM) non riescono a convergere efficacemente quando vengono addestrati direttamente su queste rappresentazioni.

Il lavoro precedente, specificamente i Representation Autoencoders (RAE), ha attribuito questo fallimento a un collo di bottiglia della capacità, proponendo che i modelli richiedano una "scala di larghezza" (aumento della larghezza del transformer per corrispondere alla dimensionalità del latente) per risolvere il problema. Questo articolo mette in discussione tale ipotesi, sostenendo che la causa principale non sia l'insufficienza della capacità del modello, ma un disallineamento geometrico fondamentale, denominato Interferenza Geometrica.

Il problema centrale nasce dal fatto che gli encoder di rappresentazione proiettano le caratteristiche su una varietà ipersferica ($S^{d-1}$) con varianza radiale quasi nulla (una geometria a "guscio rigido"). Al contrario, il normale Flow Matching Euclideo assume un percorso di probabilità lineare (una corda) tra i dati e il rumore Gaussiano. Questa interpolazione lineare taglia attraverso l'interno a bassa densità dell'ipersfera, costringendo il modello a imparare campi di velocità in regioni non definite dello spazio delle caratteristiche. Di conseguenza, il modello spreca la propria capacità minimizzando gli errori radiali (che sono strutturalmente nulli sulla varietà) invece di apprendere le transizioni angolari semantiche.

2. Metodologia: Riemannian Flow Matching con Regolarizzazione di Jacobi (RJF)

Per risolvere l'Interferenza Geometrica senza modifiche architettoniche come la scala di larghezza, gli autori propongono RJF, un framework che allinea il processo generativo con la geometria intrinseca dello spazio latente.

2.1 Riemannian Flow Matching (Correzione della Traiettoria)

Invece dell'interpolazione lineare nello spazio Euclideo, RJF definisce il percorso di probabilità direttamente sulla varietà ipersferica utilizzando l'Interpolazione Lineare Sferica (SLERP).

• Percorsi Geodetici: La traiettoria tra un punto dato $x$ e il rumore $\epsilon$ (entrambi proiettati sulla sfera unitaria) segue la geodetica (arco di cerchio massimo) piuttosto che una corda.
• Velocità nello Spazio Tangente: Il campo di velocità target è calcolato differenziando il percorso SLERP. Ciò assicura che il vettore di velocità predetto giaccia strettamente nello spazio tangente della varietà ($v_t \cdot x_t = 0$), eliminando la necessità di apprendere componenti radiali.
• Campionamento: Durante l'inferenza, il modello utilizza un integratore Geodetico (Mappa Esponenziale) per aggiornare lo stato latente, garantendo che la traiettoria rimanga sulla superficie della varietà durante tutto il processo di generazione.

2.2 Regolarizzazione di Jacobi (Correzione della Curvatura)

Gli autori osservano che semplicemente correggere la traiettoria non è sufficiente perché lo standard Mean Squared Error (MSE) tratta gli errori di velocità in modo uniforme. Su un'ipersfera a curvatura positiva, le geodetiche si focalizzano (convergono), il che significa che gli errori di velocità si propagano in modo non lineare.

• Campi di Jacobi: Per tenere conto di ciò, gli autori introducono un fattore di ponderazione derivato dai campi di Jacobi, che quantificano la separazione tra le geodetiche causata dalle perturbazioni di velocità.
• Ripesatura della Perdita: La funzione di perdita è moltiplicata per un fattore di ponderazione geometrica $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$, dove $\Omega$ è la distanza geodetica. Questo riduce il peso degli errori vicino ai dati ($t=0$), dove il focalizzamento mitiga le perturbazioni, e dà priorità alla precisione vicino al rumore ($t=1$), effettuando efficacemente un annealing del segnale di apprendimento per adattarlo alla curvatura della varietà.

3. Contributi Chiave

1. Analisi Geometrica del Fallimento della Convergenza: L'articolo identifica l'Interferenza Geometrica come il collo di bottiglia primario per la diffusione standard sugli encoder di rappresentazione. Dimostra che il fallimento deriva da obiettivi Euclidei che costringono i modelli a imparare attraverso l'interno off-manifold dell'ipersfera, piuttosto che da una mancanza di capacità del modello.
2. Framework RJF: La proposta di un framework di addestramento geometricamente coerente che combina il Riemannian Flow Matching (per l'allineamento della traiettoria) e la Regolarizzazione di Jacobi (per la propagazione dell'errore consapevole della curvatura).
3. Efficienza senza Scala di Larghezza: La dimostrazione che le architetture DiT standard (es. DiT-B con 131M parametri) possono convergere efficacemente su latenti semantici ad alta dimensionalità senza la costosa scala di larghezza proposta dai lavori precedenti.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-1K a una risoluzione di 256×256 utilizzando varie scale di modello (DiT-B, DiT-L, DiT-XL) e diversi encoder di rappresentazione (DiNOv2-B, SigLIP, MAE).

• Convergenza e Prestazioni FID:
  • DiT-B (131M): Il normale Flow Matching Euclideo sulle caratteristiche DiNOv2-B è fallito nella convergenza (FID ~24.21). RJF ha ottenuto un FID di 6.77 in 80 epoche e 3.37 (con guidance) in 200 epoche.
  • DiT-XL (677M): RJF ha ottenuto un FID di 3.62 in 80 epoche senza guidance, superando il baseline Euclideo (FID 4.28) e il LightingDiT basato su VAE (FID 4.29).
  • Scaling: Il metodo ha mostrato miglioramenti consistenti attraverso le dimensioni del modello, raggiungendo Inception Scores allo stato dell'arte (IS 186.2) e Precision (0.82).
• Studi di Ablazione:
  • Radiale vs Angolare: Mascherare la perdita radiale ha fornito solo guadagni marginali (FID 21.99), confermando che il problema principale è il percorso lineare stesso (la corda).
  • Traiettoria: Passare al solo Riemannian Flow Matching (SLERP) ha migliorato drasticamente le prestazioni (FID 7.06), validando l'ipotesi del disallineamento della traiettoria.
  • Regolarizzazione: L'aggiunta della Regolarizzazione di Jacobi ha ulteriormente raffinato i risultati (FID 6.77).
• Generalizzazione: Il metodo è riuscito a convergere con successo su altri encoder con vincoli ipersferici, inclusi SigLIP (FID 10.39) e MAE (FID 19.82), dove i metodi standard sono falliti completamente.

5. Significato e Rivendicazioni

L'articolo afferma che la difficoltà di addestrare modelli di diffusione sulle rappresentazioni degli encoder è fondamentalmente geometrica, non architettonica. Rispettando la topologia ipersferica dello spazio latente, i normali Diffusion Transformer possono apprendere semantiche ad alta fedeltà in modo efficiente.

Gli autori asseriscono che RJF sblocca il potenziale delle architetture DiT standard, permettendo loro di eguagliare o superare le prestazioni di modelli molto più grandi o di pipeline complesse basate su VAE. Il lavoro stabilisce che l'efficace modellazione generativa sulle varietà di rappresentazione richiede un allineamento geometrico (traiettorie corrette e perdite consapevoli della curvatura) piuttosto che una scalatura bruta della capacità. Questo approccio consente la sintesi ad alta fedeltà con meno parametri e una convergenza più veloce, come dimostrato dal modello DiT-B che ha raggiunto un FID di 3.37 laddove i metodi precedenti non riuscivano a convergere.