A deep learning framework for jointly solving transient Fokker-Planck equations with arbitrary parameters and initial distributions

Xiaolong Wang; Jing Feng; Qi Liu; Chengli Tan; Yuanyuan Liu; Yong Xu

physics.comp-phApr 7, 2026

A deep learning framework for jointly solving transient Fokker-Planck equations with arbitrary parameters and initial distributions

Authors: Xiaolong Wang, Jing Feng, Qi Liu, Chengli Tan, Yuanyuan Liu, Yong Xu

Organizations: School of Mathematics and Statistics, Shaanxi Normal University, Xi’an, 710119, China · School of Mathematics and Statistics, Northwestern Polytechnical University, Xi’an, 710129, China · MOE Key Laboratory for Complexity Science in Aerospace, Northwestern Polytechnical University, Xi’an, 710072, China · School of Science, Xi’an University of Posts and Telecommunications, Xi’an, 710121, China ·

Abstract

Efficiently solving the Fokker-Planck equation (FPE) is central to analyzing complex parameterized stochastic systems. However, current numerical methods lack parallel computation capabilities across varying conditions, severely limiting comprehensive parameter exploration and transient analysis. This paper introduces a deep learning-based pseudo-analytical probability solution (PAPS) that, via a single training process, simultaneously resolves transient FPE solutions for arbitrary multi-modal initial distributions, system parameters, and time points. The core idea is to unify initial, transient, and stationary distributions via Gaussian mixture distributions (GMDs) and develop a constraint-preserving autoencoder that bijectively maps constrained GMD parameters to unconstrained, low-dimensional latent representations. In this representation space, the panoramic transient dynamics across varying initial conditions and system parameters can be modeled by a single evolution network. Extensive experiments on paradigmatic systems demonstrate that the proposed PAPS maintains high accuracy while achieving inference speeds four orders of magnitude faster than GPU-accelerated Monte Carlo simulations. This efficiency leap enables previously intractable real-time parameter sweeps and systematic investigations of stochastic bifurcations. By decoupling representation learning from physics-informed transient dynamics, our work establishes a scalable paradigm for probabilistic modeling of multi-dimensional, parameterized stochastic systems.

Figures & tables

Figure 1: The architecture of the TPAPS $q(\mathbf{x},t;\bm{\Theta},p_{\text{I}})$ for jointly solving parameterized FPEs with diverse initial distributions and control parameters.

ResNet ID	EMB-RN	REP-RN	EVO-RN	DEC-RN
Input	$\{(\mu_{ij},\sigma_{ij})\}_{i=1}^{D_{\text{STA}}}$	Embedding $\mathbf{h}$	$(t,\bm{\Theta},\mathbf{H}_{\text{I}})$	Transient $\mathbf{H}_{t}(\bm{\Theta},\mathbf{H}_{\text{I}})$
$D_{\text{In}}$ / $D_{\text{Mid}}$ / $D_{\text{Out}}$ of RB#1	$2D_{\text{STA}}$ /50/50	$D_{\text{EMB}}$ /100/100	$1+D_{\text{PAR}}+D_{\text{REP}}$ /100/100	$D_{\text{REP}}$ /100/100
$D_{\text{In}}$ / $D_{\text{Mid}}$ / $D_{\text{Out}}$ of RB#2	50/50/50	100/100/100	100/100/100	100/100/100
$D_{\text{In}}$ / $D_{\text{Mid}}$ / $D_{\text{Out}}$ of RB#3	50/50/ $D_{\text{EMB}}$	100/100/ $D_{\text{REP}}$	100/100/100	100/100/100
$D_{\text{In}}$ / $D_{\text{Mid}}$ / $D_{\text{Out}}$ of RB#4			100/100/100	100/100/100
$D_{\text{In}}$ / $D_{\text{Mid}}$ / $D_{\text{Out}}$ of RB#5			100/100/100	100/100/100

Table 1: The architectures of the ResNets in the TPAPS

Figure 2: Visualization of the GMD autoencoder. (a) Six GMDs and their reconstructions. (b) The PCA and distance visualizations of the 50-D embedding space. Each point represents a 2-D Gaussian distribution in Eq. ( 45 ). The three lines consist of continuously changing distributions, where $\bm{\mu}_{1}=[-2,-2]^{\top}$ , $\bm{\mu}_{2}=[2,2]^{\top}$ , and $\mathbf{I}$ is the 2-D identity matrix.

Figure 3: The transient solutions of the 1-D system ( 49 ) by TPAPS (colored circles) and MCS (black lines) with different initial distributions and system parameters.

Figure 4: The training losses and test errors of the TPAPS on the 1-D system. The boxplots in (b) are summarized on 2,500 transient solutions with random initial distributions and system parameters.

Time	$10^{3}$ batches	$10^{4}$ batches	$10^{5}$ batches	$10^{6}$ batches
$t=0.5$	0.4257/0.3561/0.3161	0.2062/0.1906/0.1476	0.0595/0.0923/0.0345	0.0212/0.0477/0.0133
$t=1.5$	0.3504/0.3701/0.2056	0.1794/0.2354/0.0798	0.0607/0.1291/0.0232	0.0217/0.0645/0.0102
$t=3.0$	0.3637/0.3745/0.2050	0.1648/0.2340/0.0685	0.0619/0.1410/0.0230	0.0234/0.0717/0.0098

Table 2: The mean, SD, and median (MEAN/SD/MED) of the L1 errors between the TPAPS and MCS results across 2,500 random transient solutions of the 1-D system ( 49 ) at various training stages.

Figure 5: Bifurcation study of the 1-D system ( 49 ) by using TPAPS and MCS. The bifurcation parameters are (a) the combination weight $\lambda$ in the initial distribution, (b) the system parameter $a$ , and (c) the noise amplitude $\sigma$ . The resolution of each figure is $100\times 100$ .

Figure 6: Four transient solutions of the 2-D subcritical system ( 52 ) by TPAPS and MCS with different initial distributions and system parameters. The L1 errors are displayed on the residual images.

Figure 7: Visualization of transient solutions of the 2-D subcritical system ( 52 ) with $t\in[0,20]$ in the 1st, 2nd, and 3rd PCs of the 50-D embedding space. The 9 transient solutions initialize from the distributions $p_{\text{I}A}$ , $p_{\text{I}B}$ , and $p_{\text{I}D}$ in Fig. 6 with system parameters $\bm{\Theta}_{A}$ , $\bm{\Theta}_{B}$ , and $\bm{\Theta}_{D}$ and converge to three stationary distributions $p_{\text{S}A}$ , $p_{\text{S}B}$ , and $p_{\text{S}D}$ , which are attaching to the respective system parameters. For better illustration, similar to Fig. 2 , the points in the background are embedding features of random Gaussian distributions $\mathcal{N}([\mu_{1},\mu_{2}]^{\top},\text{diag}\{\sigma_{1}^{2},\sigma_{2}^{2}\})$ , and the two lines of distributions $G_{1}$ and $G_{2}$ are defined in Eq. ( 47 ). The grey dashed lines indicate the axes of the three leading PCs.

Figure 8: The training losses and test errors of the TPAPS on the 2-D subcritical system ( 52 ).

Time	$10^{3}$ batches	$10^{4}$ batches	$10^{5}$ batches	$10^{6}$ batches
$t=0.5$	0.3442/0.1681/0.3007	0.1501/0.0679/0.1348	0.0846/0.0365/0.0750	0.0689/0.0302/0.0610
$t=1.5$	0.3393/0.2051/0.2731	0.1535/0.0891/0.1273	0.0954/0.0451/0.0848	0.0661/0.0365/0.0558
$t=3.5$	0.3417/0.2314/0.2566	0.1511/0.1115/0.1119	0.0985/0.0535/0.0842	0.0670/0.0438/0.0547
$t=6.5$	0.2978/0.2333/0.1985	0.1444/0.1347/0.0912	0.1000/0.0712/0.0748	0.0662/0.0503/0.0503
$t=10.5$	0.2439/0.2101/0.1502	0.1431/0.1650/0.0801	0.1036/0.0926/0.0666	0.0649/0.0551/0.0462

Table 3: The mean, SD, and median (MEAN/SD/MED) of the L1 errors between TPAPS and MCS results across 2,500 random transient distributions of the 2-D subcritical system ( 52 ) at various training stages.

Figure 9: Efficient study of transient and near-stationary ( $t=20$ ) dynamics of the 2-D subcritcal system ( 52 ) using the TPAPS. (a) shows an example of marginal transient distributions of $x$ at different times. (b) shows the Hopf bifurcation of stationary distributions triggered by the system parameter $\mu$ by visualizing the marginal stationary distributions of $x$ . (c) investigates two-parameter bifurcation of stationary distributions depending on noise intensity $D$ and system parameter $\mu$ in terms of the SD $\sigma_{x}$ of the marginal distribution of $x$ and (d) demonstrates the two-parameter bifurcation of the non-ergodic distributions depending on $\mu$ and the x-coordinate $x_{0}$ of the center of initial distributions. (e) compares several groups of the SD $\sigma_{x}$ of the distributions obtained by TPAPS and the MCS.

Figure 10: Four transient solutions of the 2-D Duffing system ( 53 ) by TPAPS and MCS with different initial distributions and system parameters. The L1 errors are displayed on the residual images.

Figure 11: The training losses and test errors of the TPAPS on the 2-D Duffing system ( 53 ).

Time	$10^{3}$ batches	$10^{4}$ batches	$10^{5}$ batches	$10^{6}$ batches
$t=0.5$	0.9115/0.2776/0.9150	0.5292/0.2185/0.5053	0.3340/0.1317/0.3168	0.2331/0.1051/0.2166
$t=1.5$	0.7175/0.3149/0.6793	0.4036/0.2224/0.3551	0.2549/0.1435/0.2194	0.1892/0.1140/0.1589
$t=3.5$	0.4466/0.2799/0.3853	0.2747/0.2043/0.2143	0.1967/0.1413/0.1585	0.1524/0.1159/0.1160
$t=6.5$	0.2624/0.2023/0.1971	0.1528/0.1415/0.1027	0.1139/0.0995/0.0810	0.0963/0.0950/0.0630
$t=10.5$	0.1693/0.1353/0.1152	0.0969/0.0860/0.0648	0.0666/0.0586/0.0457	0.0600/0.0623/0.0364

Table 4: The mean, SD, and median (MEAN/SD/MED) of the L1 errors between TPAPS and MCS results across 2,500 random transient distributions of the Duffing system ( 53 ) at various training stages.

Figure 12: The TPAPS results of the Duffing system ( 53 ) at early training stages. The initial distribution, system parameters, and the results at $10^{6}$ training batches have been detailed in Fig. 10 (d).

Figure 13: Efficient study of transient and near-stationary ( $t=20$ ) dynamics of the Duffing system ( 53 ) using the TPAPS. (a) shows marginal transient distributions of $x$ at different times. (b) shows the impact of the noise amplitude $\sigma$ on the marginal stationary distribution of $x$ at $t=20$ . (c) investigates the dynamics of stationary distributions depending on the damping parameters $\eta$ and the stiffness parameter $\beta$ . The true transient solutions are obtained by MCS in (a2) and the true stationary results in (b2) and (c2) are calculated by Eq. ( 54 ). The resolution of all subfigures is $200\times 100$ .

Method	$10^{5}$ training batches					$10^{6}$ training batches					seconds/batch
Method	$t=0.5$	$t=1.5$	$t=3.5$	$t=6.5$	$t=10.5$	$t=0.5$	$t=1.5$	$t=3.5$	$t=6.5$	$t=10.5$	seconds/batch
Baseline	0.3340	0.2549	0.1967	0.1139	0.0666	0.2331	$\mathbf{0.1892}$	0.1524	0.0963	0.0600	1.11
TPAPS-A	$\mathbf{0.2296}$	$\mathbf{0.2290}$	$\mathbf{0.1768}$	0.1301	0.1098	$\mathbf{0.1896}$	0.1994	$\mathbf{0.1459}$	0.1204	0.1037	1.16
TPAPS-B	0.3263	0.2598	0.1968	0.1098	0.0613	0.2532	0.2035	0.1616	0.0952	0.0541	0.48
TPAPS-C	0.3402	0.2695	0.2114	0.1340	0.0781	0.2875	0.2275	0.1695	0.0996	0.0596	0.67
TPAPS-D	0.3583	0.2835	0.2007	$\mathbf{0.1067}$	$\mathbf{0.0547}$	0.2910	0.2261	0.1606	$\mathbf{0.0845}$	$\mathbf{0.0446}$	1.08

Table 5: The mean L1 errors between several TPAPS and MCS results across 2,500 random transient distributions of the Duffing system ( 53 ) at different times, while the last column shows the training speed. Best values are bold.

PDE Solving

PDE solving develops numerical and learned methods to approximate solutions of partial differential equations across scientific applications.

See all 249 papers

Operator learning for solving Fokker-Planck equations with various initial conditions
Jun 8, 2026Li Zeng, Xiaoliang Wan, Yaobin Wang +2PDE Solving PDE Operator Learning
School of Mathematics and Statistics, Fuzhou University, Fuzhou, China · Department of Mathematics and Center for Computation and Technology, Louisiana State University, Baton Rouge 70803, USA · Faculty of Science and Technology, Beijing Normal-Hong Kong Baptist University, Zhuhai 519087, China. +2
Adaptive Probability Flow Residual Minimization for High-Dimensional Fokker-Planck Equations
Dec 22, 2025Xiaolong Wu, Qifeng LiaoNeural PDE Solvers PDE Solving
School of Information Science and Technology, ShanghaiTech University, Shanghai, 201210, China
Neural Galerkin Normalizing Flow for Transition Probability Density Functions of Diffusion Models
Mar 19, 2026Riccardo Saporiti, Fabio NobileNormalizing Flows Neural PDE Solvers
CSQI, École Polytechnique Féderale de Lausanne, Switzerland

Jun 8, 2026J/K move · Enter open · S save

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

Li Zeng, Xiaoliang Wan, Yaobin Wang, Fabio Nobile+1

School of Mathematics and Statistics, Fuzhou University, Fuzhou, China · Department of Mathematics and Center for Computation and Technology, Louisiana State University, Baton Rouge 70803, USA · Faculty of Science and Technology, Beijing Normal-Hong Kong Baptist University, Zhuhai 519087, China.+2

PDE Solving PDE Operator Learning Physics-Informed ML Fokker–Planck Equation

Open paper AskPDF arXiv

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.

Abstract

PDE Solving

Explore similar work

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

Adaptive Probability Flow Residual Minimization for High-Dimensional Fokker-Planck Equations

Neural Galerkin Normalizing Flow for Transition Probability Density Functions of Diffusion Models