Score-Based Physics-Informed Neural Networks for High-Dimensional Fokker-Planck Equations
Zheyuan Hu, Zhongqiang Zhang, George Em Karniadakis, Kenji Kawaguchi
TL;DR
This work addresses the difficulty of solving high-dimensional Fokker-Planck equations by introducing a score-based SDE solver that learns the score function $s_t(x)=\nabla_x \log p_t(x)$ to recover the log-likelihood $q_t(x)$ and the density $p_t(x)$ via a LL ODE. It proposes three fitting strategies—Score Matching (SM), Sliced Score Matching (SSM), and Score-PINN—each enabling a two-stage workflow: learn the score from SDE trajectories or the Score PDE, then solve the LL ODE for LL and PDF. Through extensive experiments on anisotropic OU processes, anisotropic Brownian motion, non-Gaussian initial conditions (Cauchy, Laplace), and anisotropic Log-normal GBM, the results demonstrate that SM and SSM are fast and robust in moderate dimensions, while Score-PINN achieves superior accuracy by leveraging higher-order PDE information, particularly in challenging settings. The approach effectively mitigates CoD, delivering stable performance as dimension grows and offering faster, more reliable LL inference and SDE sampling than direct LL-PINN or Monte Carlo methods in many high-dimensional scenarios.
Abstract
The Fokker-Planck (FP) equation is a foundational PDE in stochastic processes. However, curse of dimensionality (CoD) poses challenge when dealing with high-dimensional FP PDEs. Although Monte Carlo and vanilla Physics-Informed Neural Networks (PINNs) have shown the potential to tackle CoD, both methods exhibit numerical errors in high dimensions when dealing with the probability density function (PDF) associated with Brownian motion. The point-wise PDF values tend to decrease exponentially as dimension increases, surpassing the precision of numerical simulations and resulting in substantial errors. Moreover, due to its massive sampling, Monte Carlo fails to offer fast sampling. Modeling the logarithm likelihood (LL) via vanilla PINNs transforms the FP equation into a difficult HJB equation, whose error grows rapidly with dimension. To this end, we propose a novel approach utilizing a score-based solver to fit the score function in SDEs. The score function, defined as the gradient of the LL, plays a fundamental role in inferring LL and PDF and enables fast SDE sampling. Three fitting methods, Score Matching (SM), Sliced SM (SSM), and Score-PINN, are introduced. The proposed score-based SDE solver operates in two stages: first, employing SM, SSM, or Score-PINN to acquire the score; and second, solving the LL via an ODE using the obtained score. Comparative evaluations across these methods showcase varying trade-offs. The proposed method is evaluated across diverse SDEs, including anisotropic OU processes, geometric Brownian, and Brownian with varying eigenspace. We also test various distributions, including Gaussian, Log-normal, Laplace, and Cauchy. The numerical results demonstrate the score-based SDE solver's stability, speed, and performance across different settings, solidifying its potential as a solution to CoD for high-dimensional FP equations.