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

TL;DR

This paper introduces a neural pushforward method for solving the Fokker-Planck equation on embedded Riemannian manifolds, enabling mesh-free, geometry-aware probability distribution modeling with demonstrated effectiveness on curved spaces.

Contribution

It extends the neural pushforward approach to Riemannian manifolds, incorporating geometric constraints and deriving differential operators directly from the embedding.

Findings

Successfully captures multimodal invariant distributions on curved spaces

Provides a mesh-free, chart-free algorithm for Fokker-Planck equations on manifolds

Demonstrates effectiveness on a double-well problem on the two-sphere

Abstract

In this paper, we extend the Weak Adversarial Neural Pushforward Method to the Fokker--Planck equation on compact embedded Riemannian manifolds. The method represents the solution as a probability distribution via a neural pushforward map that is constrained to the manifold by a retraction layer, enforcing manifold membership and probability conservation by construction. Training is guided by a weak adversarial objective using ambient plane-wave test functions, whose intrinsic differential operators are derived in closed form from the geometry of the embedding, yielding a fully mesh-free and chart-free algorithm. Both steady-state and time-dependent formulations are developed, and numerical results on a double-well problem on the two-sphere demonstrate the capability of the method in capturing multimodal invariant distributions on curved spaces.