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

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.

Figures (2)

• Figure 1: Left: the double-well potential $V(x,y,z) = \alpha(x^2-1)^2 + \beta z^2$ on $S^2$, shown in longitude--latitude projection. The two potential minima ($V = 0$) are located at longitude $0^\circ$ and $\pm 180^\circ$ on the equator, corresponding to $(\pm 1, 0, 0)$. Right: density histogram of $8{,}000$ samples generated by the trained neural pushforward map in the same coordinate system. Mass concentrates near both well bottoms and is absent near the poles, consistent with the Gibbs distribution $\rho_\infty \propto e^{-2V/\sigma^2}$ for $\sigma = 0.5$.
• Figure 2: Training and sampling diagnostics for the $S^2$ double-well experiment. Left: adversarial loss versus training step (log scale). The loss rises during early generator exploration and then decreases as training converges to a stationary min-max point near step $2{,}000$. Center: $8{,}000$ generated samples on $S^2$, coloured by their $x$-coordinate value (red near $x=+1$, blue near $x=-1$); gold stars mark the two well bottoms $(\pm 1, 0, 0)$. The generator correctly concentrates mass around both minima symmetrically. Right: Mollweide equal-area projection of the same samples; dashed gold lines indicate the longitudes of the well bottoms. Samples cluster near $0^\circ$ and $\pm 180^\circ$ longitude, confirming two-well capture, with near-uniform spreading in longitude about each minimum and strong equatorial confinement consistent with the $\beta z^2$ term.

Theorems & Definitions (3)

• Remark 4.1: Flat-space recovery
• Remark 4.2: Non-eigenfunction character
• Remark 5.1: Parametrization alternative