Riemannian Diffusion Models on General Manifolds via Physics-Informed Neural Networks

Sintesi Tecnica: Modelli di Diffusione Riemanniana su Manifold Generali tramite Reti Neurali Informate dalla Fisica

1. Definizione del Problema

I modelli di diffusione riemanniana generalizzano la modellazione generativa basata sul punteggio (score-based) a dati che risiedono su varietà (manifold) lisce, definendo processi di diffusione in avanti tramite il moto Browniano riemanniano. Un collo di bottiglia critico nell'addestramento e nel campionamento di questi modelli è la necessità di accedere al kernel del calore del manifold (la densità di transizione del moto Browniano) e ai suoi gradienti. Nello specifico, l'obiettivo di addestramento (Denoising Score Matching) richiede la valutazione del punteggio condizionato $\nabla_x \log p_{t|0}(x|x_0)$, e il campionamento in avanti richiede il campionamento dal kernel del calore.

Sebbene il kernel del calore sia noto in forma chiusa per manifold altamente simmetrici (ad esempio, spazio euclideo, sfere, tori), esso è raramente disponibile analiticamente per manifold riemanniani generali. Le approssimazioni esistenti, come le asintotiche a tempo breve (formula di Varadhan) o le espansioni spettrali, sono spesso limitate a regimi specifici (tempo $t$ piccolo) o richiedono forti assunzioni strutturali (ad esempio, gruppi di Lie compatti), limitando la loro applicabilità pratica a diversi contesti geometrici.

2. Metodologia

Gli autori propongono un framework generale che approssima il kernel del calore risolvendo direttamente l'equazione del calore del manifold utilizzando Reti Neurali Informate dalla Fisica (PINN). Il metodo procede attraverso quattro fasi principali:

A. Selezione del Sistema di Coordinate ed Embedding
Invece di utilizzare patch di mappe di coordinate locali, il metodo inserisce (embed) il target manifold $M$ in un manifold ambiente $\tilde{M}$ più grande (tipicamente $\mathbb{R}^N$) come submanifold o manifold quoziente.

• Submanifold: Definiti da vincoli $f_i(x)=0$. La metrica e gli operatori differenziali sono derivati tramite una matrice di proiezione $P(x)$ che mappa lo spazio ambiente allo spazio tangente di $M$.
• Manifold Quoziente: Definiti identificando i punti sotto un'azione di gruppo (ad esempio, simmetria di permutazione). Il metodo tratta i punti dell'orbita come stati identici, utilizzando spesso rappresentanti canonici.

B. Derivazione dell'Equazione del Calore
Utilizzando le coordinate ambiente e la matrice di proiezione, gli autori derivano la forma esplicita dell'operatore di Laplace-Beltrami $\Delta_M$ e la corrispondente equazione di Fokker-Planck (del calore):
$$\partial_t p_{t|0}(x|x_0) = \Delta_M p_{t|0}(x|x_0)$$
Questa equazione è espressa interamente in termini della metrica ambiente e della matrice di proiezione, consentendo il calcolo basato su coordinate senza richiedere mappe di coordinate intrinseche.

C. Approssimazione della Condizione Iniziale
Poiché la condizione iniziale di Dirac $p_{0|0} = \delta_{x_0}$ non può essere imposta numericamente, gli autori inizializzano la PINN a un piccolo tempo $t_0 > 0$ utilizzando un'approssimazione basata sulle asintotiche a tempo breve. Utilizzano il metodo del parametro (ricorsione di Minakshisundaram–Pleijel) per ottenere un'espansione di ordine superiore del kernel del calore:
$$p_{t_0}(x|x_0) \approx G_t^{(n)}(x, x_0) = G_t(x, x_0) \left( u_0 + u_1 t + \dots + u_{n-1} t^{n-1} \right)$$
dove $G_t$ è il termine simile a una Gaussiana e $u_k$ sono coefficienti derivati dalla geometria (ad esempio, distorsione del volume e curvatura).

D. Addestramento della PINN
Una rete neurale $\phi_\theta(t, x)$ viene addestrata per approssimare il log-kernel del calore $\log p_{t|0}$. L'obiettivo di addestramento minimizza una combinazione di:

1. Perdita Iniziale/Al Confine: La differenza tra l'output della rete e l'inizializzatore asintotico al tempo $t=t_0$.
2. Perdita del Residuo della PDE: Il residuo dell'equazione del calore $\partial_t \phi - \Delta_M \phi - \|\nabla \phi\|^2 = 0$ (riscritta per la log-densità).

Gli autori impiegano una pesatura adattiva della perdita e l'apprendimento curricolare (espandendo progressivamente l'intervallo temporale) per garantire la stabilità.

E. Integrazione nei Modelli di Diffusione
La PINN addestrata funge da surrogato per il kernel del calore, abilitando due operazioni chiave:

1. Campionamento in Avanti: I campioni vengono generati tramite metodi MCMC (Markov Chain Monte Carlo, ad esempio, cammini casuali geodetici) utilizzando la densità valutata dalla PINN.
2. Valutazione del Punteggio (Score): Il punteggio condizionato $\nabla_x \log p_{t|0}$ viene calcolato differenziando la PINN, il che viene poi utilizzato per addestrare il modello di diffusione tempo-inverso tramite Denoising Score Matching (DSM).

3. Contributi Chiave

• Surrogato Generale Basato su PDE: Il lavoro introduce un metodo per approssimare i kernel del calore su arbitrarie varietà risolvendo numericamente l'equazione del calore, bypassando la necessità di soluzioni in forma chiusa o restrizioni spettrali.
• Derivazione Indipendente dalle Coordinate: Il framework fornisce un modo sistematico per derivare l'equazione del calore per manifold definiti come submani o quozienti di spazi euclidei, rendendolo applicabile a geometrie complesse come $S^2$, $SO(3)$, $SPD(n)$ e nuvole di punti con permutazione-quoziente.
• Analisi della Stabilità: Gli autori forniscono una giustificazione teorica (Teorema 4.1) dimostrando che la risoluzione dell'equazione del calore da una condizione iniziale approssimata produce una soluzione con errore limitato per tutti i $t \ge t_0$, garantendo l'affidabilità del surrogato PINN.
• Diffusione Permutazione-Invariante: Il metodo è applicato al manifold quoziente $\mathbb{R}^{k \times n}/S_n$, consentendo processi di diffusione che eliminano esplicitamente le permutazioni degli atomi per la generazione di molecole, un contesto in cui la diffusione euclidea standard non rispetta le simmetrie strutturali.

4. Risultati Sperimentali

Il metodo è stato valutato su diversi manifold e dataset:

• Dati Sferici ($S^2$): Stima della densità di eventi geofisici (vulcani, terremoti, inondazioni, incendi boschivi). Il metodo proposto ha ottenuto log-likelihood di test competitivi o superiori rispetto ai modelli generativi basati sullo score riemanniano (RSGM) che si affidano ad approssimazioni di Varadhan e metodi spettrali.
• Dati di Rotazione ($SO(3)$): La modellazione di miscele sintetiche ha mostrato log-likelihood migliorati rispetto ai baseline RSGM.
• Matrici Definite Positive ($SPD(n)$):
  • Analisi del Traffico ($SPD(10)$): Generazione condizionata di matrici di flusso del traffico dai dati dei taxi di NYC. Il metodo ha ridotto la distanza di Frobenius rispetto alla verità di base rispetto a SPD-DDPM.
  • Connettività Cerebrale ($SPD(n)$): Generazione condizionata di matrici di connettività EEG. Il metodo ha ottenuto precisione e punteggi F1 elevati, competendo con DiffeoCFM (che utilizza mappature diffeomorfiche nello spazio euclideo) e superando LowTriDDPM.
• Generazione di Molecole ($\mathbb{R}^{k \times n}/S_n$): Generazione di molecole nello spazio quoziente di permutazione. Il metodo ha prodotto molecole con maggiore validità e stabilità rispetto a un baseline EDM a varianza esplosiva (VE), riflettendo un noto compromesso tra validità e novità.

Validazione: Gli autori hanno validato i surrogati PINN confrontando le log-densità e gli score appresi con le approssimazioni di riferimento (ove disponibili) e riportando i residui della PDE normalizzati e gli errori al confine, confermando che i surrogati soddisfano l'equazione del calore.

5. Significato e Rivendicazioni

Il lavoro sostiene che l'approssimazione del kernel del calore basata su PDE offra un'alternativa pratica e generale ai metodi esistenti per la diffusione riemanniana. Trattando il kernel del calore come la soluzione di una PDE nota, l'approccio estende la diffusione riemanniana priva di drift a manifold dove i kernel analitici non sono disponibili.

Gli autori posizionano questo lavoro come una prospettiva "risolvi l'equazione del calore" che amplia l'applicabilità dei modelli di diffusione su manifold oltre le geometrie semplici. Tuttavia, riconoscono con modestia alcuni limiti:

• Costo Computazionale: L'addestramento della PINN richiede risorse significative a causa dell'autodifferenziazione di alto ordine e delle valutazioni di collocazione.
• Stabilità: L'addestramento può essere instabile in dimensioni elevate o per geometrie complesse, richiedendo potenzialmente bias induttivi più forti.
• Complessità: La pipeline è più complessa rispetto ad alternative come il Flow Matching, che evita completamente la stima del kernel del calore.

Il lavoro non pretende di sostituire il Flow Matching, ma piuttosto di fornire un backend necessario per i modelli basati sullo score su manifold dove l'uso di interpolazione geodetica o mappe esponenziali nel Flow Matching potrebbe essere insufficiente o dove la specifica struttura del kernel del calore è richiesta.