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

1. Problem Statement

The paper addresses the numerical solution of the Fokker–Planck Equation (FPE) governing the time evolution of probability densities on compact, smoothly embedded Riemannian manifolds $(M, g) \subset \mathbb{R}^n$. The equation is given by:
$$\frac{\partial \rho}{\partial t} = \frac{\sigma^2}{2} \Delta_g \rho - \text{div}_g(\rho b)$$
where $\Delta_g$ is the Laplace–Beltrami operator, $b$ is a tangential drift vector field, and $\sigma$ is the diffusion coefficient.

Challenges:

• Curse of Dimensionality: Traditional grid-based methods (finite difference/element) fail in moderate-to-high dimensions.
• Geometric Complexity: Constructing explicit triangulations or charts for complex manifolds (e.g., spheres, Stiefel manifolds) is computationally expensive and often intractable.
• Jacobian Computation: Standard normalizing flow methods require invertible maps and explicit Jacobian determinants, which are difficult to compute on curved manifolds without singularities.

2. Methodology: Weak Adversarial Neural Pushforward (WANPF)

The authors extend the Weak Adversarial Neural Pushforward Method (previously developed for flat Euclidean space) to Riemannian manifolds. The core philosophy is to represent the solution distribution $\rho$ not as a pointwise density, but as the pushforward of a simple base distribution (e.g., Gaussian) through a neural network generator.

A. The Pushforward Architecture

Instead of learning a map $F_\vartheta: \mathbb{R}^d \to M$ directly, the network learns an unconstrained map $\tilde{F}_\vartheta: \mathbb{R}^d \to \mathbb{R}^n$ and composes it with a manifold retraction $\Pi_M$:
$$F_\vartheta(r) = \Pi_M(\tilde{F}_\vartheta(r))$$

• Retraction ($\Pi_M$): A projection operator that maps points in the ambient space $\mathbb{R}^n$ back onto the manifold $M$ (e.g., normalization for spheres, wrapping for tori).
• Constraint Enforcement: This ensures that all generated samples lie exactly on $M$ by construction, satisfying the manifold constraint without requiring invertibility or Jacobian calculations.

B. Weak Formulation and Adversarial Training

The method utilizes the weak formulation of the FPE. By integrating against test functions $f$ and applying integration by parts, the differential operators act on the test functions rather than the unknown density $\rho$.

• Objective: A min-max game where the generator $F_\vartheta$ minimizes the residual of the weak form, while adversarial test functions maximize it.
• Loss Function: The training objective is to minimize the squared expectation of the backward Kolmogorov operator applied to test functions:
  $$\min_\vartheta \max_{\{f_k\}} \sum_k \left( \mathbb{E}_{r \sim P_{base}} [L_{ss} f_k(F_\vartheta(r))] \right)^2$$

C. Ambient Plane-Wave Test Functions

A critical innovation is the use of ambient plane waves defined on the embedding space $\mathbb{R}^n$:
$$f^{(k)}(x) = \sin(w^{(k)} \cdot x + b^{(k)})$$
The authors derive a closed-form expression for the Laplace–Beltrami operator $\Delta_g$ acting on these plane waves using only ambient quantities:
$$\Delta_g f^{(k)}|_M = -\sin(\phi) (w^\top P(x) w) - \cos(\phi) (H(x) \cdot w)$$
where:

• $P(x)$ is the orthogonal projection onto the tangent space $T_x M$.
• $H(x)$ is the mean-curvature vector of the embedding.
• $\phi = w \cdot x + b$.

Significance: This formula allows the computation of $\Delta_g f$ without automatic differentiation through the intrinsic geometry or mesh. It relies solely on algebraic operations on the ambient vectors $x$ and $w$.

3. Key Contributions

1. Manifold Extension of WANPF: A complete framework for solving steady-state and time-dependent FPEs on embedded Riemannian manifolds using neural pushforward maps.
2. Retraction-Based Generator: A novel architecture that enforces manifold membership via projection ($\Pi_M$), avoiding the need for invertible flows or local charts.
3. Closed-Form Laplace–Beltrami Operators: Derivation of explicit formulas for $\Delta_g$ acting on plane waves for general embedded manifolds, specifically for the Unit Sphere ($S^{n-1}$) and Flat Torus ($T^n$).
4. Mesh-Free and Jacobian-Free Training: The method requires no spatial discretization, no chart atlases, and no Jacobian determinant calculations, making it scalable to high dimensions.
5. Numerical Validation: Successful application to a double-well FPE on the 2-sphere ($S^2$), demonstrating the ability to capture multimodal stationary distributions.

4. Numerical Results

The authors tested the method on a double-well FPE on $S^2$ (the unit sphere in $\mathbb{R}^3$).

• Setup: The potential $V(x) = \alpha(x^2-1)^2 + \beta z^2$ creates two minima at $(\pm 1, 0, 0)$. The drift is derived from this potential.
• Performance:
  • The trained generator successfully concentrated probability mass near the two potential minima $(\pm 1, 0, 0)$.
  • The distribution correctly suppressed mass near the poles (due to the $z^2$ term), matching the theoretical Gibbs distribution $\rho_\infty \propto e^{-2V/\sigma^2}$.
  • The adversarial loss converged to a stationary level, and the generated samples exhibited the expected symmetry and confinement.
• Efficiency: The training was mesh-free and utilized closed-form derivatives, avoiding the computational overhead of geometric differentiation.

5. Significance and Impact

• Scalability: By avoiding grids and charts, this method overcomes the curse of dimensionality, making it applicable to high-dimensional manifolds (e.g., $S^{n-1}$ for $n \ge 10$) where traditional methods fail.
• Geometric Flexibility: The reliance on ambient space representations ($P(x)$ and $H(x)$) allows the method to handle complex geometries (Stiefel manifolds, Lie groups) as long as their projection and curvature properties are known.
• Theoretical Connection: The work bridges neural PDE solvers with Geometric Measure Theory, specifically utilizing the identity that relates the Laplace–Beltrami operator to ambient quantities (varifold theory).
• Future Applications: The framework is directly applicable to problems in molecular dynamics (protein folding on manifolds), robotics (motion planning on configuration spaces), and directional statistics.

In summary, this paper presents a robust, mesh-free numerical framework for solving stochastic differential equations on curved spaces, leveraging the geometry of the embedding space to simplify the computation of intrinsic differential operators.