Operator learning for solving Fokker-Planck equations with various initial conditions

2026-06-08

AI Summary

This work addresses the numerical challenges in solving the Fokker–Planck equation under Dirac delta initial conditions and small time steps by proposing a physics-informed neural network framework based on conditional normalizing flows. The approach leverages the Chapman–Kolmogorov equation to reformulate the solution operator learning task as an approximation of transition probability density functions, using the analytical solution of a linearized stochastic differential equation as the base distribution for the normalizing flow. To mitigate instability at small time scales, a time-weighted loss function is introduced, while analytical priors are incorporated to eliminate the initial singularity. Experimental results demonstrate that the method achieves high accuracy, strong generalization, and robustness in learning solution operators across diverse complex initial conditions and small time steps.

Abstract

The Fokker-Planck equation (FPE) plays a pivotal role in describing the time evolution of probability density functions (PDFs) for systems governed by stochastic dynamics. In this work, we propose a conditional normalizing flow-based physics-informed neural network (PINN) framework for efficiently approximating the solution operator of the FPE for a whole range of initial conditions. Leveraging the Chapman-Kolmogorov equation for Markovian stochastic processes, the problem is reformulated into approximating a transition PDF starting at initial time from a Dirac mass centered at an arbitrary point. The PDF of an associated linearized stochastic differential equation (SDE) is employed as the base distribution for the normalizing flow, providing a good approximation of the target PDF, especially for small times, and thereby avoiding the singularity of the map associated with the Dirac delta initial distribution. Furthermore, a time-weighted loss function is introduced to mitigate numerical instabilities arising at small times, achieving a balance between causality and training difficulty as time progresses. A variety of numerical experiments are presented to illustrate the effectiveness and robustness of the proposed method.

Problem

• Fokker-Planck equation
• solution operator
• initial conditions
• probability density function
• stochastic dynamics

Innovation

• conditional normalizing flow
• physics-informed neural network
• Fokker-Planck equation
• solution operator learning
• Chapman-Kolmogorov equation