Self-Consistent Probability Flow for High-Dimensional Fokker-Planck Equations

TL;DR

This work tackles the curse of dimensionality in solving high-dimensional Fokker-Planck equations by reformulating the second-order FP PDE into a first-order Probability Flow ODE and solving it via a residual-based objective. The method, Self-Consistent Probability Flow (SCPF), leverages Continuous Normalizing Flows with Hutchinson Trace Estimation to compute density divergences in linear time, and employs generative adaptive sampling to focus collocation points on high-probability regions. Theoretical results connect the residual minimization to convergence to the FP solution and provide a Wasserstein error bound that scales with the adaptive residual, justifying the sampling strategy. Empirically, SCPF achieves high accuracy and constant or near-constant wall-clock time up to 100 dimensions across various benchmarks, including time-varying diffusion and non-Gaussian distributions, substantially outperforming the baseline tKRnet in both accuracy and efficiency.

Abstract

Solving high-dimensional Fokker-Planck (FP) equations is a challenge in computational physics and stochastic dynamics, due to the curse of dimensionality (CoD) and the bottleneck of evaluating second-order diffusion terms. Existing deep learning approaches, such as Physics-Informed Neural Networks (PINNs), face computational challenges as dimensionality increases, driven by the $O(D^2)$ complexity of automatic differentiation for second-order derivatives. While recent probability flow approaches bypass this by learning score functions or matching velocity fields, they often involve serial computational operations or depend on sampling efficiency in complex distributions. To address these issues, we propose the Self-Consistent Probability Flow (SCPF) method. We reformulate the second-order FP equation into an equivalent first-order deterministic Probability Flow ODE (PF-ODE) constraint. Unlike score matching or velocity matching, SCPF solves this problem by minimizing the residual of the PF-ODE continuity equation, which avoids explicit Hessian computation. We leverage Continuous Normalizing Flows (CNF) combined with the Hutchinson Trace Estimator (HTE) to reduce the training complexity to linear scale $O(D)$, achieving an effective $O(1)$ wall-clock time on GPUs. To address data sparsity in high dimensions, we apply a generative adaptive sampling strategy and theoretically prove that dynamically aligning collocation points with the evolving probability mass is a necessary condition to bound the approximation error. Experiments on diverse benchmarks -- ranging from anisotropic Ornstein-Uhlenbeck (OU) processes and high-dimensional Brownian motions with time-varying diffusion terms, to Geometric OU processes featuring non-Gaussian solutions -- demonstrate that SCPF effectively mitigates the CoD, maintaining high accuracy and constant computational cost for problems up to 100 dimensions.

Figures (7)

• Figure 5.1: Loss convergence during the training process, one-dimensional test problem.
• Figure 5.2: Error convergence during the training process, one-dimensional test problem.
• Figure 5.3: Exact solution, SCPF solution (top) and tKRnet solution (bottom), one-dimensional test problem.
• Figure 5.4: Error convergence during the training process, two-dimensional unimodal test problem.
• Figure 5.5: Exact solution, SCPF Abs.Error and tKRnet Abs.Error, two-dimensional unimodal test problem.

Theorems & Definitions (11)

• Proposition 4.1: Validity of the First-Order Residual
• proof
• Lemma 4.2: Coupling Lemma
• Lemma 4.3: Differential Inequality Lemma
• proof
• Lemma 4.4: Stability of the Score Function
• proof
• Remark 4.5: Justification of assumptions
• Theorem 4.6: Error Bound via Residual
• proof