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

2025-12-22

AI Summary

To address the curse of dimensionality and the computational bottleneck of Hessian evaluation in solving high-dimensional Fokker–Planck (FP) equations, this paper proposes the Self-Consistent Probability Flow (SCPF) framework. SCPF reformulates the FP equation as a first-order ordinary differential equation (ODE) governing the probability current, enabling efficient, Hessian-free solution via residual minimization of the continuity equation. We introduce a generative adaptive sampling strategy that dynamically aligns collocation points and rigorously prove its necessity for bounded approximation error. To our knowledge, this is the first implementation achieving near O(1) wall-clock time for high-dimensional FP solving on GPUs. The method integrates continuous normalizing flows (CNFs), the Hutchinson trace estimator (HTE), and residual-driven optimization. Evaluated on benchmark problems—including 100-dimensional anisotropic Ornstein–Uhlenbeck (OU), time-varying diffusion Brownian motion, and non-Gaussian geometric OU processes—SCPF maintains high accuracy while reducing computational complexity to O(D), substantially alleviating the curse of dimensionality.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Probabilistic Circuits and Graphical Models

Application Category

• Security and Privacy: Large-scale security measurements
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Solves high-dimensional Fokker-Planck equations efficiently by avoiding second-order derivative computations.
• Reduces training complexity to linear scale using continuous normalizing flows and trace estimation.
• Addresses data sparsity in high dimensions with generative adaptive sampling for error control.

Innovation

• Reformulates FP equation into first-order PF-ODE to avoid Hessian computation
• Uses CNF with HTE to reduce training complexity to linear O(D) scale
• Applies generative adaptive sampling to bound error in high dimensions