Neural Pushforward Samplers for the Fokker-Planck Equation on Embedded Riemannian Manifolds

Titel: Neural Pushforward Samplers voor de Fokker-Planck-vergelijking op Inbedde Riemannse Variëteiten

Auteurs: Andrew Qing He en Wei Cai
Datum: Maart 2026

---

1. Probleemstelling

Veel stochastische processen in moleculaire dynamica, robotica en klimatologie evolueren niet in een vlakke Euclidische ruimte, maar op een constraint-oppervlak: een Riemannse variëteit $(M, g)$, zoals een bol, een torus of een Stiefel-variëteit. De tijdsontwikkeling van de waarschijnlijkheidsdichtheid $\rho$ op deze variëteit wordt beschreven door de Fokker-Planck-vergelijking (FPE):

$$\frac{\partial \rho}{\partial t} = \frac{\sigma^2}{2} \Delta_g \rho - \text{div}_g(\rho b)$$

Waarbij $\Delta_g$ de Laplace-Beltrami-operator is, $b$ een tangentiële drift is, en $\sigma$ de diffusiecoëfficiënt.

Uitdagingen:

• Dimensieprobleem: Klassieke numerieke methoden (zoals eindige elementen of differenties) vereisen een expliciete triangulatie van $M$. Dit wordt prohibitief duur in matige tot hoge dimensies.
• Koordinatenprobleem: Het gebruik van lokale kaarten (charts) introduceert singulariteiten en complexiteit, vooral voor compacte variëteiten zonder rand.
• Bestaande methoden: Normalizing flows vereisen inverteerbare kaarten en een gelijke dimensie van de basisruimte en de variëteit, wat de representatieve capaciteit beperkt.

2. Methodologie: Weak Adversarial Neural Pushforward (WANPF)

De auteurs breiden de bestaande "Weak Adversarial Neural Pushforward Method" (WANPF) uit van de vlakke ruimte naar Riemannse variëteiten. De kern van de methode is het benaderen van de oplossing $\rho$ niet als een puntsgewijze dichtheid, maar als een pushforward-maatschappij gegenereerd door een neurale netwerk $F_\vartheta$.

A. Architectuur en Manifold-beperking

In plaats van een inverteerbare stroom te gebruiken, definieert de methode een generator $F_\vartheta$ die een eenvoudige basisverdeling (bijv. een Gaussische ruis $r \in \mathbb{R}^d$) afbeeldt naar de variëteit $M$.

• Retractie: Om te garanderen dat alle monsters op de variëteit blijven, wordt een retractie $\Pi_M$ toegepast na een onbeperkt MLP:
  $$F_\vartheta(r) = \Pi_M(\tilde{F}_\vartheta(r))$$
  Voor een bol $S^{n-1}$ is dit bijvoorbeeld normalisatie: $v \mapsto v/\|v\|$. Dit zorgt ervoor dat de massa-conservering en het verblijf op de variëteit "by construction" worden afgedwongen.
• Niet-inverteerbaarheid: De dimensie van de ruis $d$ hoeft niet gelijk te zijn aan $\dim(M)$. Dit biedt meer representatieve capaciteit dan traditionele flows.

B. Zwakke Formulering en Adversariaal Training

De methode gebruikt de zwakke formulering van de FPE. Door partiële integratie (waarbij randtermen verdwijnen door compactheid) wordt de operator overgebracht naar de testfuncties $f$.
De voorwaarde voor een stationaire oplossing is:
$$\mathbb{E}_{x \sim \rho} [L^{ss} f(x)] = 0$$
Waarbij $L^{ss}$ de achterwaartse Kolmogorov-operator is.

• Testfuncties: Er worden omgevende vlakke golven (ambient plane waves) gebruikt als testfuncties: $f(x) = \sin(w \cdot x + b)$.
• Adversariaal doel: Een min-max spel waarbij de generator de loss minimaliseert (zodat de verwachting nul is) en de testfuncties (de "adversary") de parameters $w, b$ maximaliseren om de afwijking van de FPE te detecteren.

C. De Laplace-Beltrami Operator in Omgeving

Het cruciale technische inzicht is dat de Laplace-Beltrami-operator op een ingebedde variëteit exact kan worden berekend in de omgevende ruimte $\mathbb{R}^n$ zonder automatische differentiatie door de intrinsieke geometrie.
Voor een testfunctie $f$ geldt:
$$\Delta_g f(x) = \nabla^2 f(x) : P(x) - \nabla f(x) \cdot H(x)$$
Waarbij:

• $P(x) = I - xx^\top$ de projectie op de raakruimte is.
• $H(x)$ de gemiddelde krommingsvector is.
  Voor vlakke golven kunnen deze termen in gesloten vorm worden berekend, wat training zonder mesh en zonder dure autodiff door de geometrie mogelijk maakt.

3. Belangrijkste Bijdragen

1. Uitbreiding naar Riemannse Variëteiten: Een volledige formulering van WANPF voor stationaire en tijdsafhankelijke FPE's op compacte, ingebedde variëteiten.
2. Gesloten Formules: Afleiding van expliciete formules voor de actie van de Laplace-Beltrami-operator op vlakke golven voor de eenheidsbol $S^{n-1}$ en de platte torus $T^n$.
3. Retractie-gebaseerde Generator: Een architectuur die monsters dwingt op de variëteit te blijven via een projectie, waardoor de initialisatie en randvoorwaarden automatisch worden voldaan.
4. Mesh-vrije Training: De methode vereist geen discretisatie van de variëteit, geen kaarten (charts) en geen berekening van Jacobianen.

4. Numerieke Resultaten

De methode werd getest op een stationaire FPE met een dubbelputpotentiaal op de bol $S^2$.

• Opstelling: Een potentiaal $V(x,y,z) = \alpha(x^2-1)^2 + \beta z^2$ creëert twee minima bij $(\pm 1, 0, 0)$ op de evenaar van de bol.
• Resultaat: De getrainde generator produceerde een verdeling die de massa correct concentreerde rond de twee minima, in overeenstemming met de theoretische Gibbs-verdeling ($\rho \propto e^{-2V/\sigma^2}$).
• Observaties:
  • De massa was sterk onderdrukt bij de polen (door de $\beta z^2$ term).
  • De symmetrie tussen de twee putten werd behouden.
  • De training convergerde na ongeveer 2000 stappen naar een stationair punt in het min-max spel.

5. Betekenis en Toekomstperspectief

Deze paper biedt een krachtige, schaalbare oplossing voor het oplossen van FPE's in hoge dimensies op complexe meetkundige structuren.

• Efficiëntie: Door het gebruik van gesloten formules voor de Laplace-Beltrami-operator en het vermijden van meshes, is de methode veel efficiënter dan traditionele PDE-oplossers voor hoge dimensies.
• Generaliteit: De aanpak is toepasbaar op elke variëteit waarvoor de projectie $P(x)$ en krommingsvector $H(x)$ analytisch of numeriek berekend kunnen worden (bijv. Stiefel-variëteiten, Lie-groepen).
• Toekomst: De methode kan worden uitgebreid naar niet-isotrope diffusie, fractionele FPE's, en andere meetkundige structuren zoals SO(3) voor robotica.

Conclusie: De auteurs hebben een mesh-vrije, kaart-vrije en Jacobian-vrije methode ontwikkeld die de complexiteit van het oplossen van stochastische processen op Riemannse variëteiten aanzienlijk verlaagt, met succesvolle toepassing op de bol $S^2$.