Learning Neural Pushforward Samplers for Distributions from Fokker-Planck Equations by Weak Adversarial Training

Abstract

This paper presents a new method for solving Fokker-Planck equations (FPE) by learning a neural sampler for the distribution given by the FPE via an adversarial training based on a weak formulation of the FPE where the adjoint operator of FPE acts on the test function. Such a weak formulation transforms the PDE solution problem into a Monte Carlo importance sampling problem where the FPE solution-distribution is learned through a neural pushforward map, avoiding some of the limitations of direct PDE based methods. Moreover, by using simple plane-wave test functions, derivatives on the test functions can be explicitly computed. This approach produces a natural importance sampling strategy for the FPE solution distribution with probability conservation, from which the FPE solution can be easily constructed.

Figures (16)

• Figure 1: Numerical results for the one-dimensional Fokker-Planck equation. (a) The pushforward network $F_{{\boldsymbol{\vartheta}}}$ successfully learns to transform a uniform distribution into the target Gaussian distribution $\mathcal{N}(2,1)$, as the learned PDF shows excellent agreement with the analytical solution. (b) The training dynamics show a rapid and stable decrease in the loss value $\mathcal{L}_{\text{total}}$, indicating the model is effectively minimizing the violation of the PDE's weak form.
• Figure 2: Numerical results for the one-dimensional double-peak Fokker-Planck equation. (a) The learned distribution successfully captures the bimodal structure with peaks at $x \approx \pm 1$, showing excellent agreement with the analytical Boltzmann distribution. (b) The adversarial training exhibits stable convergence with the loss decreasing over several orders of magnitude.
• Figure 3: Training results for the two-dimensional double-peak Fokker-Planck equation. Left: Scatter plot of samples generated by the learned pushforward map, showing clear clustering at the target peaks $(-1,-1)$ and $(1,1)$ (marked with red crosses). Middle: Density heatmap revealing two distinct high-probability regions at the correct locations. Right: Training loss convergence over 100,000 epochs, demonstrating initial rapid decrease followed by a stable regime, with a transient instability around epoch 50,000--60,000.
• Figure 4: Comparison of analytical and learned probability distributions for the 2D double-peak problem. Left: Analytical Boltzmann distribution $\rho_{\text{true}}(\boldsymbol{x}) \propto \exp(-2V(\boldsymbol{x})/\sigma^2)$ with $\sigma=0.4$. Right: Learned distribution obtained from the pushforward network. The two distributions show excellent agreement in both peak locations and overall shape, validating the method's accuracy in capturing multimodal distributions.
• Figure 5: Training results for the two-dimensional ring potential with rotational drift. Left: Scatter plot of generated samples clearly showing concentration along the target ring at $r = r_0 = 2$ (red dashed circle). Middle: Density heatmap revealing the learned ring structure. Right: Training loss convergence showing initial adversarial dynamics followed by stable convergence.