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

Technische Samenvatting: Riemanniaanse Diffusiemodellen op Algemene Manifolds via Physics-Informed Neural Networks

1. Probleemstelling

Riemanniaanse diffusiemodellen generaliseren score-gebaseerde generatieve modellering naar data die zich op gladde manifolds bevindt door voorwaartse diffusieprocessen te definiëren via de Riemanniaanse Brownse beweging. Een kritieke bottleneck bij het trainen en samplen van deze modellen is de vereiste om de manifold heat kernel (de transitiedichtheid van de Brownse beweging) en de bijbehorende gradiënten te kunnen berekenen. Specifiek vereist het trainingsdoel (Denoising Score Matching) het evalueren van de conditionele score $\nabla_x \log p_{t|0}(x|x_0)$, en voorwaartse sampling vereist het trekken van samples uit de heat kernel.

Hoewel de heat kernel in gesloten vorm bekend is voor zeer symmetrische manifolds (bijv. de Euclidische ruimte, sferen, tori), is deze voor algemene Riemanniaanse manifolds zelden analytisch beschikbaar. Bestaande benaderingen, zoals korttijd-asymptotica (Varadhans formule) of spectrale expansies, zijn vaak beperkt tot specifieke regimes (kleine tijd $t$) of vereisen sterke structurele aannames (bijv. compacte Lie-groepen), wat de praktische toepasbaarheid op diverse geometrische instellingen beperkt.

2. Methodologie

De auteurs stellen een algemeen framework voor dat de heat kernel benadert door direct de manifold-warmtevergelijking op te lossen met behulp van Physics-Informed Neural Networks (PINNs). De methode verloopt via vier hoofdfasen:

A. Selectie van Coördinatensysteem en Embedding
In plaats van lokale coördinatenkaarten (coordinate charts) te gebruiken, embed de methode de doel-manifold $M$ in een grotere omgeving-manifold $\tilde{M}$ (meestal $\mathbb{R}^N$) als een sub-manifold of quotient-manifold.

• Submanifolds: Gedefinieerd door restricties $f_i(x)=0$. De metriek en differentiaaloperatoren worden afgeleid met behulp van een projectiematrix $P(x)$ die de omgeving-ruimte naar de raakruimte (tangent space) van $M$ mapt.
• Quotient Manifolds: Gedefinieerd door punten te identificeren onder een groepswerking (bijv. permutatiesymmetrie). De methode behandelt orbit-punten als identieke toestanden, waarbij vaak canonieke representanten worden gebruikt.

B. Afleiding van de Warmtevergelijking
Gebruikmakend van de omgeving-coördinaten en de projectiematrix, leiden de auteurs de expliciete vorm van de Laplace-Beltrami-operator $\Delta_M$ en de bijbehorende Fokker-Planck (warmte)vergelijking af:
$$\partial_t p_{t|0}(x|x_0) = \Delta_M p_{t|0}(x|x_0)$$
Deze vergelijking wordt volledig uitgedrukt in termen van de omgeving-metriek en de projectiematrix, waardoor coördinaat-gebaseerde berekeningen mogelijk zijn zonder dat er intrinsieke coördinatenkaarten nodig zijn.

C. Benadering van de Initiële Conditie
Omdat de Dirac-delta initiële conditie $p_{0|0} = \delta_{x_0}$ numeriek niet kan worden opgelegd, initialiseren de auteurs de PINN op een kleine tijd $t_0 > 0$ met behulp van een gefundeerde korttijd-asymptotische benadering. Zij maken gebruik van de parametrix-methode (Minakshisundaram–Pleijel recursie) om een hogere-orde expansie van de heat kernel te verkrijgen:
$$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)$$
waarbij $G_t$ de Gaussian-achtige term is en $u_k$ coëfficiënten zijn afgeleid van de geometrie (bijv. volumedistorsie en kromming).

D. PINN Training
Een neuraal netwerk $\phi_\theta(t, x)$ wordt getraind om de log-heat kernel $\log p_{t|0}$ te benaderen. Het trainingsdoel minimaliseert een combinatie van:

1. Initial/Boundary Loss: Het verschil tussen de netwerkoutput en de asymptotische initializer bij $t=t_0$.
2. PDE Residual Loss: De residu van de warmtevergelijking $\partial_t \phi - \Delta_M \phi - \|\nabla \phi\|^2 = 0$ (herschreven voor de log-densiteit).

De auteurs maken gebruik van adaptieve loss-weging en curriculum learning (het progressief uitbreiden van het tijdsinterval) om stabiliteit te garanderen.

E. Integratie in Diffusiemodellen
De getrainde PINN dient als een surrogaat voor de heat kernel, wat twee belangrijke operaties mogelijk maakt:

1. Voorwaartse Sampling: Samples worden gegenereerd via Markov Chain Monte Carlo (MCMC)-methoden (bijv. geodetische random walks) met behulp van de door de PINN geëvalueerde dichtheid.
2. Score Evaluatie: De conditionele score $\nabla_x \log p_{t|0}$ wordt berekend door de PINN te differentiëren, wat vervolgens wordt gebruikt om het reverse-time diffusiemodel te trainen via Denoising Score Matching (DSM).

3. Belangrijkste Bijdragen

• Algemene PDE-gebaseerde Surrogaat: Het artikel introduceert een methode om heat kernels op willekeurige manifolds te benaderen door de warmtevergelijking numeriek op te lossen, waarmee de noodzaak voor closed-form oplossingen of restrictieve spectrale aannames wordt omzeild.
• Coördinaat-vrije Afleiding: Het framework biedt een systematische manier om de warmtevergelijking af te leiden voor manifolds die gedefinieerd zijn als submanifolds of quotients van Euclidische ruimtes, wat het toepasbaar maakt op complexe geometrieën zoals $S^2$, $SO(3)$, $SPD(n)$ en permutatie-quotiented puntwolken.
• Stabiliteitsanalyse: De auteurs bieden een theoretische rechtvaardiging (Theorema 4.1) die aantoont dat het oplossen van de warmtevergelijking vanuit een benaderde initiële conditie een oplossing oplevert met een gebonden fout voor alle $t \ge t_0$, wat de betrouwbaarheid van de PINN-surrogaat waarborgt.
• Permutatie-invariante Diffusie: De methode wordt toegepast op de quotient-manifold $\mathbb{R}^{k \times n}/S_n$, wat diffusieprocessen mogelijk maakt die expliciet atomische permutaties wegmoduleren voor molecuulgeneratie, een setting waarin standaard Euclidische diffusie de structurele symmetrieën niet respecteert.

4. Experimentele Resultaten

De methode werd geëvalueerd op diverse manifolds en datasets:

• Sferische Data ($S^2$): Dichtheidschatting van geofysische gebeurtenissen (vulkanen, aardbevingen, overstromingen, bosbranden). De voorgestelde methode behaalde competitieve of superieure test log-likelihoods vergeleken met Riemannian Score-based Generative Models (RSGM) die steunen op Varadhan-benaderingen en spectrale methoden.
• Rotatie Data ($SO(3)$): Synthetische mixture modeling toonde verbeterde log-likelihoods vergeleken met RSGM-baselines.
• Positief-Definiete Matrices ($SPD(n)$):
  • Verkeersanalyse ($SPD(10)$): Conditionele generatie van verkeersstroommatrices vanuit NYC taxi-data. De methode verminderde de Frobenius-afstand tot de grondwaarheid vergeleken met SPD-DDPM.
  • Breinconnectiviteit ($SPD(n)$): Conditionele generatie van EEG-connectiviteitsmatrices. De methode behaalde sterke precisie en F1-scores, waarbij zij concurreert met DiffeoCFM (dat diffeomorfe mappingen naar de Euclidische ruimte gebruikt) en presteert boven LowTriDDPM.
• Molecuulgeneratie ($\mathbb{R}^{k \times n}/S_n$): Generatie van moleculen op de permutatie-quotient ruimte. De methode produceerde moleculen met een hogere validiteit en stabiliteit dan een variance-exploding (VE) EDM-baseline, hoewel met een lagere nieuwheid (novelty), wat een bekend validiteit-nieuwheid trade-off reflecteert.

Validatie: De auteurs hebben de PINN-surrogaten gevalideerd door de geleerde log-densiteiten en scores te vergelijken met referentie-benaderingen (waar beschikbaar) en door genormaliseerde PDE-residuen en boundary-fouten te rapporteren, wat bevestigt dat de surrogaten aan de warmtevergelijking voldoen.

5. Betekenis en Claims

Het artikel claimt dat PDE-gebaseerde heat-kernel benadering een praktisch en algemeen alternatief biedt voor bestaande methoden voor Riemanniaanse diffusie. Door de heat kernel te behandelen als de oplossing van een bekende PDE, breidt de aanpak drift-vrije Riemanniaanse diffusie uit naar manifolds waar analytische kernels niet beschikbaar zijn.

De auteurs positioneren dit werk als een "solve-the-heat-equation" perspectief dat de toepasbaarheid van manifold diffusiemodellen vergroot voorbij eenvoudige geometrieën. Zij erkennen echter bescheiden de beperkingen:

• Computationele Kosten: PINN-training vereist aanzienlijke middelen vanwege hoge-orde autodifferentiatie en collocatie-evaluaties.
• Stabiliteit: Training kan instabiel zijn in hoge dimensies of voor complexe geometrieën, wat mogelijk sterkere inductieve biases vereist.
• Complexiteit: De pipeline is complexer dan alternatieven zoals Flow Matching, die de schatting van de heat kernel volledig omzeilen.

Het werk claimt niet de vervanging van flow matching, maar biedt eerder een noodzakelijke backend voor score-gebaseerde modellen op manifolds waar de afhankelijkheid van flow matching van geodetische interpolatie of exponentiële maps ontoereikend kan zijn, of waar de specifieke structuur van de heat kernel vereist is.